# The speed of train become so slow when i use the loss_function i defined

class DiceCoeff(Function):
“”“Dice coeff for individual examples”""

``````def forward(self, input, target):
self.save_for_backward(input, target)
self.inter = torch.dot(input.view(-1), target.view(-1)) + 0.0001
self.union = torch.sum(input) + torch.sum(target) + 0.0001

t = 2 * self.inter.float() / self.union.float()
return t

# This function has only a single output, so it gets only one gradient

input, target = self.saved_variables

grad_input = grad_output * 2 * (target * self.union + self.inter) \
/ self.union * self.union

``````

def dice_coeff(input, target):
“”“Dice coeff for batches”""
if input.is_cuda:
s = torch.FloatTensor(1).cuda().zero_()
else:
s = torch.FloatTensor(1).zero_()

``````for i, c in enumerate(zip(input, target)):
s = s + DiceCoeff().forward(c[0], c[1])

return s / (i + 1)
``````

it is so so so so so slower then i use BCEloss.it almost can not to be used.

why do you define it as autograd function? The way you used it, you always create a new instance of your Function. Why not defining it like this:

``````class DiceCoeff(torch.nn.Module):
def __init__(self, eps=0.0001)
super().__init__()
self._eps = eps

def forward(self, prediction, target):
inter = torch.dot(prediction.view(-1), target.view(-1)) + self._eps
union = prediction.sum() + target.sum() + self._eps

return 2*inter.float()/union.float()
``````

and using it like this:

``````dice_crit = DiceCoeff()
total_dice_coeff = 0
single_dices = [dice_crit(_x, _y) for _x, _y in zip(prediction, target)]

for _dice in single_dices:
total_dice_coeff += _dice

total_dice_coeff = total_dice_coeff / len(single_dices)
``````

This way you only create the criterion once, autograd handles the gradient computing (in a very efficient way) and you don’t have to create the s tensor yourself every time. From my expierience custom loss functions are almost as fast as the native ones, if they are implemented in `torch.nn.Module`

Can you test if this approach speeds up your training?

EDIT: Assuming prediction and target are minibatch-tensors you could do something like this:

``````class DiceCoeff(torch.nn.Module):
def __init__(self, eps=0.0001)
super().__init__()
self._eps = eps

def forward(self, prediction, target):
flattened_pred = prediction.view(prediction.size(0), -1)
flattened_target = target.view(target.size(0), -1)
inter = torch.dot(flattened_pred, flattened_target) + self._eps
union = flattened_pred.sum(-1) + flattened_target.sum(-1) + self._eps

return (2*inter.float()/union.float()).mean()
``````

and simply pass the whole batches like this:

``````dice_crit = DiceCoeff()
total_dice_coeff = dice_crit(predictions, targets)
``````

and you will directly get the mean dice for the current batch

2 Likes

Thank you very very much！It run faster than before.Your guidance solve my problem perfectly！