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

1

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
def backward(self, grad_output):

    input, target = self.saved_variables
    grad_input = grad_target = None

    if self.needs_input_grad[0]:
        grad_input = grad_output * 2 * (target * self.union + self.inter) \
                     / self.union * self.union
    if self.needs_input_grad[1]:
        grad_target = None

    return grad_input, grad_target

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.

2

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
3

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