def _create_trainer(config, model, optimizer, lr_scheduler, loss_criterion,
eval_criterion, loaders):
assert 'trainer' in config, 'Could not find trainer configuration'
trainer_config = config['trainer']
resume = trainer_config.get('resume', None)
pre_trained = trainer_config.get('pre_trained', None)
# get tensorboard formatter
tensorboard_formatter = get_tensorboard_formatter(
trainer_config.pop('tensorboard_formatter', None))
# get sample plotter
sample_plotter = create_sample_plotter(
trainer_config.pop('sample_plotter', None))
if resume is not None:
# continue training from a given checkpoint
return UNet3DTrainer.from_checkpoint(
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
loss_criterion=loss_criterion,
eval_criterion=eval_criterion,
loaders=loaders,
tensorboard_formatter=tensorboard_formatter,
sample_plotter=sample_plotter,
**trainer_config)
elif pre_trained is not None:
# fine-tune a given pre-trained model
return UNet3DTrainer.from_pretrained(
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
loss_criterion=loss_criterion,
eval_criterion=eval_criterion,
tensorboard_formatter=tensorboard_formatter,
sample_plotter=sample_plotter,
device=config['device'],
loaders=loaders,
**trainer_config)
else:
# start training from scratch
return UNet3DTrainer(model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
loss_criterion=loss_criterion,
eval_criterion=eval_criterion,
device=config['device'],
loaders=loaders,
tensorboard_formatter=tensorboard_formatter,
sample_plotter=sample_plotter,
**trainer_config)
def create_trainer(config):
# Create the model
model = get_model(config['model'])
# use DataParallel if more than 1 GPU available
device = config['device']
if torch.cuda.device_count() > 1 and not device.type == 'cpu':
model = nn.DataParallel(model)
logger.info(f'Using {torch.cuda.device_count()} GPUs for training')
# put the model on GPUs
logger.info(f"Sending the model to '{config['device']}'")
model = model.to(device)
# Log the number of learnable parameters
logger.info(
f'Number of learnable params {get_number_of_learnable_parameters(model)}'
)
# Create loss criterion
loss_criterion = get_loss_criterion(config)
# Create evaluation metric
eval_criterion = get_evaluation_metric(config)
# Create data loaders
loaders = get_train_loaders(config)
# Create the optimizer
optimizer = create_optimizer(config['optimizer'], model)
# Create learning rate adjustment strategy
lr_scheduler = create_lr_scheduler(config.get('lr_scheduler', None),
optimizer)
trainer_config = config['trainer']
# Create tensorboard formatter
tensorboard_formatter = get_tensorboard_formatter(
trainer_config.pop('tensorboard_formatter', None))
# Create trainer
resume = trainer_config.pop('resume', None)
pre_trained = trainer_config.pop('pre_trained', None)
return UNet3DTrainer(model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
loss_criterion=loss_criterion,
eval_criterion=eval_criterion,
tensorboard_formatter=tensorboard_formatter,
device=config['device'],
loaders=loaders,
resume=resume,
pre_trained=pre_trained,
**trainer_config)
def _create_trainer(config, model, optimizer, lr_scheduler, loss_criterion, eval_criterion, loaders):
assert 'trainer' in config, 'Could not find trainer configuration'
trainer_config = config['trainer']
resume = trainer_config.get('resume', None)
pre_trained = trainer_config.get('pre_trained', None)
skip_train_validation = trainer_config.get('skip_train_validation', False)
# get tensorboard formatter
tensorboard_formatter = get_tensorboard_formatter(trainer_config.get('tensorboard_formatter', None))
if resume is not None:
# continue training from a given checkpoint
return UNet3DTrainer.from_checkpoint(resume, model,
optimizer, lr_scheduler, loss_criterion,
eval_criterion, loaders, tensorboard_formatter=tensorboard_formatter)
elif pre_trained is not None:
# fine-tune a given pre-trained model
return UNet3DTrainer.from_pretrained(pre_trained, model, optimizer, lr_scheduler, loss_criterion,
eval_criterion, device=config['device'], loaders=loaders,
max_num_epochs=trainer_config['epochs'],
max_num_iterations=trainer_config['iters'],
validate_after_iters=trainer_config['validate_after_iters'],
log_after_iters=trainer_config['log_after_iters'],
eval_score_higher_is_better=trainer_config['eval_score_higher_is_better'],
tensorboard_formatter=tensorboard_formatter,
skip_train_validation=skip_train_validation)
else:
# start training from scratch
return UNet3DTrainer(model, optimizer, lr_scheduler, loss_criterion, eval_criterion,
config['device'], loaders, trainer_config['checkpoint_dir'],
max_num_epochs=trainer_config['epochs'],
max_num_iterations=trainer_config['iters'],
validate_after_iters=trainer_config['validate_after_iters'],
log_after_iters=trainer_config['log_after_iters'],
eval_score_higher_is_better=trainer_config['eval_score_higher_is_better'],
tensorboard_formatter=tensorboard_formatter,
skip_train_validation=skip_train_validation)
