Hi!
I am working on a segmentation problem and wanted to try custom loss functions (namely dice and a variant of BCE). The functions seem to work fine as of the forward pass, but when doing backward they converge very fast to a local minimum which consists in zeroing the output mask. When I check the gradients, they tend to be very high to begin with, and then get stuck to 0 or nan when the local minima is reached. Gradient clipping doesn’t help (it is just solwer to get to the minimum). It may be specific to my problem but I believe there is something I don’t get in pytorch’s implementation, as a simple BCE doesn’t provoke this issue, while my custom version does. Below is the code for my functions (they are of course then called through a nn.Module):
def bce_loss(input, target, reduction='mean', beta=None, eps=1e-12):
    n = input.size(0)
    iflat = torch.sigmoid(input).view(n, -1).clamp(eps, 1-eps)
    tflat = target.view(n, -1)
    if not beta:
        beta = 0.5
    bce = -2*(beta*tflat*iflat.log()+(1-beta)*(1-tflat)*(1-iflat).log()).mean(-1)
    if reduction == 'mean':
        return bce.mean()
    elif reduction == 'sum':
        return bce.sum()
    else:
        return bce

def dice(input, target, smooth=1., reduction='mean'):
    assert input.shape==target.shape, "input and target must have same shape"
    iflat = torch.sigmoid(input).view(input.size(0), -1)
    tflat = target.view(target.size(0), -1)
    intersection = (iflat * tflat).sum(-1)
    dice = 1 - (2. * intersection + smooth)/(iflat.sum(-1) + tflat.sum(-1) + smooth)
    if reduction=='mean':
        return dice.mean()
    elif reduction=='sum':
        return dice.sum()
    else:
        return dice

I use mean reduction for everything. Even when I try using bce with beta=0.5 (which is standard bce in theory), it doesn’t work. Is there any magic in autograd that I don’t get ? Do I need to define my custom backward to avoid this ?
Thanks !
EDIT: I also used gradient accumulation which may not help, but that doesn’t change the fact that pytorch’s BCE works perfectly fine, while mine doesn’t.