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)