history = my_callback.Histories()
callbacks = [checkpoint, early_stopping, history]
## *********************** Train settings **************************************
steps_train = 5000
steps_val = 1700
epochs = 1000
## ************************** Model fit ****************************************
model.fit_generator(
train_generator,
steps_per_epoch=steps_train,
validation_data=val_generator,
validation_steps=steps_val,
epochs=epochs,
verbose=2,
workers=16,
use_multiprocessing=True,
callbacks=callbacks)
## ************************* save curves ***************************************
history_dictionary_loss = history.loss
history_dictionary_val_loss = history.val_loss
np.save(FOLDER+'loss.npy', history_dictionary_loss)
np.save(FOLDER+'val_loss.npy', history_dictionary_val_loss)
print('Ending phase 1')
## ------------- End phase 1 - with batch normalization -----------------------
## ---------- Phase 2 - batch normalization off - lr 5e-5 ---------------------
# Instantiate the model
model = PConvUnet(vgg_weights='./data/pytorch_vgg16.h5')
FOLDER = './data/logs/C2/'
# Run training for certain amount of epochs
model.fit_generator(
train_generator,
steps_per_epoch=500,
validation_data=val_generator,
validation_steps=50,
epochs=50,
verbose=0,
callbacks=[
TensorBoard(
log_dir=FOLDER,
write_graph=True
),
ModelCheckpoint(
FOLDER+'weights.{epoch:02d}-{loss:.2f}.h5',
monitor='val_loss',
save_best_only=True,
save_weights_only=True
),
LambdaCallback(
on_epoch_end=lambda epoch, logs: plot_callback(model)
),
TQDMNotebookCallback()
]
)
model.load(args.checkpoint)
elif args.stage == 'finetune':
model.load(args.checkpoint, train_bn=False, lr=0.00005)
# Fit model
model.fit_generator(
train_generator,
steps_per_epoch=10000,
validation_data=val_generator,
validation_steps=1000,
epochs=args.epochs,
verbose=0,
callbacks=[
TensorBoard(
log_dir=os.path.join(args.log_path, args.name+'_'+args.stage),
write_graph=False
),
ModelCheckpoint(
os.path.join(args.log_path, args.name+'_'+args.stage, 'weights.{epoch:02d}-{loss:.2f}.h5'),
monitor='val_loss',
save_best_only=True,
save_weights_only=True
),
LambdaCallback(
on_epoch_end=lambda epoch, logs: plot_callback(model, args.test_path)
),
TQDMCallback()
]
)