I have a network which I’m trying to train a network for 2-class pixel-wise segmentation. To handle skew in the classes, I’m using the Dice loss. It works well with a baseline network that just predicts the probability of the pixel being 1. But in a second network, the outputs for each pixel are parameters of a Beta distribution, and samples are taken from it. The mean of these samples is used to calculate the dice loss. This works upto a certain point, from which it starts giving NaNs in the gradients. The loss value isn’t nan, however. The NaNs in gradients also do not occur when I train with BCE loss. I’m using a smooth dice loss with smoothness parameter = 1.
Here are some snippets:
import torch
import torch.nn as nn
from torch.nn import functional as F
from torch.autograd import Variable

class WNetBeta(nn.Module):
	def __init__(self, in_ch, out_ch, C0=32):
		super(WNetBeta, self).__init__()
		self.conv0 = conv(in_ch, C0, (1, 3, 3), padding=(0, 1, 1))

		self.resblock1 = ResBlock(C0, C0)
		self.resblock2 = ResBlock(C0, C0)
		self.tconv1 = tconv(C0, C0)
		
		self.avgsample = nn.AvgPool3d((1, 2, 2), stride=(1, 2, 2))

		self.resblock3 = ResBlock(C0, C0)
		self.resblock4 = ResBlock(C0, C0)

		self.tconv2 = tconv(C0, C0)

		self.pconv1 = pconv(C0, out_ch, (1, 3, 3), padding=(0, 1, 1))

	def forward(self, x):
		data = self.conv0(x)
		data = self.resblock1(data)
		data = self.resblock2(data)
		data = self.tconv1(data)
		data = self.avgsample(data)
		
		data = self.resblock3(data)
		data = self.resblock4(data)
		data = self.tconv2(data)
		out = self.pconv1(data)

		## get alpha, beta
		logalpha = out[:, 0]
		logbeta  = out[:, 1]

		alpha = torch.exp(logalpha)
		beta  = torch.exp(logbeta)

		# Get samples now
		T = 10
		alphar = alpha.unsqueeze(1).repeat(1, T, 1, 1, 1, 1)
		betar  =  beta.unsqueeze(1).repeat(1, T, 1, 1, 1, 1)
		m = Beta(alphar, betar)
		p = m.rsample().mean(1)

		return p, logalpha, logbeta

Think of the blocks ResBlock, pconv, etc. as conv blocks with different parameters. I have not pasted their snippets for brevity, and because they’re probably not relevant here.
The baseline is the same except that logalpha and logbeta are replaced by logits of the segmentation. The dice loss (rather 1- dice_loss) is given by the following code:
	def dice_loss(pred, target, smooth=1.):
		"""This definition generalize to real valued pred and target vector.
		This should be differentiable.
		pred: tensor with first dimension as batch
		target: tensor with first dimension as batch
		"""
		# have to use contiguous since they may from a torch.view op
		# iflat = B * N
		# tflat = B * N
		iflat = torch.flatten(pred.contiguous(), 1)
		tflat = torch.flatten(target.contiguous(), 1)
		intersection = (iflat * tflat).sum(1)
		A_sum = iflat.sum(1)
		B_sum = tflat.sum(1)
		
		loss = 1 - ((2. * intersection + smooth) / (A_sum + B_sum + smooth)).mean()
		return loss

In the case of network with beta parameterization, the dice loss is taken with the mean of the samples.
Can anyone tell me what is going wrong in here?  Thank you.