def define_criterion(self, name):
if name.lower() == 'bce+dice':
self.criterion = Loss.BCE_Dice()
elif name.lower() == 'dice':
self.criterion = Loss.DiceLoss()
elif name.lower() == 'bce':
self.criterion = nn.BCEWithLogitsLoss()
elif name.lower() == 'robustfocal':
self.criterion = Loss.RobustFocalLoss2d()
elif name.lower() == 'lovasz-hinge' or name.lower() == 'lovasz':
self.criterion = Loss.Lovasz_Hinge(per_image=True)
elif name.lower() == 'bce+lovasz':
self.criterion = Loss.BCE_Lovasz(per_image=True)
else:
raise NotImplementedError(
'Loss {} is not implemented'.format(name))
def define_criterion(self, name):
if name.lower() == 'bce+dice':
self.criterion = Loss.BCE_Dice()
elif name.lower() == 'dice':
self.criterion = Loss.DiceLoss()
elif name.lower() == 'bce':
# self.criterion = nn.BCEWithLogitsLoss()
# self.criterion = nn.CrossEntropyLoss(size_average=False, reduction='sum')
self.criterion = nn.CrossEntropyLoss()
# self.criterion = Loss.CE_SOFT()
elif name.lower() == 'robustfocal':
self.criterion = Loss.RobustFocalLoss2d()
elif name.lower() == 'lovasz-hinge' or name.lower() == 'lovasz':
self.criterion = Loss.Lovasz_Hinge(per_image=True)
elif name.lower() == 'bce+lovasz':
self.criterion = Loss.BCE_Lovasz(per_image=True)
else:
raise NotImplementedError(
'Loss {} is not implemented'.format(name))
