def __init__(self, plans_file, fold, local_rank, output_folder=None, dataset_directory=None, batch_dice=True, stage=None, unpack_data=True, deterministic=True, distribute_batch_size=False, fp16=False): super().__init__(plans_file, fold, output_folder, dataset_directory, batch_dice, stage, unpack_data, deterministic, fp16) self.init_args = (plans_file, fold, local_rank, output_folder, dataset_directory, batch_dice, stage, unpack_data, deterministic, distribute_batch_size, fp16) self.distribute_batch_size = distribute_batch_size np.random.seed(local_rank) torch.manual_seed(local_rank) if torch.cuda.is_available(): torch.cuda.manual_seed_all(local_rank) self.local_rank = local_rank if torch.cuda.is_available(): torch.cuda.set_device(local_rank) dist.init_process_group(backend='nccl', init_method='env://') self.loss = None self.ce_loss = RobustCrossEntropyLoss() self.global_batch_size = None # we need to know this to properly steer oversample
def __init__(self, soft_dice_kwargs, ce_kwargs, aggregate="sum", square_dice=False, weight_ce=1, weight_dice=1, log_dice=False, ignore_label=None): """ CAREFUL. Weights for CE and Dice do not need to sum to one. You can set whatever you want. :param soft_dice_kwargs: :param ce_kwargs: :param aggregate: :param square_dice: :param weight_ce: :param weight_dice: """ super(DC_and_CE_loss, self).__init__() if ignore_label is not None: assert not square_dice, 'not implemented' ce_kwargs['reduction'] = 'none' self.log_dice = log_dice self.weight_dice = weight_dice self.weight_ce = weight_ce self.aggregate = aggregate self.ce = RobustCrossEntropyLoss(**ce_kwargs) self.ignore_label = ignore_label if not square_dice: self.dc = SoftDiceLoss(apply_nonlin=softmax_helper, **soft_dice_kwargs) else: self.dc = SoftDiceLossSquared(apply_nonlin=softmax_helper, **soft_dice_kwargs)
def __init__(self, plans_file, fold, output_folder=None, dataset_directory=None, batch_dice=True, stage=None, unpack_data=True, deterministic=True, fp16=False): super().__init__(plans_file, fold, output_folder, dataset_directory, batch_dice, stage, unpack_data, deterministic, fp16) self.loss = RobustCrossEntropyLoss()
def __init__(self, soft_dice_kwargs, topk_kwargs, ce_kwargs, aggregate="sum", square_dice=False): super(DC_topk_ce_loss, self).__init__() self.aggregate = aggregate self.topk = TopKLoss(**topk_kwargs) self.ce = RobustCrossEntropyLoss(**ce_kwargs) if not square_dice: self.dc = SoftDiceLoss(apply_nonlin=softmax_helper, **soft_dice_kwargs) else: self.dc = SoftDiceLossSquared(apply_nonlin=softmax_helper, **soft_dice_kwargs)
def __init__(self, ce_kwargs, ignore_label=None): """ CAREFUL. Weights for CE and Dice do not need to sum to one. You can set whatever you want. :param soft_dice_kwargs: :param ce_kwargs: :param aggregate: :param square_dice: :param weight_ce: :param weight_dice: """ super(Only_CE_loss, self).__init__() if ignore_label is not None: ce_kwargs['reduction'] = 'none' self.ce = RobustCrossEntropyLoss(**ce_kwargs) self.ignore_label = ignore_label
def __init__(self, input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage=2, feat_map_mul_on_downscale=2, conv_op=nn.Conv2d, norm_op=nn.BatchNorm2d, norm_op_kwargs=None, dropout_op=nn.Dropout2d, dropout_op_kwargs=None, nonlin=nn.LeakyReLU, nonlin_kwargs=None, deep_supervision=True, dropout_in_localization=False, weightInitializer=InitWeights_He(1e-2), pool_op_kernel_sizes=None, conv_kernel_sizes=None, upscale_logits=False, convolutional_pooling=False, convolutional_upsampling=False, max_num_features=None): """ As opposed to the Generic_UNet, this class will compute parts of the loss function in the forward pass. This is useful for GPU parallelization. The batch DICE loss, if used, must be computed over the whole batch. Therefore, in a naive implementation, all softmax outputs must be copied to a single GPU which will then do the loss computation all by itself. In the context of 3D Segmentation, this results in a lot of overhead AND is inefficient because the DICE computation is also kinda expensive (Think 8 GPUs with a result of shape 2x4x128x128x128 each.). The DICE is a global metric, but its parts can be computed locally (TP, FP, FN). Thus, this implementation will compute all the parts of the loss function in the forward pass (and thus in a parallelized way). The results are very small (batch_size x num_classes for TP, FN and FP, respectively; scalar for CE) and copied easily. Also the final steps of the loss function (computing batch dice and average CE values) are easy and very quick on the one GPU they need to run on. BAM. final_nonlin is lambda x:x here! """ super(Generic_UNet_DP, self).__init__( input_channels, base_num_features, num_classes, num_pool, num_conv_per_stage, feat_map_mul_on_downscale, conv_op, norm_op, norm_op_kwargs, dropout_op, dropout_op_kwargs, nonlin, nonlin_kwargs, deep_supervision, dropout_in_localization, lambda x: x, weightInitializer, pool_op_kernel_sizes, conv_kernel_sizes, upscale_logits, convolutional_pooling, convolutional_upsampling, max_num_features) self.ce_loss = RobustCrossEntropyLoss()
def __init__(self, plans_file, fold, output_folder=None, dataset_directory=None, batch_dice=True, stage=None, unpack_data=True, deterministic=True, fp16=False): super(sauNetTrainer, self).__init__(plans_file, fold, output_folder, dataset_directory, batch_dice, stage, unpack_data, deterministic, fp16) self.loss = RobustCrossEntropyLoss() self.max_num_epochs = 350 self.initial_lr = 1e-3 self.num_classes = 1
def __init__(self, gdl_dice_kwargs, ce_kwargs, aggregate="sum"): super(GDL_and_CE_loss, self).__init__() self.aggregate = aggregate self.ce = RobustCrossEntropyLoss(**ce_kwargs) self.dc = GDL(softmax_helper, **gdl_dice_kwargs)