def train(self): self.ckpt.restore(self.ckpt_manager.latest_checkpoint) if self.ckpt_manager.latest_checkpoint: print("Restored from {}".format( self.ckpt_manager.latest_checkpoint)) else: print("Initializing from scratch.") dg = DataGenerator(task_version=self.task_version, action='train') validation_batch_num = self.batch_num // 10 over_all_performace = [] print('Start to train') print('#' * 20) for epoch_i in range(self.epoch_num): print('Epoch ' + str(epoch_i)) # Train # train_loss_all = 0 for batch_i in range(self.batch_num): decs, masks, inputs, outputs = next(dg) with tf.GradientTape() as tape: logits, _ = self.ei_rnn(inputs, [self.init_state], training=True) logits = tf.transpose(logits, perm=[0, 2, 1]) train_loss = self.loss_fun(outputs, logits, masks) train_loss += sum(self.ei_rnn.losses) train_loss_all += train_loss.numpy() all_weights = self.ei_rnn.trainable_weights + [self.init_state] grads = tape.gradient(train_loss, all_weights) grads, _ = tf.clip_by_global_norm(grads, self.grad_clip) self.optimizer.apply_gradients( zip(grads, self.ei_rnn.trainable_weights)) train_loss_all = train_loss_all / self.batch_num print('train loss:', train_loss_all) with self.train_summary_writer.as_default(): tf.summary.scalar('loss', train_loss_all, step=epoch_i) # Validation # validation_loss_all = 0 validation_acc_all = 0 for v_batch_i in range(validation_batch_num): _, v_masks, v_inputs, v_outputs = dg.get_valid_test_datasets( ) # dg.get_valid_test_datasets() v_logits, _ = self.ei_rnn(v_inputs, [self.init_state]) acc_v_logits = v_logits.numpy() acc_v_outputs = tf.transpose(v_outputs, perm=[0, 2, 1]).numpy() validation_acc = self.get_accuracy(acc_v_logits, acc_v_outputs, v_masks) validation_acc_all += validation_acc v_logits = tf.transpose(v_logits, perm=[0, 2, 1]) validation_loss = self.loss_fun(v_outputs, v_logits, v_masks) validation_loss += sum(self.ei_rnn.losses) validation_loss_all += validation_loss.numpy() validation_loss_all = validation_loss_all / validation_batch_num validation_acc_all = validation_acc_all / validation_batch_num over_all_performace.append(validation_acc_all) print('validation loss:', validation_loss_all) print('validation acc:', validation_acc_all) with self.validation_summary_writer.as_default(): tf.summary.scalar('loss', validation_loss_all, step=epoch_i) tf.summary.scalar('acc', validation_acc_all, step=epoch_i) cm_image = plot.plot_confusion_matrix(self.get_w_rec_m()) tf.summary.image('M_rec', cm_image, step=epoch_i) win_image = plot.plot_confusion_matrix( funs.rectify( self.rnn_cell.W_in.numpy()[:, :int(UNITS_SIZE * EI_RATIO)]), False) tf.summary.image('M_in', win_image, step=epoch_i) wout_image = plot.plot_confusion_matrix( self.get_w_out_m()[:, :int(UNITS_SIZE * EI_RATIO)], False) tf.summary.image('M_out', wout_image, step=epoch_i) print('spr: ', funs.spectral_radius(self.get_w_rec_m().T)) if epoch_i > PERFORMANCE_CHECK_REGION and \ np.mean(over_all_performace[-PERFORMANCE_CHECK_REGION:]) > PERFORMANCE_LEVEL: print( 'Overall performance level is satisfied, training is terminated\n' ) break # Save Model self.ckpt.step.assign_add(1) self.ckpt_manager.save() # self.reset_all_weights() # todo: may uncomment # print('Remove all weights below ' + str(SGD_p['mini_w_threshold'])) # print('\n') print('Training is done') print('#' * 20) print('\n') # Test # self.test()
def test(self, test_batch_num=50): print('Start to test') print('#' * 20) self.ckpt.restore(self.ckpt_manager.latest_checkpoint) if self.ckpt_manager.latest_checkpoint: print("Restored from {}".format( self.ckpt_manager.latest_checkpoint)) else: print("Initializing from scratch.") dg = DataGenerator(task_version=self.task_version, action='test') # todo: should be test for action psycollection = {'coh': [], 'perc': []} for batch_index in range(20): descs, test_masks, test_inputs, test_outputs = dg.get_valid_test_datasets( ) test_logits, _ = self.ei_rnn(test_inputs, [self.init_state], training=False) acc_test_logits = test_logits.numpy() acc_test_outputs = tf.transpose(test_outputs, perm=[0, 2, 1]).numpy() test_acc = self.get_accuracy(acc_test_logits, acc_test_outputs, test_masks) test_logits = tf.transpose(test_logits, perm=[0, 2, 1]) test_loss = self.loss_fun(test_outputs, test_logits, test_masks) test_loss += sum(self.ei_rnn.losses) print('test loss:', test_loss.numpy()) print('test acc:', test_acc) tmp_data = self.get_psychometric_data(descs, test_logits.numpy()) psycollection['coh'] += tmp_data['coh'] psycollection['perc'] += tmp_data['perc'] with self.test_summary_writer.as_default(): tf.summary.scalar('loss', test_loss, step=batch_index) tf.summary.scalar('acc', test_acc, step=batch_index) curve_image = plot.plot_dots(psycollection['coh'], psycollection['perc']) tf.summary.image('psycollection', curve_image, step=batch_index) cm_image = plot.plot_confusion_matrix(self.get_w_rec_m()) tf.summary.image('M_rec', cm_image, step=batch_index) win_image = plot.plot_confusion_matrix( funs.rectify( self.rnn_cell.W_in.numpy()[:, :int(UNITS_SIZE * EI_RATIO)]), False) tf.summary.image('M_in', win_image, step=batch_index) wout_image = plot.plot_confusion_matrix( self.get_w_out_m()[:, :int(UNITS_SIZE * EI_RATIO)], False) tf.summary.image('M_out', wout_image, step=batch_index)