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)
# Create target Directory if don't exist
dirName = config.modelfn
if not os.path.exists(dirName):
os.mkdir(dirName)
print("Directory ", dirName, " Created ")
else:
print("Directory ", dirName, " already exists")
# model.weights[2].assign(2 * model.weights[2].numpy())
model.weights[2].assign((1234.0, 567.0))
model.weights[3].assign((1234.0, 567.0))
model.save_weights(config.modelfn + '/ckpt')
print('model weights ... ')
print(model.weights[2])
print(model.weights[3])
loaded_model = setup_model(config, A=p.A)
loaded_model.load_weights(config.modelfn + '/ckpt')
print('loaded model weights ... ')
print(loaded_model.weights[2])
print(loaded_model.weights[3])
# loaded_model.load_weights('ckpt')
# model.load_weights('my_model.h5', by_name=False, skip_mismatch=False)
"""Set up input."""
config.SNR = np.inf if config.SNR == 'inf' else float(config.SNR)
data_set = DataSet.DataSet(config, p)
"""Set up training."""
stages = train.setup_sc_training(model, data_set, None, config.init_lr,
config.decay_rate, config.lr_decay)
tfconfig = tf.compat.v1.ConfigProto(allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
# start timer
start = time.time()
# train model
model.do_training(stages, data_set, 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)))
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, name='new_global')
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))
lnmse = []
# 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)
res = dict(nmse=np.asarray(lnmse))
np.savez(config.resfn, **res)
def run_sc_test_with_created_session(config, sess, model):
"""
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(True, p, xt.shape[1], None, False,
config.supp_prob, config.SNR,
config.magdist, **config.distargs))
xhs_ = model.inference(input_, None)
nmse_denom = np.sum(np.square(xt))
lnmse = []
# 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)
res = dict(nmse=np.asarray(lnmse))
# print(lnmse)
np.savez(config.resfn, **res)
# end of test
return lnmse
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)))
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)
config.SNR = np.inf if config.SNR == 'inf' else float(config.SNR)
tfconfig = tf.ConfigProto(allow_soft_placement=True)
tfconfig.gpu_options.allow_growth = True
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 model."""
global_model = setup_model(config, A=p.A, name='global')
comm_rounds = 2
num_clients = config.num_cl
do_3_stage_training = True
new_Layer_model = setup_model(config, A=p.A, name='layer')
client_data = y_.shape[1] // num_clients
client_val_data = y_val_.shape[1] // num_clients
client_models_dict = {}
client_stages_dict = {}
for i in range(num_clients):
client_models_dict[i] = setup_model(config,
A=p.A,
name='client%d' % (i))
client_stages_dict[i] = train.setup_sc_training(
client_models_dict[i], y_[:,
i * client_data:(i + 1) * client_data],
x_[:, i * client_data:(i + 1) * client_data],
y_val_[:, i * client_val_data:(i + 1) * client_val_data],
x_val_[:, i * client_val_data:(i + 1) * client_val_data], None,
config.init_lr, config.decay_rate, config.lr_decay, i,
do_3_stage_training)
Layer_wise_lnmse = []
with tf.Session(config=tfconfig) as sess:
nmse_for_all_rounds = []
sess.run(tf.global_variables_initializer())
for layer in range(global_model._T):
print("Layer ", layer + 1)
for rounds in range(comm_rounds):
global_layer = global_model.vars_in_layer[layer]
global_weights = get_weight_obj(global_layer[0].eval(sess),
global_layer[1].eval(sess),
global_layer[2].eval(sess))
client_weight_list = get_weight_obj([], [], [])
for client in range(num_clients):
client_model = client_models_dict[client]
client_model.set_weights_at_layer(global_weights, layer,
sess)
print(
'--------------------------------------------------------------'
)
print(f'Round: {rounds+1:02} | client no: {client+1:02}')
stages = client_stages_dict[client]
start = time.time()
if do_3_stage_training:
for stage_number in range(layer * 3, layer * 3 + 3):
client_model.do_training_one_stage(
sess, stages[stage_number], config.modelfn,
config.scope, config.val_step, config.maxit,
config.better_wait)
else:
client_model.do_training_one_stage(
sess, stages[layer], config.modelfn, config.scope,
config.val_step, config.maxit, config.better_wait)
end = time.time()
elapsed = end - start
print("elapsed time of training = " +
str(timedelta(seconds=elapsed)))
client_layers = client_model.vars_in_layer[layer]
client_weights = get_weight_obj(
client_layers[0].eval(sess),
client_layers[1].eval(sess),
client_layers[2].eval(sess))
client_weight_list['B'].append(client_weights['B'])
client_weight_list['W'].append(client_weights['W'])
client_weight_list['theta'].append(client_weights['theta'])
new_weights = {}
new_weights['B'] = tf.convert_to_tensor(
np.mean(client_weight_list['B'], axis=0))
new_weights['W'] = tf.convert_to_tensor(
np.mean(client_weight_list['W'], axis=0))
new_weights['theta'] = tf.convert_to_tensor(
np.mean(client_weight_list['theta'], axis=0))
global_model.set_weights_at_layer(new_weights, layer, sess)
if do_3_stage_training:
for layer_in in range(layer + 1):
client_weight_list = get_weight_obj([], [], [])
for client in range(num_clients):
client_layer = client_models_dict[
client].vars_in_layer[layer_in]
client_weights = get_weight_obj(
client_layer[0].eval(sess),
client_layer[1].eval(sess),
client_layer[2].eval(sess))
client_weight_list['B'].append(client_weights['B'])
client_weight_list['W'].append(client_weights['W'])
client_weight_list['theta'].append(
client_weights['theta'])
new_weights = {}
new_weights['B'] = tf.convert_to_tensor(
np.mean(client_weight_list['B'], axis=0))
new_weights['W'] = tf.convert_to_tensor(
np.mean(client_weight_list['W'], axis=0))
new_weights['theta'] = tf.convert_to_tensor(
np.mean(client_weight_list['theta'], axis=0))
new_Layer_model.set_weights_at_layer(
new_weights, layer_in, sess)
lnmse2 = run_sc_test_with_created_session(
config, sess, new_Layer_model)
Layer_wise_lnmse.append(lnmse2[layer + 1])
np.savez('Layer_lnmse_' + str(num_clients) + "_" + str(config.maxit),
Layer_wise_lnmse)
print("Layer wise performance")
print(Layer_wise_lnmse)
if do_3_stage_training:
new_global_model = setup_model(config, A=p.A, name='new_global')
for layer in range(new_global_model._T):
client_weight_list = get_weight_obj([], [], [])
for client in range(num_clients):
client_layer = client_models_dict[client].vars_in_layer[
layer]
client_weights = get_weight_obj(client_layer[0].eval(sess),
client_layer[1].eval(sess),
client_layer[2].eval(sess))
client_weight_list['B'].append(client_weights['B'])
client_weight_list['W'].append(client_weights['W'])
client_weight_list['theta'].append(client_weights['theta'])
new_weights = {}
new_weights['B'] = tf.convert_to_tensor(
np.mean(client_weight_list['B'], axis=0))
new_weights['W'] = tf.convert_to_tensor(
np.mean(client_weight_list['W'], axis=0))
new_weights['theta'] = tf.convert_to_tensor(
np.mean(client_weight_list['theta'], axis=0))
new_global_model.set_weights_at_layer(new_weights, layer, sess)
lnmse2 = run_sc_test_with_created_session(config, sess,
new_global_model)
np.savez(
'lnmse_new_global' + str(num_clients) + "_" +
str(config.maxit), lnmse2)
print("New Global model performance")
print(lnmse2)
save_trainable_variables(sess, config.modelfn, config.scope)
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, False)
"""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)
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)))
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))
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)))
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)
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()