def test_training(data, loss):
train_gen = data.batch_generator(batch_size=5)
eval_gen = data.batch_generator(batch_size=5)
test_gen = data.batch_generator(batch_size=2, subset='test')
num_epochs = 2
box_size = data.box_size
input_channels = data.x_channels
output_channels = data.y_channels
inputs = Input((box_size, box_size, box_size, input_channels))
outputs = Convolution3D(filters=output_channels,
kernel_size=1,
activation='sigmoid')(inputs)
model = UNet(inputs=inputs, outputs=outputs)
model.compile(optimizer=Adam(lr=1e-6),
loss=loss,
metrics=[dice, dice_loss, ovl, ovl_loss])
model.fit_generator(train_gen,
steps_per_epoch=2,
epochs=num_epochs,
verbose=0)
for scores in (model.evaluate_generator(eval_gen, steps=2),
model.evaluate_generator(test_gen, steps=1)):
assert np.allclose(scores[1], -scores[2])
assert np.allclose(scores[3], -scores[4])
loss_change = model.history.history['loss']
assert len(loss_change) == num_epochs
assert (loss_change[0] != loss_change[1:]).all()
def main():
args = parse_args()
if args.output is None:
args.output = 'output_' + time.strftime('%Y-%m-%d')
if not os.path.exists(args.output):
os.makedirs(args.output)
if not os.access(args.output, os.W_OK):
raise IOError(
'Cannot create files inside %s (check your permissions).' %
args.output)
if args.train_ids:
with open(args.train_ids) as f:
train_ids = list(filter(None, f.read().split('\n')))
else:
train_ids = None
if args.test_ids:
with open(args.test_ids) as f:
test_ids = list(filter(None, f.read().split('\n')))
else:
test_ids = None
if train_ids:
if test_ids:
all_ids = sorted(set(train_ids) | set(test_ids))
else:
all_ids = train_ids
else:
all_ids = None
data = DataWrapper(args.input,
test_set=test_ids,
pdbids=all_ids,
load_data=args.load)
if args.model:
model = UNet.load_model(args.model, data_handle=data)
else:
model = UNet(data_handle=data)
model.compile(optimizer=Adam(lr=1e-6),
loss=dice_loss,
metrics=[dice, ovl, 'binary_crossentropy'])
train_batch_generator = data.batch_generator(batch_size=args.batch_size)
callbacks = [
ModelCheckpoint(os.path.join(args.output, 'checkpoint.hdf'),
save_best_only=False)
]
if test_ids:
val_batch_generator = data.batch_generator(batch_size=args.batch_size,
subset='test')
num_val_steps = max(args.steps_per_epoch // 5, 1)
callbacks.append(
ModelCheckpoint(os.path.join(args.output, 'best_weights.hdf'),
save_best_only=True))
else:
val_batch_generator = None
num_val_steps = None
model.fit_generator(train_batch_generator,
steps_per_epoch=args.steps_per_epoch,
epochs=args.epochs,
verbose=args.verbose,
callbacks=callbacks,
validation_data=val_batch_generator,
validation_steps=num_val_steps)
history = pd.DataFrame(model.history.history)
history.to_csv(os.path.join(args.output, 'history.csv'))
model.save(os.path.join(args.output, 'model.hdf'))