training package

Submodules

training.utils module

training.utils.concat_all_gather(tensor)[source]

Performs all_gather operation on the provided tensor * Warning *: torch.distributed.all_gather has no gradient.

training.utils.exp_lr_scheduler_with_warmup(optimizer, init_lr, epoch, warmup_epoch, max_epoch)[source]
training.utils.filter_validation_results(dice_list, ASD_list, HD_list, args)[source]
training.utils.get_optimizer(args, net)[source]
training.utils.log_evaluation_result(writer, dice_list, ASD_list, HD_list, name, epoch, args)[source]
training.utils.multistep_lr_scheduler_with_warmup(optimizer, init_lr, epoch, warmup_epoch, lr_decay_epoch, max_epoch, gamma=0.1)[source]
training.utils.remove_wrap_arounds(tensor, ranks)[source]

Due to the DistributedSampler will pad samples for evenly distribute samples to gpus, the padded samples need to be removed for right evaluation. Need to turn shuffle to False for the dataloader.

training.utils.unwrap_model_checkpoint(net, ema_net, args)[source]
training.utils.update_ema_variables(model, ema_model, alpha, global_step)[source]

Module contents