Custom Loss Function Gradient Vector nan

username: snoopkat
cooked: <p>I am writing a simplified version of the YOLO v1 object detection model for face detection. The input for this model is an n x 3 x 300 x 300 tensor of RGB images. The output is a n x 5 x 7 x 7 tensor, where the 5 channels represent x, y, w, h, and prob(center of a face) for 49 different regions. I want the yolo model to produce tensors of the same shape, with the same types of channels.</p>
<p>The simplified yolo module is defined as the following:</p>
<pre><code class="lang-auto">class YOLO(nn.Module):
  def __init__(self):
    super(YOLO, self).__init__()
    self.mp = nn.MaxPool2d(kernel_size=2, stride=2)
    self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
    self.conv2 = nn.Conv2d(64, 192, kernel_size=3, padding=1)
    self.conv3 = nn.Conv2d(192, 256, kernel_size=3, padding=1)
    self.conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1)
    self.conv5 = nn.Conv2d(512, 1024, kernel_size=3, padding=1)
    self.conv6 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
    self.conv7 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
    self.conv8 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
    self.conv9 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
    self.conv10 = nn.Conv2d(1024, 1024, kernel_size=3)

    self.relu_downsize = nn.Sequential(
        nn.MaxPool2d(kernel_size=2, stride=2),
        nn.LeakyReLU(.01, inplace=False)
    )
    
    self.linear_layers = nn.Sequential(
        nn.Flatten(),
        nn.Linear(7 * 7 * 1024,  4096),
        nn.LeakyReLU(.01, inplace=False),
        nn.Linear(4096, 7 * 7 * 5),
        nn.LeakyReLU(.01, inplace=False)
    )

  def forward(self, x):

    x = self.relu_downsize(self.conv1(x))
    x = self.relu_downsize(self.conv2(x))
    x = self.relu_downsize(self.conv3(x))
    x = F.leaky_relu(self.conv4(x), .01)
    x = self.relu_downsize(self.conv5(x))
    x = F.leaky_relu(self.conv6(x), .01)
    x = self.relu_downsize(self.conv7(x))
    x = F.leaky_relu(self.conv8(x))
    x = F.leaky_relu(self.conv9(x))
    x = F.leaky_relu(self.conv10(x))
    x = torch.sigmoid(self.linear_layers(x))
    x = x.view(-1, 5, 7, 7)
    return x
</code></pre>
<p>And the simplified yolo loss as the following:</p>
<pre><code class="lang-auto">class YoloLoss(nn.Module):
  def __init__(self):
    super(YoloLoss, self).__init__()
    self.lmbda_coord = torch.autograd.Variable(torch.tensor(.5), requires_grad=True)
    self.lmbda_noobj = torch.autograd.Variable(torch.tensor(.5), requires_grad=True)

  def forward(self, yhat, y):

    cp_dim = (1, 5, 1, 1)

    region_yhat = yhat * torch.tile(y[:, 4, :, :].unsqueeze(1), cp_dim)
    nonregion_yhat = (1 - yhat) * torch.tile((1 - y[:, 4, :, :].unsqueeze(1)), cp_dim)

    err_loc = F.mse_loss(region_yhat[:, :2, :, :], y[:, :2, :, :])
    err_size = F.mse_loss(torch.sqrt(region_yhat[:, 2:4, :, :]), torch.sqrt(y[:, 2:4, :, :]))
    err_inclass = F.mse_loss(region_yhat[:, 4, :, :], y[:, 4, :, :])
    err_outclass = F.mse_loss(nonregion_yhat[:, 4, :, :], 1 - y[:, 4, :, :])

    return self.lmbda_coord * (err_loc + err_size) + err_inclass + self.lmbda_noobj + err_outclass
</code></pre>
<p>The problem is that the gradient vector seems not to be calculated correctly. During training, the first loss is correctly calculated, but after the first gradient descent step, the model only returns nans. The following code block produces the error. Keep in mind that the <code>get_images_and_boxes</code> function returns a tuple of tensor images (input to a model) and tensor bounding boxes (actual output to predict).</p>
<p>Code:</p>
<pre><code class="lang-auto">mod = YOLO()
loss_f = YoloLoss()
optimizer = torch.optim.Adam(params=mod.parameters(), lr=1e-3)

x1, y1 = get_images_and_boxes(range(4))
yhat1 = mod(x1)
print(x1.shape, y1.shape, yhat1.shape)

loss = loss_f(yhat1, y1)
print(loss)

loss.backward()
optimizer.step()

x2, y2 = get_images_and_boxes(range(4))
yhat2 = mod(x2)
print(x2.shape, y2.shape, yhat2.shape)
print(yhat2)
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">torch.Size([4, 3, 300, 300]) torch.Size([4, 5, 7, 7]) torch.Size([4, 5, 7, 7])
tensor(0.7533, grad_fn=&lt;AddBackward0&gt;)
torch.Size([4, 3, 300, 300]) torch.Size([4, 5, 7, 7]) torch.Size([4, 5, 7, 7])
tensor([[[[nan, nan, nan, nan, nan, nan, nan],
          [nan, nan, nan, nan, nan, nan, nan],
          [nan, nan, nan, nan, nan, nan, nan],
          [nan, nan, nan, nan, nan, nan, nan],
          [nan, nan, nan, nan, nan, nan, nan],
          [nan, nan, nan, nan, nan, nan, nan],
          [nan, nan, nan, nan, nan, nan, nan]],
          ...
</code></pre>
<p>I am not sure why backpropagation is not working on my custom loss function.</p>

username: lskook47
cooked: <p>Hi there<br>
Check whether your argument for torch.sqrt() has no negative element or not.<br>
That was the problem, for my case</p>