def train_VAE():
model = VAE()
model.summary()
history = []
x_train = import_image(folder='train')
x_test = import_image(folder='test')
#x_train = import_image('test')
arbitrary = np.zeros([x_train.shape[0], 1024 * 2])
model.summary()
history = model.fit(x_train, [x_train, arbitrary],
epochs=200,
batch_size=128,
shuffle=True,
verbose=1,
validation_data=(x_test,
[x_test,
arbitrary[:x_test.shape[0]]]))
print("history =", history.history.keys())
output_loss1 = history.history['recons_loss']
output_loss2 = history.history['KLD_loss']
np.save('VAE2_recons_loss', output_loss1)
np.save('VAE2_KLD_loss', output_loss2)
model.save_weights('saved_model/VAE2_epochs60_weights.h5')
# Define MNIST encoder / decoder
encoder = tfk.Sequential([
tfkl.Conv2D(filters=16, kernel_size=3, strides=2, padding='same', activation='relu'),
tfkl.Conv2D(filters=32, kernel_size=3, strides=2, padding='same', activation='relu'),
tfkl.Flatten(),
tfkl.Dense(units=2 * 10), # no activation
])
decoder = tfk.Sequential([
tfkl.Dense(units=7*7*32, activation='relu'),
tfkl.Reshape(target_shape=(7, 7, 32)),
tfkl.Conv2DTranspose(filters=16, kernel_size=3, strides=2, padding='same', activation='relu'),
tfkl.Conv2DTranspose(filters=1, kernel_size=3, strides=2, padding='same'), # no activation
])
# Define model
model = VAE(encoder, decoder, args.kl_weight)
model.compile(optimizer='adam', loss='mse')
# Callbacks
time = datetime.now().strftime('%Y%m%d-%H%M%S')
log_dir = os.path.join('.', 'logs', 'vae', time)
tensorboard_clbk = tfk.callbacks.TensorBoard(log_dir=log_dir)
plot_clbk = PlotReconstructionCallback(logdir=log_dir, test_ds=test_ds, nex=4)
samples_clbk = PlotSamplesCallback(logdir=log_dir, nex=4)
callbacks = [tensorboard_clbk, plot_clbk, samples_clbk]
# Fit
model.fit(train_ds, validation_data=test_ds, epochs=EPOCHS, callbacks=callbacks)
