SRCNN = Model(inputs=inputs, outputs=x, name='SRCNN')
SRCNN.compile(Adam(lr=0.0002), loss='mse')
#%% log model
logger.callback.set_model(SRCNN) # record in tensorboard
logger.save_model(SRCNN, "SRCNN_model")
logger.load_pretrained_weights(SRCNN, 'SRCNN_Weights.h5')
logger.plot_model(SRCNN, 'SRCNN')
log('Create Networks Successfully')
###################################################################
######################## Start training ###########################
###################################################################
batch_steps = n_images // FLAGS.batch_size
for epoch in range(FLAGS.epochs):
log("Epoch {}".format(epoch + 1))
testindex = np.random.choice(len(dataset.test_names), 1)
testname = os.path.join(dataset.test_dir, dataset.test_names[testindex[0]])
img_ycbcr = np.squeeze(dataset.imagefiles2arrs([testname]))
utils.predict(img_ycbcr / 255.0, SRCNN, epoch)
loss_sum = 0
for i in range(batch_steps):
x_batch, y_batch = dataset.train_next_batch()
loss = SRCNN.train_on_batch(x_batch, y_batch)
loss_sum += loss
utils.progressbar(i, batch_steps)
logger.write_tensorboard(['train_loss'], [loss_sum / i], epoch)
logger.save_weights(SRCNN, "SRCNN_Weights")
