def train(self):
# Placeholders for Tensorboard
adv_tensor = tf.placeholder(tf.float32, shape = [], name = "Advantage")
test_rmse = tf.placeholder(tf.float32, shape = [], name = "Test_RMSE")
test_ssim = tf.placeholder(tf.float32, shape = [], name = "Test_SSIM")
test_psnr = tf.placeholder(tf.float32, shape = [], name = "Test_psnr")
closs = tf.placeholder(tf.float32, shape = [], name = "Controller_Loss")
im_result = tf.placeholder(tf.uint8, shape = [1, 8 * self.image_shape[0], 3 * self.image_shape[1], self.image_shape[2]], name = "Input_Placeholder")
# Histogram summaries
softmax_histograms, tensors = softmax_summaries()
# Summaries
summ_rmse = tf.summary.scalar('RMSE', test_rmse)
summ_ssim = tf.summary.scalar('SSIM', test_ssim)
summ_psnr = tf.summary.scalar('PSNR', test_psnr)
summ_rew = tf.summary.scalar('Advantage', adv_tensor)
summ_closs = tf.summary.scalar('Controller_Loss', closs)
summ_denoising = tf.summary.image('Denoising_Results', im_result)
merged = tf.summary.merge([summ_rmse, summ_psnr, summ_ssim, summ_rew])
tf.global_variables_initializer().run(session = self.session)
self.writer = tf.summary.FileWriter(os.path.join(LOG_DIR, "training_monitor"), graph = self.session.graph)
# Learning loop
mem_softmaxes = [] # Softmax memory
mem_rewards = [] # Reward memory
mem_advantage = [] # Advantage memory
old_softmax = np.zeros([1, *INPUT_SHAPE])
print("Starting Controller training")
for epoch in range(self.n_epochs):
start_epoch = time.time()
print("Epoch\t{}/{}".format(epoch, self.n_epochs) + 149 * "-")
# Policy prediction
softmaxes, subpolicies = self.sample_policy(old_softmax)
# Updates old softmax
old_softmax = np.expand_dims(np.concatenate(softmaxes[-6:], axis = 1), axis = -1)
# Softmax memory
mem_softmaxes.append(softmaxes)
# Creates new child and trains it
child = Autoencoder(self.dataset, n_epochs = 250, n_stages = 20, test_runs = 50, lr = 0.02, decay = 0.1)
train_start = time.time()
child.train(policy = subpolicies, do_valid = False)
elapsed = time.time() - train_start
print("Training time:\t{}".format(timedelta(seconds = elapsed)))
# Logging training data into tensorboard
train_start = time.time()
loss, ssim, rmse, psnr = child.test() # Evaluates training results
elapsed = time.time() - train_start
reward = psnr # Epoch reward
mem_rewards.append(reward) # Keeps a memory of rewards
baseline = get_baseline(mem_rewards) # Simple moving average
advantage = reward - baseline # Computes advantage
mem_advantage.append(advantage) # Keeps a memory of advantages
print("Eval time:\t{}".format(timedelta(seconds = elapsed)))
# Saving child model
self.saver(child.model, subpolicies, mem_rewards[-1])
# Weight update
if len(mem_softmaxes) > 5:
# Perform backprop only if we have enough softmaxes in memory
controller_loss = self.fit(mem_softmaxes, mem_advantage, epoch)
self.model.save_weights(os.path.join(LOG_DIR, "controller_model/model.h5"))
closs_summary = self.session.run(summ_closs, feed_dict = {closs : controller_loss})
self.writer.add_summary(closs_summary, epoch)
# Logging data into tensorboard
# Scalar summaries
feed_dict = {
test_rmse : rmse,
test_ssim : 1 - ssim,
test_psnr : psnr,
adv_tensor : advantage
}
summ = self.session.run(merged, feed_dict = feed_dict)
self.writer.add_summary(summ, epoch)
print("SSIM: \t\t{}".format(1 - ssim))
print("RMSE: \t\t{}".format(rmse))
print("PSNR: \t\t{}".format(psnr))
print("Reward: \t{}".format(mem_rewards[-1]))
# Images and histograms
if epoch % 10 == 0:
# Image summary
x, y, r = child.inference()
# Converts 8 first samples to uint8
noise = np.uint8(255 * x[0:8])
original = np.uint8(255 * y[0:8])
denoised = np.uint8(255 * r[0:8])
# Stacking images per batch
a = np.vstack([noise[i] for i in range(len(noise))])
b = np.vstack([original[i] for i in range(len(original))])
c = np.vstack([denoised[i] for i in range(len(denoised))])
# Stacking batches side-by-side
img = np.expand_dims(np.hstack([a, b, c]), 0)
img_summ = self.session.run(summ_denoising, feed_dict = {im_result: img})
self.writer.add_summary(img_summ, epoch)
# Histograms
type_softmaxes = [softmaxes[0], softmaxes[3]]
prob_softmaxes = [softmaxes[1], softmaxes[4]]
magn_softmaxes = [softmaxes[2], softmaxes[5]]
feed_hist = feed_histogram_summaries(*tensors, type_softmaxes, prob_softmaxes, magn_softmaxes)
soft_summ = self.session.run(softmax_histograms, feed_dict = feed_hist)
self.writer.add_summary(soft_summ, epoch)
print("Epoch took:\t{}".format(timedelta(seconds = time.time() - start_epoch)))
weights_ = model.get_weights()
for idx, w in weights_.items():
with open("results/{2}/weight_{1}_{0}.pickle".format("_".join([str(x) for x in params['network_shape'][1:-1]]), idx, MODEL_NAME), 'wb') as f:
data_w = np.sum(np.abs(w), 1)
pickle.dump(data_w, f)
activation_dicts = {k: {} for k in params['activation_pattern_layers']}
sum_of_act_dicts = {k: [] for k in params['activation_pattern_layers']}
for epoch in range(inference_epochs):
total_batch=int(n_test_samples / batch_size)
# Loop over all batches
for i in range(total_batch):
batch_xs, _ = get_random_block_from_data(X_test, Y_test, batch_size)
cost = model.inference(batch_xs)
activation_patterns = model.get_activations(batch_xs)
for j, act_patt in activation_patterns.items():
current_act_dict = activation_dicts[j]
sum_of_act_dicts[j] += np.mean(act_patt, 1).tolist()
for single_act_patt in act_patt:
nonzero_tup = tuple(np.nonzero(single_act_patt)[0])
if nonzero_tup in current_act_dict:
current_act_dict[nonzero_tup] += 1
else:
current_act_dict[nonzero_tup] = 1
activation_dicts[j] = current_act_dict
with open("results/{1}/sum_act_{0}.pickle".format("_".join([str(x) for x in params['network_shape'][1:-1]]), MODEL_NAME), 'wb') as sum_act_f:
pickle.dump(sum_of_act_dicts, sum_act_f)
