#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
file : main.py
author: Xiaohan Chen
email : chernxh@tamu.edu
last_modified: 2018-10-13
Main script. Start running model from main.py.
"""
import os , sys
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # BE QUIET!!!!
# timing
import time
from datetime import timedelta
from config import get_config
import utils.prob as problem
import utils.data as data
import utils.train as train
import numpy as np
import tensorflow as tf
try :
from sklearn.feature_extraction.image \
import extract_patches_2d, reconstruct_from_patches_2d
except Exception as e :
pass
def setup_model(config , **kwargs) :
untiedf = 'u' if config.untied else 't'
coordf = 'c' if config.coord else 's'
if config.net == 'LISTA' :
"""LISTA"""
config.model = ("LISTA_T{T}_lam{lam}_{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, untiedf=untiedf,
coordf=coordf, exp_id=config.exp_id))
from models.LISTA import LISTA
model = LISTA (kwargs['A'], T=config.T, lam=config.lam,
untied=config.untied, coord=config.coord,
scope=config.scope)
if config.net == 'LAMP' :
"""LAMP"""
config.model = ("LAMP_T{T}_lam{lam}_{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, untiedf=untiedf,
coordf=coordf, exp_id=config.exp_id))
from models.LAMP import LAMP
model = LAMP (kwargs['A'], T=config.T, lam=config.lam,
untied=config.untied, coord=config.coord,
scope=config.scope)
if config.net == 'LIHT' :
"""LIHT"""
from models.LIHT import LIHT
model = LIHT (p, T=config.T, lam=config.lam, y_=p.y_ , x0_=None ,
untied=config.untied , cord=config.coord)
if config.net == 'LISTA_cp' :
"""LISTA-CP"""
config.model = ("LISTA_cp_T{T}_lam{lam}_{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, untiedf=untiedf,
coordf=coordf, exp_id=config.exp_id))
from models.LISTA_cp import LISTA_cp
model = LISTA_cp (kwargs['A'], T=config.T, lam=config.lam,
untied=config.untied, coord=config.coord,
scope=config.scope)
if config.net == 'LISTA_ss' :
"""LISTA-SS"""
config.model = ("LISTA_ss_T{T}_lam{lam}_p{p}_mp{mp}_"
"{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, p=config.percent,
mp=config.max_percent, untiedf=untiedf,
coordf=coordf, exp_id=config.exp_id))
from models.LISTA_ss import LISTA_ss
model = LISTA_ss (kwargs['A'], T=config.T, lam=config.lam,
percent=config.percent, max_percent=config.max_percent,
untied=config.untied , coord=config.coord,
scope=config.scope)
if config.net == 'LISTA_cpss' :
"""LISTA-CPSS"""
config.model = ("LISTA_cpss_T{T}_lam{lam}_p{p}_mp{mp}_"
"{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, p=config.percent,
mp=config.max_percent, untiedf=untiedf,
coordf=coordf, exp_id=config.exp_id))
from models.LISTA_cpss import LISTA_cpss
model = LISTA_cpss (kwargs['A'], T=config.T, lam=config.lam,
percent=config.percent, max_percent=config.max_percent,
untied=config.untied , coord=config.coord,
scope=config.scope)
if config.net == 'TiLISTA':
"""TiLISTA"""
config.model = ("TiLISTA_T{T}_lam{lam}_p{p}_mp{mp}_"
"{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam,
p=config.percent, mp=config.max_percent,
coordf=coordf, exp_id=config.exp_id))
from models.TiLISTA import TiLISTA
# Note that TiLISTA is just LISTA-CPSS with tied weight in all layers.
model = TiLISTA(kwargs['A'], T=config.T, lam=config.lam,
percent=config.percent, max_percent=config.max_percent,
coord=config.coord, scope=config.scope)
if config.net == "ALISTA":
"""ALISTA"""
config.model = ("ALISTA_T{T}_lam{lam}_p{p}_mp{mp}_{W}_{coordf}_{exp_id}"
.format(T=config.T, lam=config.lam,
p=config.percent, mp=config.max_percent,
W=os.path.basename(config.W),
coordf=coordf, exp_id=config.exp_id))
W = np.load(config.W)
print("Pre-calculated weight W loaded from {}".format(config.W))
from models.ALISTA import ALISTA
model = ALISTA(kwargs['A'], T=config.T, lam=config.lam, W=W,
percent=config.percent, max_percent=config.max_percent,
coord=config.coord, scope=config.scope)
if config.net == 'LISTA_cs':
"""LISTA-CS"""
config.model = ("LISTA_cs_T{T}_lam{lam}_llam{llam}_"
"{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, llam=config.lasso_lam,
untiedf=untiedf, coordf=coordf,
exp_id=config.exp_id))
from models.LISTA_cs import LISTA_cs
model = LISTA_cs (kwargs['Phi'], kwargs['D'], T=config.T,
lam=config.lam, untied=config.untied,
coord=config.coord, scope=config.scope)
if config.net == 'LISTA_ss_cs' :
"""LISTA-SS-CS"""
config.model = ("LISTA_ss_cs_T{T}_lam{lam}_p{p}_mp{mp}_llam{llam}_"
"{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, p=config.percent,
mp=config.max_percent, llam=config.lasso_lam,
untiedf=untiedf, coordf=coordf,
exp_id=config.exp_id))
from models.LISTA_ss_cs import LISTA_ss_cs
model = LISTA_ss_cs (kwargs['Phi'], kwargs['D'], T=config.T,
lam=config.lam, percent=config.percent,
max_percent=config.max_percent,
untied=config.untied, coord=config.coord,
scope=config.scope)
if config.net == 'LISTA_cpss_cs' :
"""LISTA-CPSS-CS"""
config.model = ("LISTA_cpss_cs_T{T}_lam{lam}_p{p}_mp{mp}_llam{llam}_"
"{untiedf}_{coordf}_{exp_id}"
.format (T=config.T, lam=config.lam, p=config.percent,
mp=config.max_percent, llam=config.lasso_lam,
untiedf=untiedf, coordf=coordf,
exp_id=config.exp_id))
from models.LISTA_cpss_cs import LISTA_cpss_cs
model = LISTA_cpss_cs (kwargs['Phi'], kwargs['D'], T=config.T,
lam=config.lam, percent=config.percent,
max_percent=config.max_percent,
untied=config.untied, coord=config.coord,
scope=config.scope)
if config.net == 'LISTA_cp_conv':
"""LISTA-CP-CONV"""
config.model = ("LISTA_cp_conv_T{T}_lam{lam}_alpha{alpha}_"
"sigma{sigma}_{untiedf}_{exp_id}.npz"
.format(T=config.T, lam=config.lam, alpha=config.conv_alpha,
sigma=config.sigma, untiedf=untiedf, coordf=coordf,
exp_id=config.exp_id))
from models.LISTA_cp_conv import LISTA_cp_conv
model = LISTA_cp_conv(kwargs['filters'], T=config.T,
lam=config.lam, alpha=config.conv_alpha,
untied=config.untied, scope=config.scope)
if config.net == 'ALISTA_conv':
"""ALISTA-CONV"""
config.model = ("ALISTA_conv_T{T}_lam{lam}_alpha{alpha}_"
"sigma{sigma}_{exp_id}.npz"
.format(T=config.T, lam=config.lam, alpha=config.conv_alpha,
sigma=config.sigma, exp_id=config.exp_id))
W = np.load(config.W)
print("Pre-calculated weight W loaded from {}".format(config.W))
from models.ALISTA_conv import ALISTA_conv
model = ALISTA_conv(kwargs['filters'], W=W, T=config.T,
lam=config.lam, alpha=config.conv_alpha,
scope=config.scope)
if config.net == "AtoW_grad":
"""AtoW_grad"""
config.model = ("AtoW_grad_eT{eT}_Binit-{Binit}_eta{eta}_loss-{loss}_ps{ps}_lr{lr}_{id}"
.format(eT=config.eT, Binit=config.encoder_Binit, eta=config.eta,
loss=config.encoder_loss, ps=config.encoder_psigma,
lr=config.encoder_pre_lr, id=config.exp_id))
from models.AtoW_grad import AtoW_grad
model = AtoW_grad(config.M, config.N, config.eT, Binit=kwargs["Binit"],
eta=config.eta, loss=config.encoder_loss,
Q=kwargs["Q"], scope=config.scope)
if config.net == "robust_ALISTA":
"""Robust ALISTA"""
config.encoder = ("AtoW_grad_eT{eT}_Binit-{Binit}_eta{eta}_loss-{loss}_ps{ps}_lr{lr}_{id}"
.format(eT=config.eT, Binit=config.encoder_Binit, eta=config.eta,
loss=config.encoder_loss, ps=config.encoder_psigma,
lr=config.encoder_pre_lr, id=config.exp_id))
config.decoder = ("ALISTA_robust_T{T}_lam{lam}_p{p}_mp{mp}_{W}_{coordf}_{exp_id}"
.format(T=config.T, lam=config.lam,
p=config.percent, mp=config.max_percent,
W=os.path.basename(config.W),
coordf=coordf, exp_id=config.exp_id))
# set up encoder
from models.AtoW_grad import AtoW_grad
encoder = AtoW_grad(config.M, config.N, config.eT, Binit=kwargs["Binit"],
eta=config.eta, loss=config.encoder_loss,
Q=kwargs["Q"], scope=config.encoder_scope)
# set up decoder
from models.ALISTA_robust import ALISTA_robust
decoder = ALISTA_robust(M=config.M, N=config.N, T=config.T,
percent=config.percent, max_percent=config.max_percent,
coord=config.coord, scope=config.decoder_scope)
model_desc = ("robust_" + config.encoder + '_' + config.decoder +
"_elr{}_dlr{}_psmax{}_psteps{}_{}"
.format(config.encoder_lr, config.decoder_lr,
config.psigma_max, config.psteps, config.exp_id))
model_dir = os.path.join(config.expbase, model_desc)
config.resfn = os.path.join(config.resbase, model_desc)
if not os.path.exists(model_dir):
if config.test:
raise ValueError("Testing folder {} not existed".format(model_dir))
else:
os.makedirs(model_dir)
config.enc_load = os.path.join(config.expbase, config.encoder)
config.dec_load = os.path.join(config.expbase, config.decoder.replace("_robust", ""))
config.encoderfn = os.path.join(model_dir, config.encoder)
config.decoderfn = os.path.join(model_dir, config.decoder)
return encoder, decoder
config.modelfn = os.path.join(config.expbase, config.model)
config.resfn = os.path.join(config.resbase, config.model)
print ("model disc:", config.model)
return model
############################################################
###################### Training #######################
############################################################
def run_train(config) :
if config.task_type == "sc":
run_sc_train(config)
elif config.task_type == "cs":
run_cs_train(config)
elif config.task_type == "denoise":
run_denoise_train(config)
elif config.task_type == "encoder":
run_encoder_train(config)
elif config.task_type == "robust":
run_robust_train(config)
def run_sc_train(config) :
"""Load problem."""
if not os.path.exists(config.probfn):
raise ValueError ("Problem file not found.")
else:
p = problem.load_problem(config.probfn)
"""Set up model."""
model = setup_model (config, A=p.A)
"""Set up input."""
config.SNR = np.inf if config.SNR == 'inf' else float (config.SNR)
y_, x_, y_val_, x_val_ = (
train.setup_input_sc (
config.test, p, config.tbs, config.vbs, config.fixval,
config.supp_prob, config.SNR, config.magdist, **config.distargs))
"""Set up training."""
stages = train.setup_sc_training (
model, y_, x_, y_val_, x_val_, None,
config.init_lr, config.decay_rate, config.lr_decay)
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
# graph initialization
sess.run (tf.global_variables_initializer ())
# start timer
start = time.time ()
# train model
model.do_training(sess, stages, config.modelfn, config.scope,
config.val_step, config.maxit, config.better_wait)
# end timer
end = time.time ()
elapsed = end - start
print ("elapsed time of training = " + str (timedelta (seconds=elapsed)))
# end of run_sc_train
def run_cs_train (config) :
"""Load dictionary and sensing matrix."""
Phi = np.load (config.sensing)['A']
D = np.load (config.dict)
"""Set up model."""
model = setup_model (config, Phi=Phi, D=D)
"""Set up inputs."""
y_, f_, y_val_, f_val_ = train.setup_input_cs(config.train_file,
config.val_file,
config.tbs, config.vbs)
"""Set up training."""
stages = train.setup_cs_training (
model, y_, f_, y_val_, f_val_, None, config.init_lr, config.decay_rate,
config.lr_decay, config.lasso_lam)
"""Start training."""
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
# graph initialization
sess.run (tf.global_variables_initializer ())
# start timer
start = time.time ()
# train model
model.do_training (sess, stages, config.modelfn, config.scope,
config.val_step, config.maxit, config.better_wait)
# end timer
end = time.time ()
elapsed = end - start
print ("elapsed time of training = " + str (timedelta (seconds=elapsed)))
# end of run_cs_train
def run_denoise_train (config) :
"""Load problem."""
import utils.prob_conv as problem
if not os.path.exists (config.probfn):
raise ValueError ("Problem file not found.")
else:
p = problem.load_problem (config.probfn)
"""Set up model."""
model = setup_model (config, filters=p._fs)
"""Set up input."""
# training
clean_ = data.bsd500_denoise_inputs(config.data_folder, config.train_file, config.tbs,
config.height_crop, config.width_crop, config.num_epochs)
clean_.set_shape((config.tbs, *clean_.get_shape()[1:],))
# validation
clean_val_ = data.bsd500_denoise_inputs(config.data_folder, config.val_file, config.vbs,
config.height_crop, config.width_crop, 1)
clean_val_.set_shape((config.vbs, *clean_val_.get_shape()[1:],))
# add noise
noise_ = tf.random_normal(clean_.shape, stddev=config.denoise_std,
dtype=tf.float32)
noise_val_ = tf.random_normal(clean_val_.shape, stddev=config.denoise_std,
dtype=tf.float32)
noisy_ = clean_ + noise_
noisy_val_= clean_val_ + noise_val_
# fix validation set
with tf.name_scope ('input'):
clean_val_ = tf.get_variable(name='clean_val',
dtype=tf.float32,
initializer=clean_val_)
noisy_val_ = tf.get_variable(name='noisy_val',
dtype=tf.float32,
initializer=noisy_val_)
"""Set up training."""
stages = train.setup_denoise_training(
model, noisy_, clean_, noisy_val_, clean_val_,
None, config.init_lr, config.decay_rate, config.lr_decay)
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
# graph initialization
sess.run (tf.global_variables_initializer ())
# start timer
start = time.time ()
# train model
model.do_training(sess, stages, config.modelfn, config.scope,
config.val_step, config.maxit, config.better_wait)
# end timer
end = time.time ()
elapsed = end - start
print ("elapsed time of training = " + str (timedelta (seconds=elapsed)))
# end of run_denoise_train
def run_encoder_train(config):
"""Load problem."""
if not os.path.exists(config.probfn):
raise ValueError("Problem file not found.")
else:
p = problem.load_problem(config.probfn)
"""Load the Q reweighting matrix."""
if config.Q is None: # use default Q reweighting matrix
if "re" in config.encoder_loss: # if using reweighted loss
Q = np.sqrt((np.ones(shape=(config.N, config.N), dtype=np.float32) +
np.eye(config.N, dtype=np.float32) * (config.N - 2)))
else:
Q = None
elif os.path.exists(config.Q) and config.Q.endswith(".npy"):
Q = np.load(config.Q)
assert Q.shape == (config.N, config.N)
else:
raise ValueError("Invalid parameter `--Q`\n"
"A valid `--Q` parameter should be one of the following:\n"
" 1) omitted for default value as in the paper;\n"
" 2) path/to/your/npy/file that contains your Q matrix.\n")
"""Binit matrix."""
if config.encoder_Binit == "default":
Binit = p.A
elif config.Binit in ["uniform", "normal"]:
pass
else:
raise ValueError("Invalid parameter `--Binit`\n"
"A valid `--Binit` parameter should be one of the following:\n"
" 1) omitted for default value `p.A`;\n"
" 2) `normal` or `uniform`.\n")
"""Set up model."""
model = setup_model(config, Binit=Binit, Q=Q)
print("The trained model will be saved in {}".format(config.model))
"""Set up training."""
from utils.tf import get_loss_func, bmxbm, mxbm
with tf.name_scope ('input'):
A_ = tf.constant(p.A, dtype=tf.float32)
perturb_ = tf.random.normal(shape=(config.Abs, config.M, config.N),
mean=0.0, stddev=config.encoder_psigma,
dtype=tf.float32)
Ap_ = A_ + perturb_
Ap_ = Ap_ / tf.sqrt(tf.reduce_sum(tf.square( Ap_ ), axis=1, keepdims=True))
Apt_ = tf.transpose(Ap_, [0,2,1])
W_ = model.inference(Ap_)
"""Set up loss."""
eye_ = tf.eye(config.N, batch_shape=[config.Abs], dtype=tf.float32)
residual_ = bmxbm(Apt_, W_, batch_first=True) - eye_
loss_func = get_loss_func(config.encoder_loss, model._Q_)
loss_ = loss_func(residual_)
# fix validation set
Ap_val_ = tf.get_variable(name='Ap_val', dtype=tf.float32,
initializer=Ap_, trainable=False)
Apt_val_ = tf.transpose(Ap_val_, [0,2,1])
W_val_ = model.inference(Ap_val_)
# validation loss
residual_val_ = bmxbm(Apt_val_, W_val_, batch_first=True) - eye_
loss_val_ = loss_func(residual_val_)
"""Set up optimizer."""
global_step = tf.Variable(0, trainable=False)
lr = tf.train.exponential_decay(config.encoder_lr, global_step,
5000, 0.75, staircase=True)
learning_step = (tf.train.AdamOptimizer(lr)
.minimize(loss_, global_step=global_step))
# create session and initialize the graph
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
sess.run (tf.global_variables_initializer ())
# start timer
start = time.time ()
for i in range (config.maxit):
# training step
_, loss = sess.run([learning_step, loss_])
# validation step
if i % config.val_step == 0:
# validation step
loss_val = sess.run(loss_val_)
sys.stdout.write (
"\ri={i:<7d} | train_loss={train_loss:.6f} | "
"loss_val={loss_val:.6f}"
.format(i=i, train_loss=loss, loss_val=loss_val))
sys.stdout.flush()
# end timer
end = time.time()
elapsed = end - start
print("elapsed time of training = " + str(timedelta(seconds=elapsed)))
train.save_trainable_variables (sess, config.modelfn, config.scope)
print("model saved to {}".format(config.modelfn))
# end of run_encoder_train
def run_robust_train(config):
"""Load problem."""
if not os.path.exists(config.probfn):
raise ValueError("Problem file not found.")
else:
p = problem.load_problem(config.probfn)
"""Set up input."""
# `psigma` is a list of standard deviations for curriculum learning
psigmas = np.linspace(0, config.psigma_max, config.psteps)[1:]
psigma_ = tf.placeholder(dtype=tf.float32, shape=())
with tf.name_scope ('input'):
Ap_, y_, x_ = train.setup_input_robust(p.A, psigma_, config.msigma,
p.pnz, config.Abs, config.xbs)
if config.net != "robust_ALISTA":
# If not joint robust training
# reshape y_ into shape (m, Abs * xbs)
# reshape x_ into shape (n, Abs * xbs)
y_ = tf.reshape(tf.transpose(y_, [1, 0, 2]), (config.M, -1))
x_ = tf.reshape(tf.transpose(x_, [1, 0, 2]), (config.N, -1))
# fix validation set
Ap_val_ = tf.get_variable(name="Ap_val", dtype=tf.float32, initializer=Ap_)
y_val_ = tf.get_variable(name="y_val", dtype=tf.float32, initializer=y_)
x_val_ = tf.get_variable(name="x_val", dtype=tf.float32, initializer=x_)
"""Set up model."""
if config.net == "robust_ALISTA":
"""Load the Q reweighting matrix."""
if config.Q is None: # use default Q reweighting matrix
if "re" in config.encoder_loss: # if using reweighted loss
Q = np.sqrt((np.ones(shape=(config.N, config.N), dtype=np.float32) +
np.eye(config.N, dtype=np.float32) * (config.N - 2)))
else:
Q = None
elif os.path.exists(config.Q) and config.Q.endswith(".npy"):
Q = np.load(config.Q)
assert Q.shape == (config.N, config.N)
else:
raise ValueError("Invalid parameter `--Q`\n"
"A valid `--Q` parameter should be one of the following:\n"
" 1) omitted for default value as in the paper;\n"
" 2) path/to/your/npy/file that contains your Q matrix.\n")
"""Binit matrix."""
if config.encoder_Binit == "default":
Binit = p.A
elif config.Binit in ["uniform", "normal"]:
pass
else:
raise ValueError("Invalid parameter `--Binit`\n"
"A valid `--Binit` parameter should be one of the following:\n"
" 1) omitted for default value `p.A`;\n"
" 2) `normal` or `uniform`.\n")
encoder, decoder = setup_model(config, Q=Q, Binit=Binit)
W_ = encoder.inference(Ap_)
W_val_ = encoder.inference(Ap_val_)
xh_ = decoder.inference(y_, Ap_, W_, x0_=None)[-1]
xh_val_ = decoder.inference(y_val_, Ap_val_, W_val_, x0_=None)[-1]
else:
decoder = setup_model(config, A=p.A)
xh_ = decoder.inference(y_, None)[-1]
xh_val_ = decoder.inference(y_val_, None)[-1]
config.dec_load = config.modelfn
config.decoder = (
"robust_" + config.model + '_ps{ps}_nsteps{nsteps}_ms{ms}_lr{lr}'
.format(ps=config.psigma_max, nsteps=config.psteps,
ms=config.msigma, lr=config.decoder_lr))
config.decoderfn = os.path.join(config.expbase, config.decoder)
print("\npretrained decoder loaded from {}".format(config.modelfn))
print("trained augmented model will be saved to {}".format(config.decoderfn))
"""Set up loss."""
loss_ = tf.nn.l2_loss (xh_ - x_)
nmse_denom_ = tf.nn.l2_loss (x_)
nmse_ = loss_ / nmse_denom_
db_ = 10.0 * tf.log (nmse_) / tf.log (10.0)
# validation
loss_val_ = tf.nn.l2_loss (xh_val_ - x_val_)
nmse_denom_val_ = tf.nn.l2_loss (x_val_)
nmse_val_ = loss_val_ / nmse_denom_val_
db_val_ = 10.0 * tf.log (nmse_val_) / tf.log (10.0)
"""Set up optimizer."""
global_step_ = tf.Variable (0, trainable=False)
if config.net == "robust_ALISTA":
"""Encoder and decoder apply different initial learning rate."""
# get trainable variable for de encoder and decoder
encoder_variables_ = tf.get_collection(
key=tf.GraphKeys.TRAINABLE_VARIABLES, scope=config.encoder_scope)
decoder_variables_ = tf.get_collection(
key=tf.GraphKeys.TRAINABLE_VARIABLES, scope=config.decoder_scope)
trainable_variables_ = encoder_variables_ + decoder_variables_
# calculate gradients w.r.t. all trainable variables in the model
grads_ = tf.gradients (loss_, trainable_variables_)
encoder_grads_ = grads_[:len (encoder_variables_)]
decoder_grads_ = grads_[len (encoder_variables_):]
# define learning rates for optimizers over two parts
global_step_ = tf.Variable (0, trainable=False)
encoder_lr_ = tf.train.exponential_decay(
config.encoder_lr, global_step_, 5000, 0.75, staircase=False)
encoder_opt_ = tf.train.AdamOptimizer(encoder_lr_)
decoder_lr_ = tf.train.exponential_decay(
config.decoder_lr, global_step_, 5000, 0.75, staircase=False)
decoder_opt_ = tf.train.AdamOptimizer(decoder_lr_)
# define training operator
encoder_op_ = encoder_opt_.apply_gradients(
zip(encoder_grads_, encoder_variables_))
decoder_op_ = decoder_opt_.apply_gradients(
zip(decoder_grads_, decoder_variables_))
learning_step_ = tf.group(encoder_op_, decoder_op_)
else:
lr_ = tf.train.exponential_decay(config.decoder_lr, global_step_,
5000, 0.75, staircase=False)
learning_step_ = (tf.train.AdamOptimizer(lr_)
.minimize(loss_, global_step=global_step_))
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
# graph initialization
sess.run (tf.global_variables_initializer (),
feed_dict={psigma_: psigmas[0]})
# load pre-trained model(s)
if config.net == "robust_ALISTA":
encoder.load_trainable_variables(sess, config.enc_load)
decoder.load_trainable_variables(sess, config.dec_load)
# start timer
start = time.time ()
for psigma in psigmas:
print ('\ncurrent sigma: {}'.format (psigma))
global_step_.initializer.run ()
for i in range (config.maxit):
db, loss, _ = sess.run([db_, loss_, learning_step_],
feed_dict={psigma_: psigma})
if i % config.val_step == 0:
db_val, loss_val = sess.run([db_val_, loss_val_],
feed_dict={psigma_: psigma})
sys.stdout.write(
"\ri={i:<7d} | loss_train={loss_train:.6f} | "
"db_train={db_train:.6f} | loss_val={loss_val:.6f} | "
"db_val={db_val:.6f}".format(
i=i, loss_train=loss, db_train=db,
loss_val=loss_val, db_val=db_val))
sys.stdout.flush()
if config.net == "robust_ALISTA":
encoder.save_trainable_variables(sess, config.encoderfn)
decoder.save_trainable_variables(sess, config.decoderfn)
# end timer
end = time.time()
elapsed = end - start
print("elapsed time of training = " + str(timedelta(seconds=elapsed)))
# end of run_robust_train
############################################################
###################### Testing ########################
############################################################
def run_test (config):
if config.task_type == "sc":
run_sc_test (config)
elif config.task_type == "cs":
run_cs_test (config)
elif config.task_type == "denoise":
run_denoise_test(config)
elif config.task_type == "robust":
run_robust_test(config)
def run_sc_test (config) :
"""
Test model.
"""
"""Load problem."""
if not os.path.exists (config.probfn):
raise ValueError ("Problem file not found.")
else:
p = problem.load_problem (config.probfn)
"""Load testing data."""
xt = np.load (config.xtest)
"""Set up input for testing."""
config.SNR = np.inf if config.SNR == 'inf' else float (config.SNR)
input_, label_ = (
train.setup_input_sc (config.test, p, xt.shape [1], None, False,
config.supp_prob, config.SNR,
config.magdist, **config.distargs))
"""Set up model."""
model = setup_model (config , A=p.A)
xhs_ = model.inference (input_, None)
"""Create session and initialize the graph."""
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
# graph initialization
sess.run (tf.global_variables_initializer ())
# load model
model.load_trainable_variables (sess , config.modelfn)
nmse_denom = np.sum (np.square (xt))
supp_gt = xt != 0
lnmse = []
lspar = []
lsperr = []
lflspo = []
lflsne = []
# test model
for xh_ in xhs_ :
xh = sess.run (xh_ , feed_dict={label_:xt})
# nmse:
loss = np.sum (np.square (xh - xt))
nmse_dB = 10.0 * np.log10 (loss / nmse_denom)
print (nmse_dB)
lnmse.append (nmse_dB)
supp = xh != 0.0
# intermediate sparsity
spar = np.sum (supp , axis=0)
lspar.append (spar)
# support error
sperr = np.logical_xor(supp, supp_gt)
lsperr.append (np.sum (sperr , axis=0))
# false positive
flspo = np.logical_and (supp , np.logical_not (supp_gt))
lflspo.append (np.sum (flspo , axis=0))
# false negative
flsne = np.logical_and (supp_gt , np.logical_not (supp))
lflsne.append (np.sum (flsne , axis=0))
res = dict (nmse=np.asarray (lnmse),
spar=np.asarray (lspar),
sperr=np.asarray (lsperr),
flspo=np.asarray (lflspo),
flsne=np.asarray (lflsne))
np.savez (config.resfn , **res)
# end of test
def run_cs_test (config) :
from utils.cs import imread_CS_py, img2col_py, col2im_CS_py
from skimage.io import imsave
"""Load dictionary and sensing matrix."""
Phi = np.load (config.sensing) ['A']
D = np.load (config.dict)
# loading compressive sensing settings
M = Phi.shape [0]
F = Phi.shape [1]
N = D.shape [1]
assert M == config.M and F == config.F and N == config.N
patch_size = int (np.sqrt (F))
assert patch_size ** 2 == F
"""Set up model."""
model = setup_model (config, Phi=Phi, D=D)
"""Inference."""
y_ = tf.placeholder (shape=(M, None), dtype=tf.float32)
_, fhs_ = model.inference (y_, None)
"""Start testing."""
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
# graph initialization
sess.run (tf.global_variables_initializer ())
# load model
model.load_trainable_variables (sess , config.modelfn)
# calculate average NMSE and PSRN on test images
test_dir = './data/test_images/'
test_files = os.listdir (test_dir)
avg_nmse = 0.0
avg_psnr = 0.0
overlap = 0
stride = patch_size - overlap
out_dir = "./data/recon_images"
if 'joint' in config.net :
D = sess.run (model.D_)
for test_fn in test_files :
# read in image
out_fn = test_fn[:-4] + "_recon_{}.png".format(config.sample_rate)
out_fn = os.path.join(out_dir, out_fn)
test_fn = os.path.join (test_dir, test_fn)
test_im, H, W, test_im_pad, H_pad, W_pad = \
imread_CS_py (test_fn, patch_size, stride)
test_fs = img2col_py (test_im_pad, patch_size, stride)
# remove dc from features
test_dc = np.mean (test_fs, axis=0, keepdims=True)
test_cfs = test_fs - test_dc
test_cfs = np.asarray (test_cfs) / 255.0
# sensing signals
test_ys = np.matmul (Phi, test_cfs)
num_patch = test_ys.shape [1]
rec_cfs = sess.run (fhs_ [-1], feed_dict={y_: test_ys})
rec_fs = rec_cfs * 255.0 + test_dc
# patch-level NMSE
patch_err = np.sum (np.square (rec_fs - test_fs))
patch_denom = np.sum (np.square (test_fs))
avg_nmse += 10.0 * np.log10 (patch_err / patch_denom)
rec_im = col2im_CS_py (rec_fs, patch_size, stride,
H, W, H_pad, W_pad)
# image-level PSNR
image_mse = np.mean (np.square (np.clip(rec_im, 0.0, 255.0) - test_im))
avg_psnr += 10.0 * np.log10 (255.**2 / image_mse)
num_test_ims = len (test_files)
print ('Average Patch-level NMSE is {}'.format (avg_nmse / num_test_ims))
print ('Average Image-level PSNR is {}'.format (avg_psnr / num_test_ims))
# end of cs_testing
def run_denoise_test(config) :
import glob
from PIL import Image
"""Load problem."""
import utils.prob_conv as problem
if not os.path.exists(config.probfn):
raise ValueError("Problem file not found.")
else:
p = problem.load_problem(config.probfn)
"""Set up model."""
model = setup_model(config, filters=p._fs)
"""Set up input."""
orig_clean_ = tf.placeholder(dtype=tf.float32, shape=(None, 256, 256, 1))
clean_ = orig_clean_ * (1.0 / 255.0)
mean_ = tf.reduce_mean(clean_, axis=(1,2,3,), keepdims=True)
demean_ = clean_ - mean_
"""Add noise."""
noise_ = tf.random_normal (tf.shape (demean_), stddev=config.denoise_std,
dtype=tf.float32)
noisy_ = demean_ + noise_
"""Inference."""
_, recons_ = model.inference(noisy_, None)
recon_ = recons_[-1]
# denormalize
recon_ = (recon_ + mean_) * 255.0
"""PSNR."""
mse2_ = tf.reduce_mean(tf.square(orig_clean_ - recon_), axis=(1,2,3,))
psnr_ = 10.0 * tf.log(255.0 ** 2 / mse2_) / tf.log (10.0)
avg_psnr_ = tf.reduce_mean(psnr_)
"""Load test images."""
test_images = []
filenames = []
types = ("*.tif", "*.png", "*.jpg", "*.gif",)
for type in types:
filenames.extend(glob.glob(os.path.join(config.test_dir, type)))
for filename in filenames:
im = Image.open(filename)
if im.size != (256, 256):
im = im.resize((256, 256))
test_images.append(np.asarray (im).astype(np.float32))
test_images = np.asarray(test_images).reshape((-1, 256, 256, 1))
tfconfig = tf.ConfigProto(allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session(config=tfconfig) as sess:
# graph initialization
sess.run(tf.global_variables_initializer())
# load model
model.load_trainable_variables(sess, config.modelfn)
# testing
psnr, avg_psnr = sess.run([psnr_, avg_psnr_],
feed_dict={orig_clean_:test_images})
print('file names\t| PSNR/dB')
for fname, p in zip(filenames, psnr):
print(os.path.basename (fname), '\t', p)
print("average PSNR = {} dB".format(avg_psnr))
print("full PSNR records on testing set are stored in {}".format(config.resfn))
np.save(config.resfn, psnr)
sum_time = 0.0
ntimes = 200
for i in range(ntimes):
# start timer
start = time.time()
# testing
sess.run(recon_, feed_dict={orig_clean_:test_images})
# end timer
end = time.time()
sum_time = sum_time + end - start
print("average elapsed time for one image inference = " +
str(timedelta(seconds=sum_time/ntimes/test_images.shape[0])))
# start timer
start = time.time()
# end of run_denoise_test
def run_robust_test(config):
"""Load problem."""
print(config.probfn)
if not os.path.exists(config.probfn):
raise ValueError("Problem file not found.")
else:
p = problem.load_problem(config.probfn)
"""Set tesing data."""
test_As = np.load('./data/robust_test_A.npz')
x = np.load('./data/xtest_n500_p10.npy')
"""Set up input."""
psigmas = sorted([float(k) for k in test_As.keys()])
psigma_ = tf.placeholder(dtype=tf.float32, shape=())
with tf.name_scope ('input'):
Ap_ = tf.placeholder (dtype=tf.float32, shape=(250, 500))
x_ = tf.placeholder (dtype=tf.float32, shape=(500, None))
## measure y_ from x_ using Ap_
y_ = tf.matmul (Ap_, x_)
"""Set up model."""
if config.net == "robust_ALISTA":
"""Load the Q reweighting matrix."""
if config.Q is None: # use default Q reweighting matrix
if "re" in config.encoder_loss: # if using reweighted loss
Q = np.sqrt((np.ones(shape=(config.N, config.N), dtype=np.float32) +
np.eye(config.N, dtype=np.float32) * (config.N - 2)))
else:
Q = None
elif os.path.exists(config.Q) and config.Q.endswith(".npy"):
Q = np.load(config.Q)
assert Q.shape == (config.N, config.N)
else:
raise ValueError("Invalid parameter `--Q`\n"
"A valid `--Q` parameter should be one of the following:\n"
" 1) omitted for default value as in the paper;\n"
" 2) path/to/your/npy/file that contains your Q matrix.\n")
"""Binit matrix."""
if config.encoder_Binit == "default":
Binit = p.A
elif config.Binit in ["uniform", "normal"]:
pass
else:
raise ValueError("Invalid parameter `--Binit`\n"
"A valid `--Binit` parameter should be one of the following:\n"
" 1) omitted for default value `p.A`;\n"
" 2) `normal` or `uniform`.\n")
encoder, decoder = setup_model(config, Q=Q, Binit=Binit)
W_ = tf.squeeze(encoder.inference(tf.expand_dims(Ap_, axis=0)), axis=0)
xh_ = decoder.inference(y_, Ap_, W_, x0_=None)[-1]
else:
decoder = setup_model(config, A=p.A)
xh_ = decoder.inference(y_, None)[-1]
config.decoder = (
"robust_" + config.model + '_ps{ps}_nsteps{nsteps}_ms{ms}_lr{lr}'
.format(ps=config.psigma_max, nsteps=config.psteps,
ms=config.msigma, lr=config.decoder_lr))
config.decoderfn = os.path.join(config.expbase, config.decoder)
print("\ntrained augmented model loaded from {}".format(config.decoderfn))
"""Set up loss."""
loss_ = tf.nn.l2_loss (xh_ - x_)
nmse_denom_ = tf.nn.l2_loss (x_)
nmse_ = loss_ / nmse_denom_
db_ = 10.0 * tf.log (nmse_) / tf.log (10.0)
tfconfig = tf.ConfigProto (allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
with tf.Session (config=tfconfig) as sess:
# graph initialization
sess.run (tf.global_variables_initializer (),
feed_dict={psigma_: psigmas[0]})
# load pre-trained model(s)
if config.net == "robust_ALISTA":
encoder.load_trainable_variables(sess, config.encoderfn)
decoder.load_trainable_variables(sess, config.decoderfn)
# start timer
start = time.time ()
res = dict (sigmas=np.array (psigmas))
avg_dBs = []
print ('sigma\tnmse')
sum_time = 0.0
tcounter = 0
for psigma in psigmas:
Aps = test_As[str(psigma)]
sum_dB = 0.0
counter = 0
for Ap in Aps:
db = sess.run(db_, feed_dict={x_:x, Ap_:Ap})
# start timer
start = time.time ()
# inference
sess.run (xh_, feed_dict={x_:x, Ap_:Ap})
# end timer
end = time.time ()
elapsed = end - start
sum_time = sum_time + elapsed
tcounter = tcounter + 1
sum_dB = sum_dB + db
counter = counter + 1
avg_dB = sum_dB / counter
print(psigma, '\t', avg_dB)
avg_dBs.append (avg_dB)
print("average elapsed time of inference =", str (timedelta (seconds=sum_time/tcounter)))
res['avg_dBs'] = np.asarray(avg_dBs)
print('saving results to', config.resfn)
np.savez (config.resfn, **res)
# end of run_robust_test
############################################################
####################### Main #########################
############################################################
def main ():
# parse configuration
config, _ = get_config()
# set visible GPUs
os.environ['CUDA_VISIBLE_DEVICES'] = config.gpu
if config.test:
run_test (config)
else:
run_train (config)
# end of main
if __name__ == "__main__":
main ()
