###################################################################
# hyper-parameters
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_index
epochs = FLAGS.epochs
eval_freq = FLAGS.eval_frequency
batch_size = FLAGS.batch_size
batch_steps = int(n_train_imgs // batch_size)
# create networks
Models = CNNModels(FLAGS, dataset)
model = Models.Classifier()
# log network
logger.callback.set_model(model) # record in tensorboard
logger.save_model(model, "Model")
logger.load_pretrained_weights(model, 'Model_weights.h5')
logger.plot_model(model, 'Classifier')
log('Create Networks Successfully')
#%%
###################################################################
######################## Start training ###########################
###################################################################
for epoch in range(epochs + 1):
log("Epoch {}".format(epoch + 1))
# train
loss_sum, acc_sum = 0, 0
for i in range(batch_steps):
x_batch, y_batch, showlabels = dataset.train_next_batch()
img_ch = 1
img_height, img_width = inputsize[0], inputsize[1]
inputs = Input((img_height, img_width, img_ch))
x = Conv2D(64, kernel_size=(9, 9), padding='valid')(inputs)
x = Activation('relu')(x)
x = Conv2D(32, kernel_size=(1, 1), padding='valid')(x)
x = Activation('relu')(x)
x = Conv2D(1, kernel_size=(5, 5), padding='valid')(x)
x = Activation('relu')(x)
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)
