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_ce(self):
criterion = WeightedCrossEntropyLoss()
shape = [1, 0, 30, 30, 30]
target_shape = [1, 30, 30, 30]
results = []
for C in range(1, 5):
input_shape = list(shape)
input_shape[1] = C
input_shape = tuple(input_shape)
for i in range(100):
input = torch.rand(input_shape)
target = torch.zeros(target_shape, dtype=torch.long).random_(0, C)
results.append(criterion(input, target))
results = np.array(results)
assert np.all(results >= 0)
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