config.cycle_consistency_loss, config.cycle_consistency_loss,
config.id_loss, config.id_loss
],
metrics=[utils.ssim])
def scheduler(epoch):
if epoch < config.startLRdecay:
return 2e-4
else:
epochs_passed = epoch - config.startLRdecay
decay_step = 2e-4 / (config.epochs - config.startLRdecay)
return 2e-4 - epochs_remaining * decay_step
LRscheduler = callbacks.MultiLRScheduler(scheduler, training_models=[model.d_A, model.d_B, model.combined])
# Generate Callbacks
tensorboard = tf.keras.callbacks.TensorBoard(log_dir=LOG_DIR, write_graph=True, update_freq='epoch')
start_tensorboard = callbacks.StartTensorBoard(LOG_DIR)
prog_bar = tf.keras.callbacks.ProgbarLogger(count_mode='steps', stateful_metrics=None)
log_code = callbacks.LogCode(LOG_DIR, './trainer')
copy_keras = callbacks.CopyKerasModel(MODEL_DIR, LOG_DIR)
saving = callbacks.MultiModelCheckpoint(MODEL_DIR + '/model.{epoch:02d}-{val_ssim:.10f}.hdf5',
monitor='val_ssim', verbose=1, freq='epoch', mode='max', save_best_only=False,
save_weights_only=True,
multi_models=[('g_AB', g_AB), ('g_BA', g_BA), ('d_A', d_A), ('d_B', d_B)])
restore_best_weights=True, verbose=1)
image_gen = callbacks.GenerateImages(g_AB, test_X, test_Y, LOG_DIR, interval=int(dataset_count/config.bs))
# Fit the model
model.fit(train_X, train_Y,
write_graph=True,
update_freq=write_freq)
saving = tf.keras.callbacks.ModelCheckpoint(
config.model_dir + '/d' + '/model.{epoch:02d}-{val_loss:.5f}.hdf5',
monitor='val_loss',
verbose=1,
save_freq='epoch',
save_best_only=False)
log_code = callbacks.LogCode(config.job_dir, './trainer')
#copy_keras = callbacks.CopyKerasModel(config.model_dir, config.job_dir)
#image_gen_val = callbacks.GenerateImages(generator_model, validation_dataset, config.job_dir, interval=write_freq, postfix='val')
#image_gen = callbacks.GenerateImages(generator_model, train_dataset, config.job_dir, interval=write_freq, postfix='train')
start_tensorboard = callbacks.StartTensorBoard(config.job_dir)
# Fit model
d_model.fit(
train_dataset,
steps_per_epoch=int(train_count / config.bs),
epochs=config.epochs,
validation_data=validation_dataset,
validation_steps=int(validation_count / config.bs),
verbose=1,
callbacks=[
log_code,
start_tensorboard,
tensorboard,
#image_gen,
#image_gen_val,