def _get_loss_criterion(loss):
if loss == 'bce':
return nn.BCELoss(), True
elif loss == 'ce':
return nn.CrossEntropyLoss(), False
elif loss == 'wce':
return WeightedCrossEntropyLoss(), False
else:
return GeneralizedDiceLoss(), True
def test_weighted_generalized_dice_loss(self):
shape = [1, 0, 30, 30, 30]
results = []
for C in range(1, 5):
batch_shape = list(shape)
batch_shape[1] = C
batch_shape = tuple(batch_shape)
results = results + _eval_criterion(GeneralizedDiceLoss(weight=torch.rand(C)), batch_shape)
# check that all of the coefficients belong to [0, 1]
results = np.array(results)
assert np.all(results > 0)
assert np.all(results < 1)
def _get_loss_criterion(loss_str, weight=None):
"""
Returns the loss function together with boolean flag which indicates
whether to apply an element-wise Sigmoid on the network output
"""
LOSSES = ['ce', 'bce', 'wce', 'dice']
#assert loss_str in LOSSES, f'Invalid loss string: {loss_str}'
print(loss_str)
if loss_str == 'bce':
return nn.BCELoss(weight), True
elif loss_str == 'ce':
return nn.CrossEntropyLoss(weight, ignore_index=-1), False
elif loss_str == 'wce':
return WeightedCrossEntropyLoss(ignore_index=-1), False
else:
return GeneralizedDiceLoss(), True
def test_generalized_dice_loss(self):
results = _compute_criterion(GeneralizedDiceLoss())
# check that all of the coefficients belong to [0, 1]
results = np.array(results)
assert np.all(results > 0)
assert np.all(results < 1)