K.set_session(utils.create_session())
mnist_dataset = MnistWrapper.load_default()
# ----------normal relu pretraining----------
print 'Training model with normal relu'
folder = 'test/mnist_ae_relu_inf'
ae = AutoEncoder(mnist_dataset, encode, decode, None, folder)
ae.build_models()
num_epoch = 30
lr_schedule = utils.generate_decay_lr_schedule(num_epoch, 0.1, 1)
ae.train(128, num_epoch, lr_schedule)
ae.save_models()
ae.test_models(utils.vis_mnist)
ae.log()
encoded_dataset = ae.encode(MnistWrapper)
encoded_dataset.dump_to_h5(os.path.join(folder, 'encoded_mnist.h5'))
encoded_dataset.plot_data_dist(os.path.join(folder, 'encoded_plot.png'))
# ----------truncate relu and fine-tune----------
print 'Training model with relu-%d' % RELU_MAX
new_folder = 'test/mnist_ae_relu_%d' % RELU_MAX
ae = AutoEncoder(mnist_dataset, encode, decode, RELU_MAX, new_folder)
ae.build_models(folder) # load previously trained ae
mnist_dataset.plot_data_dist(os.path.join(new_folder, 'original_mnist_plot.png'))
# num_epoch = 2
# lr_schedule = utils.generate_decay_lr_schedule(num_epoch, 0.1, 1)
