def train_model(model: SLModel,
trainset: NpDataset,
valset: NpDataset,
epochs=5,
batch_size=32):
# Create the generators
logging.info("Training model for {} epochs and {} batch size".format(
epochs, batch_size))
logging.info("Flowing the train and validation sets")
traingen = trainset.flow(
batch_size=batch_size, shuffle=True, seed=utils.get_random_seed())
valgen = valset.flow(batch_size=batch_size, shuffle=False)
# Create the callbacks
logging.info("Creating the callbacks")
callbacks = [
ModelCheckpoint(
utils.get_model_path(RUN_ID),
"val_loss",
verbose=1,
save_best_only=True),
Plotter(
"loss",
scale='log',
plot_during_train=True,
save_to_file=utils.get_plot_path(RUN_ID),
block_on_end=False),
Plotter(
"accuracy",
scale='linear',
plot_during_train=True,
save_to_file=utils.get_plot_path(RUN_ID + "_acc"),
block_on_end=False)
]
# Create the optiizer
logging.info("Creating the optimizer")
params = [param for param in model.parameters() if param.requires_grad]
# optimizer = optim.SGD(
# params
# lr=0.01,
# momentum=0.9,
# nesterov=True)
optimizer = optim.Adam(params)
logging.info("Optimizer: %r" % optimizer)
# Train the model
logs = model.fit_generator(
traingen,
traingen.steps_per_epoch,
epochs=epochs,
optimizer=optimizer,
validation_generator=valgen,
validation_steps=valgen.steps_per_epoch,
metrics=["accuracy"],
callbacks=callbacks,
verbose=1)
return logs
def train_model(model: SLModel,
trainset: NpDataset,
valset: NpDataset,
epochs=EPOCHS,
batch_size=BATCH_SIZE,
plot=True):
# Create the generators
logging.info("Training model for {epochs} epochs and {batch_size} batch "
"size".format(**locals()))
logging.info("Flowing the train and validation sets")
traingen = trainset.flow(
batch_size=batch_size, shuffle=True, seed=utils.get_random_seed())
valgen = valset.flow(batch_size=batch_size, shuffle=False)
# Create the callbacks
logging.info("Creating the callbacks")
callbacks = [
ModelCheckpoint(
utils.get_model_path(RUN_ID),
"val_loss",
verbose=1,
save_best_only=True),
Plotter(
"bce",
scale='log',
plot_during_train=plot,
save_to_file=utils.get_plot_path(RUN_ID+"_bce"),
block_on_end=False),
Plotter(
"dice",
plot_during_train=plot,
save_to_file=utils.get_plot_path(RUN_ID+"_dice"),
block_on_end=False),
Plotter(
"iou",
plot_during_train=plot,
save_to_file=utils.get_plot_path(RUN_ID + "_iou"),
block_on_end=False),
]
# Train the model
logs = model.fit_generator(
traingen,
traingen.steps_per_epoch,
epochs=epochs,
validation_data=valgen,
validation_steps=valgen.steps_per_epoch,
callbacks=callbacks,
metrics=["iou", mean_iou],
verbose=1)
return logs
def test_model(model: SLModel, test_data: NpDataset, batch_size=32):
logging.info("Testing model with batch size of {}".format(batch_size))
logging.info("Flowing the test set")
testgen = test_data.flow(batch_size=batch_size, shuffle=False)
test_preds = model.predict_generator(
testgen, testgen.steps_per_epoch, verbose=1)
return test_preds[:, 0]
def test_model(model: SLModel, test_data: NpDataset, batch_size=BATCH_SIZE):
logging.info("Testing model with batch size of {batch_size}".format(**locals()))
logging.info("Flowing the test set")
test_data.output_labels = False
testgen = test_data.flow(batch_size=batch_size, shuffle=False)
test_preds = model.predict_generator(
testgen, testgen.steps_per_epoch, verbose=1)
return test_preds.squeeze(-1)
def validate_model(model: SLModel, val_data: NpDataset, batch_size=32):
logging.info("Validating model with batch size of {}".format(batch_size))
val_data.output_labels = False
logging.info("Flowing the validation set")
valgen = val_data.flow(batch_size=batch_size, shuffle=False)
logging.info("Getting validation predictions")
val_preds = model.predict_generator(valgen, valgen.steps_per_epoch)
score = roc_auc_score(val_data.y[:, 0], val_preds[:, 0])
logging.info("Validation ROC AUC score: {}".format(score))
return score
def test_augmenter_basic():
# Try different combinations of with labels and without
data = NpDataset(x=np.ones((32, )), y=np.ones((32, )))
augmenter = ZeroAugmenter(labels=False, augment_labels=False)
assert not augmenter.labels
assert not augmenter.augment_labels
data.output_labels = False
x = next(augmenter(data.flow(batch_size=32)))
assert np.all(x == 0.)
augmenter = ZeroAugmenter(labels=False, augment_labels=True)
assert not augmenter.labels
assert augmenter.augment_labels
data.output_labels = False
x = next(augmenter(data.flow(batch_size=32)))
assert np.all(x == 0.)
augmenter = ZeroAugmenter(labels=True, augment_labels=False)
assert augmenter.labels
assert not augmenter.augment_labels
data.output_labels = True
x, y = next(augmenter(data.flow(batch_size=32)))
assert np.all(x == 0.)
assert np.all(y == 1.)
augmenter = ZeroAugmenter(labels=True, augment_labels=True)
assert augmenter.labels
assert augmenter.augment_labels
data.output_labels = True
x, y = next(augmenter(data.flow(batch_size=32)))
assert np.all(x == 0.)
assert np.all(y == 0.)
# Try a generic python generator
def datagen():
yield np.ones((32, ))
augmenter = ZeroAugmenter(labels=False, augment_labels=False)
x = next(augmenter(datagen()))
assert np.all(x == 0.)
def train_model(model,
trainset: NpDataset,
valset: NpDataset,
epochs=70,
batch_size=32,
val_batch_size=32,
plot=True,
run_id='default_model_name',
augmenter=None,
verbose=1,
debug=False):
# Create the generators
logger.info(
f'Training model for {epochs} epochs and {batch_size} batch size')
logger.info('Flowing the train and validation sets')
traingen = trainset.flow(batch_size=batch_size,
shuffle=True,
seed=utils.get_random_seed())
valgen = valset.flow(batch_size=val_batch_size, shuffle=False)
if augmenter is not None:
logger.info(f'Training with augmenter {augmenter.image_augmenter}')
augmenter.labels = True
traingen = augmenter(traingen)
# Create the callbacks
logger.info('Creating the callbacks')
callbacks = [
ModelCheckpoint(utils.get_model_path(run_id),
'val_loss',
verbose=verbose,
save_best_only=True,
save_weights_only=True),
ModelCheckpoint(utils.get_model_path(run_id + '_dice_coef'),
'val_dice_coef',
verbose=verbose,
save_best_only=True,
save_weights_only=True,
mode='max'),
Plotter('loss',
scale='log',
plot_during_train=plot,
save_to_file=utils.get_plot_path(run_id + '_loss'),
block_on_end=False),
Plotter('dice_coef',
plot_during_train=plot,
save_to_file=utils.get_plot_path(run_id + '_dice_coef'),
block_on_end=False),
]
train_steps = 3 if debug else traingen.steps_per_epoch
val_steps = 3 if debug else valgen.steps_per_epoch
epochs = 2 if debug else epochs
# Train the model
logs = model.fit_generator(traingen,
train_steps,
epochs=epochs,
validation_data=valgen,
validation_steps=val_steps,
callbacks=callbacks,
verbose=verbose,
max_queue_size=3)
return logs
def test_model(model,
test_data: NpDataset,
batch_size=32,
num_augmentations=0,
view_preds=False,
debug=False):
logger.info(f'Testing model with batch size of {batch_size}')
logger.info('Flowing the test set')
test_data.output_labels = False
testgen = test_data.flow(batch_size=batch_size, shuffle=False)
if num_augmentations:
print(f'Testing with a flip augmenter')
augmenter = FlipAugmenter(flipud=True, fliplr=True)
aug_params = [
dict(flipud=True, fliplr=True),
dict(flipud=True, fliplr=False),
dict(flipud=False, fliplr=True),
dict(flipud=False, fliplr=False)
]
augmenter.labels = False
testgen = augmenter(testgen)
else:
num_augmentations = 1
augmenter = None
test_steps = 3 if debug else testgen.steps_per_epoch
test_preds = 0.
for i in range(num_augmentations):
if augmenter is not None:
print(
f'Testing for augmentation {i+1}/{num_augmentations} with flipud={aug_params[i]["flipud"]} and fliplr={aug_params[i]["fliplr"]}'
)
augmenter.flipud = aug_params[i]['flipud']
augmenter.fliplr = aug_params[i]['fliplr']
aug_test_preds = model.predict_generator(
testgen, test_steps, verbose=1, max_queue_size=0, workers=0
) # Must set to workers=0 to maintain test prediction order
# Reverse the augmentations
# TODO: only works with flips, implement general solution for non-flips
if augmenter is not None:
print('Running reverse augmentation on predictions...')
aug_test_preds = augmenter.reverse_augment(aug_test_preds)
if view_preds:
if augmenter:
testgen.generator.restart()
display_predictions(testgen.generator, aug_test_preds)
else:
display_predictions(testgen, aug_test_preds)
test_preds = test_preds + aug_test_preds
test_preds /= num_augmentations
if debug:
filler = np.zeros(
(len(test_data) - len(test_preds), *test_preds.shape[1:]))
test_preds = np.concatenate([test_preds, filler])
if view_preds:
display_predictions(testgen, test_preds)
return test_preds.squeeze(-1)