def main():
# Create main logger
logger = get_logger('UNet3DTrainer')
# Load and log experiment configuration
config = load_config()
logger.info(config)
manual_seed = config.get('manual_seed', None)
if manual_seed is not None:
logger.info(f'Seed the RNG for all devices with {manual_seed}')
torch.manual_seed(manual_seed)
# see https://pytorch.org/docs/stable/notes/randomness.html
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Create the model
model = get_model(config)
# put the model on GPUs
logger.info(f"Sending the model to '{config['device']}'")
model = model.to(config['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)
# Cross validation
path_to_folder = config['loaders']['all_data_path'][0]
cross_walidation = CrossValidation(path_to_folder, 1, 3, 2)
train_set = cross_walidation.train_filepaths
val_set = cross_walidation.validation_filepaths
config['loaders']['train_path'] = train_set
config['loaders']['val_path'] = val_set
# Create data loaders
loaders = get_train_loaders(config)
# Create the optimizer
optimizer = _create_optimizer(config, model)
# Create learning rate adjustment strategy
lr_scheduler = _create_lr_scheduler(config, optimizer)
# Create model trainer
trainer = _create_trainer(config,
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
loss_criterion=loss_criterion,
eval_criterion=eval_criterion,
loaders=loaders,
logger=logger)
# Start training
trainer.fit()
def _train_save_load(self, tmpdir, loss, val_metric, model='UNet3D', max_num_epochs=1, log_after_iters=2,
validate_after_iters=2, max_num_iterations=4, weight_map=False):
binary_loss = loss in ['BCEWithLogitsLoss', 'DiceLoss', 'GeneralizedDiceLoss']
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
test_config = copy.deepcopy(CONFIG_BASE)
test_config['model']['name'] = model
test_config.update({
# get device to train on
'device': device,
'loss': {'name': loss, 'weight': np.random.rand(2).astype(np.float32)},
'eval_metric': {'name': val_metric}
})
test_config['model']['final_sigmoid'] = binary_loss
if weight_map:
test_config['loaders']['weight_internal_path'] = 'weight_map'
loss_criterion = get_loss_criterion(test_config)
eval_criterion = get_evaluation_metric(test_config)
model = get_model(test_config)
model = model.to(device)
if loss in ['BCEWithLogitsLoss']:
label_dtype = 'float32'
else:
label_dtype = 'long'
test_config['loaders']['transformer']['train']['label'][0]['dtype'] = label_dtype
test_config['loaders']['transformer']['test']['label'][0]['dtype'] = label_dtype
train, val = TestUNet3DTrainer._create_random_dataset((128, 128, 128), (64, 64, 64), binary_loss)
test_config['loaders']['train_path'] = [train]
test_config['loaders']['val_path'] = [val]
loaders = get_train_loaders(test_config)
optimizer = _create_optimizer(test_config, model)
test_config['lr_scheduler']['name'] = 'MultiStepLR'
lr_scheduler = _create_lr_scheduler(test_config, optimizer)
logger = get_logger('UNet3DTrainer', logging.DEBUG)
formatter = DefaultTensorboardFormatter()
trainer = UNet3DTrainer(model, optimizer, lr_scheduler,
loss_criterion, eval_criterion,
device, loaders, tmpdir,
max_num_epochs=max_num_epochs,
log_after_iters=log_after_iters,
validate_after_iters=validate_after_iters,
max_num_iterations=max_num_iterations,
logger=logger, tensorboard_formatter=formatter)
trainer.fit()
# test loading the trainer from the checkpoint
trainer = UNet3DTrainer.from_checkpoint(os.path.join(tmpdir, 'last_checkpoint.pytorch'),
model, optimizer, lr_scheduler,
loss_criterion, eval_criterion,
loaders, logger=logger, tensorboard_formatter=formatter)
return trainer
def main():
logger = get_logger('UNet3DTrainer')
config = load_config()
logger.info(config)
# Create loss criterion
loss_criterion = get_loss_criterion(config)
# Create the model
model = UNet3D(config['in_channels'], config['out_channels'],
final_sigmoid=config['final_sigmoid'],
init_channel_number=config['init_channel_number'],
conv_layer_order=config['layer_order'],
interpolate=config['interpolate'])
model = model.to(config['device'])
# Log the number of learnable parameters
logger.info(f'Number of learnable params {get_number_of_learnable_parameters(model)}')
# Create evaluation metric
eval_criterion = get_evaluation_metric(config)
loaders = get_train_loaders(config)
# Create the optimizer
optimizer = _create_optimizer(config, model)
# Create learning rate adjustment strategy
lr_scheduler = _create_lr_scheduler(config, optimizer)
if config['resume'] is not None:
trainer = UNet3DTrainer.from_checkpoint(config['resume'], model,
optimizer, lr_scheduler, loss_criterion,
eval_criterion, loaders,
logger=logger)
else:
trainer = UNet3DTrainer(model, optimizer, lr_scheduler, loss_criterion, eval_criterion,
config['device'], loaders, config['checkpoint_dir'],
max_num_epochs=config['epochs'],
max_num_iterations=config['iters'],
validate_after_iters=config['validate_after_iters'],
log_after_iters=config['log_after_iters'],
logger=logger)
trainer.fit()
def main():
# Create main logger
logger = get_logger('UNet3DTrainer')
# Load and log experiment configuration
config = load_config() # Set DEFAULT_DEVICE and config file
logger.info(config) # Log configure from train_config_4d_input.yaml
manual_seed = config.get('manual_seed', None)
if manual_seed is not None:
logger.info(f'Seed the RNG for all devices with {manual_seed}')
torch.manual_seed(manual_seed)
torch.backends.cudnn.deterministic = True # Ensure the repeatability of the experiment
torch.backends.cudnn.benchmark = False # Benchmark mode improves the computation speed, but results in slightly different network feedforward results
# Create the model
model = get_model(config)
# put the model on GPUs
logger.info(f"Sending the model to '{config['device']}'")
model = model.to(config['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: {'train': train_loader, 'val': val_loader}
loaders = get_train_loaders(config)
# Create the optimizer
optimizer = _create_optimizer(config, model)
# Create learning rate adjustment strategy
lr_scheduler = _create_lr_scheduler(config, optimizer)
# Create model trainer
trainer = _create_trainer(config, model=model, optimizer=optimizer, lr_scheduler=lr_scheduler,
loss_criterion=loss_criterion, eval_criterion=eval_criterion, loaders=loaders,
logger=logger)
# Start training
trainer.fit()
def main():
# Load and log experiment configuration
config = load_config()
# Create main logger
logfile = config.get('logfile', None)
config['logfile'] = logfile
logger = get_logger('UNet3DTrainer', logfile=logfile)
logger.info(config)
manual_seed = config.get('manual_seed', None)
if manual_seed is not None:
logger.info(f'Seed the RNG for all devices with {manual_seed}')
torch.manual_seed(manual_seed)
# see https://pytorch.org/docs/stable/notes/randomness.html
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Create the model
model = get_model(config)
# put the model on GPUs
logger.info(f"Sending the model to '{config['device']}'")
# multiple GPUs
if (torch.cuda.device_count() > 1):
logger.info("There are {} GPUs available".format(
torch.cuda.device_count()))
model = nn.DataParallel(model)
model = model.to(config['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, model)
# Create learning rate adjustment strategy
lr_scheduler = _create_lr_scheduler(config, optimizer)
# Create model trainer
trainer = _create_trainer(config,
model=model,
optimizer=optimizer,
lr_scheduler=lr_scheduler,
loss_criterion=loss_criterion,
eval_criterion=eval_criterion,
loaders=loaders,
logger=logger)
# Start training
trainer.fit()
def _train_save_load(self,
tmpdir,
loss,
val_metric,
max_num_epochs=1,
log_after_iters=2,
validate_after_iters=2,
max_num_iterations=4):
# conv-relu-groupnorm
conv_layer_order = 'crg'
final_sigmoid = loss in ['bce', 'dice']
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
test_config = dict(CONFIG_BASE)
test_config.update({
# get device to train on
'device': device,
'loss': {
'name': loss,
'weight': np.random.rand(2).astype(np.float32)
},
'eval_metric': {
'name': val_metric
}
})
loss_criterion = get_loss_criterion(test_config)
eval_criterion = get_evaluation_metric(test_config)
model = self._create_model(final_sigmoid, conv_layer_order)
channel_per_class = loss in ['bce', 'dice', 'gdl']
if loss in ['bce']:
label_dtype = 'float32'
else:
label_dtype = 'long'
test_config['loaders']['transformer']['train']['label'][0][
'dtype'] = label_dtype
test_config['loaders']['transformer']['test']['label'][0][
'dtype'] = label_dtype
train, val = TestUNet3DTrainer._create_random_dataset(
(128, 128, 128), (64, 64, 64), channel_per_class)
test_config['loaders']['train_path'] = [train]
test_config['loaders']['val_path'] = [val]
loaders = get_train_loaders(test_config)
learning_rate = 2e-4
weight_decay = 0.0001
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
lr_scheduler = MultiStepLR(optimizer, milestones=[2, 3], gamma=0.5)
logger = get_logger('UNet3DTrainer', logging.DEBUG)
trainer = UNet3DTrainer(model,
optimizer,
lr_scheduler,
loss_criterion,
eval_criterion,
device,
loaders,
tmpdir,
max_num_epochs=max_num_epochs,
log_after_iters=log_after_iters,
validate_after_iters=validate_after_iters,
max_num_iterations=max_num_iterations,
logger=logger)
trainer.fit()
# test loading the trainer from the checkpoint
trainer = UNet3DTrainer.from_checkpoint(os.path.join(
tmpdir, 'last_checkpoint.pytorch'),
model,
optimizer,
lr_scheduler,
loss_criterion,
eval_criterion,
loaders,
logger=logger)
return trainer
model = get_model(config)
# put the model on GPUs
logger.info(f"Sending the model to '{config['device']}'")
model = model.to(config['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)
config['n_iter_loader'] = len(loaders['train'])
# Create the optimizer
optimizer = _create_optimizer(config, model)
# Create learning rate adjustment strategy
lr_scheduler = _create_lr_scheduler(config, optimizer)
if config['findlr']:
from unet3d.utils import find_lr
log_lrs, losses = find_lr(model,
loaders['train'],
loss_criterion,
optimizer,