Gradient for different parts of the loss is None even if `.retain_grad()` is called

szobov: <p>Hi folks,</p>
<p>I’m in training auto-encoder on 3D objects and I have such a piece of code, to calculate the loss for the forward pass.</p>
<pre><code class="lang-auto">    def forward(self, in_pc_batch, iteration, frame, t_idx, t_nor, faces, bDebug = False):

        nbat = in_pc_batch.size(0)
        npt = in_pc_batch.size(1)
        nch = in_pc_batch.size(2)


        t_mu, t_logstd = self.net_geoenc(in_pc_batch, self.mcvcoeffsenc) # in in mm and out in dm
        t_std = t_logstd.exp()

        t_eps = torch.ones_like(t_std).normal_() #torch.FloatTensor(t_std.size()).normal_().to(device)
        t_z = t_mu + t_std * t_eps


        klloss = torch.mean(-0.5 - t_logstd + 0.5 * t_mu ** 2 + 0.5 * torch.exp(2 * t_logstd))
        out_pc_batchfull = self.net_geodec(t_z, self.mcvcoeffsdec)
        out_pc_batch = out_pc_batchfull[:,:,0:3]


        dif_pos = out_pc_batch - in_pc_batch

        faces_long = faces.long()
        vet0 = index_selection_nd(dif_pos, faces[:, :, 0], 1)
        vet1 = index_selection_nd(dif_pos, faces[:, :, 1], 1)
        vet2 = index_selection_nd(dif_pos, faces[:, :, 2], 1)

        loss_normal = ((vet1 + vet0 + vet2)/3.0 * t_nor).sum(2).pow(2).mean()

        loss_pose_l1 = self.net_loss.compute_geometric_loss_l1(in_pc_batch[:,:,0:3], out_pc_batch[:,:,0:3])
        loss_laplace_l1 = self.net_loss.compute_laplace_loss_l2(in_pc_batch[:,:,0:3], out_pc_batch[:,:,0:3])

        loss = loss_pose_l1*self.w_pose +  loss_laplace_l1 * self.w_laplace  + klloss * self.klweight + loss_normal * self.w_nor
</code></pre>
<p>But, when I try to analyze the gradients of different losses, it gives me <code>None</code>, even if I explicitly call <code>.retain_grad()</code> on it:</p>
<pre><code class="lang-auto">        loss,loss_pose_l1,loss_laplace_l1, klloss, loss_normal  = net_autoenc(pc_batch, iteration,frame, t_idx, t_nor, faces, True)

        optimizer.zero_grad()
        for l in [loss_laplace_l1, loss_pose_l1, klloss, loss_normal]:
            l.retain_grad()

        loss.register_hook(lambda grad: print("Loss_grad: ", grad))
        loss_pose_l1.register_hook(lambda grad: print("Loss_pos_l1_grad: ", grad))

        loss.backward(torch.ones(loss.size(0), device=device))

        print("Loss_pose_l1_grad: ", loss_pose_l1.grad)
        print("Loss_laplace_l1_grad: ", loss_laplace_l1.grad)
        print("klloss_grad: ", klloss.grad)
        print("Loss_normal_grad: ", loss_normal.grad)
        optimizer.step()
</code></pre>
<p>The output of this is:</p>
<pre><code class="lang-auto">Loss_grad:  tensor([1.], device='cuda:0')
Loss_pose_l1_grad:  None
Loss_laplace_l1_grad:  None
klloss_grad:  None
Loss_normal_grad:  None
</code></pre>
<p>When I tried to interactively check the tensors:</p>
<pre><code class="lang-auto">ipdb&gt; loss_pose_l1
tensor([0.0060], device='cuda:0', grad_fn=&lt;UnsqueezeBackward0&gt;)
ipdb&gt; loss_pose_l1.grad
ipdb&gt; loss_pose_l1.requires_grad
True
ipdb&gt; loss_pose_l1.is_leaf
False
ipdb&gt; loss_pose_l1.retains_grad
True
</code></pre>
<p>Could someone please help me, to understand, why the gradients are <code>None</code>?</p>
<p>My PyTorch version is <code>'1.9.0+cu111'</code>.</p>
<p>Thanks in advance!</p>

ptrblck: <p>Assuming you are returning the summed <code>loss</code> as well as all separate losses in your <code>forward</code> method, it seems to work for me:</p>
<pre><code class="lang-python">def fun(model, data):
    out = model(data)
    loss_pose_l1 = out.mean()
    loss_laplace_l1 = out.sum()
    klloss = F.mse_loss(out, torch.rand_like(out))
    loss_normal = (out**2).mean()    
    
    loss = loss_pose_l1 + loss_laplace_l1 + klloss + loss_normal
    return loss, loss_pose_l1, loss_laplace_l1, klloss, loss_normal


model = nn.Linear(1, 1)
data = torch.randn(1, 1)
loss, loss_pose_l1, loss_laplace_l1, klloss, loss_normal = fun(model, data)

for l in [loss_laplace_l1, loss_pose_l1, klloss, loss_normal]:
    l.retain_grad()
    
loss.backward()
print("Loss_pose_l1_grad: ", loss_pose_l1.grad)
print("Loss_laplace_l1_grad: ", loss_laplace_l1.grad)
print("klloss_grad: ", klloss.grad)
print("Loss_normal_grad: ", loss_normal.grad)
</code></pre>
<p>so could you post an executable code snippet to reproduce the issue?</p>

szobov: <p>Hi ptrblck,<br>
thanks a lot for the reply!</p>
<p>Yes, I already tested on such a simple block and it works, so I was super confused about what’s wrong in my case.</p>
<p>The only thing, that finally worked for me was to put <code>.register_hook</code> inside the <code>forward</code> method:</p>
<pre><code class="lang-auto">    def forward(self, in_pc_batch, iteration, frame, t_idx, t_nor, faces, bDebug = False):

        nbat = in_pc_batch.size(0)
        npt = in_pc_batch.size(1)
        nch = in_pc_batch.size(2)


        t_mu, t_logstd = self.net_geoenc(in_pc_batch, self.mcvcoeffsenc) # in in mm and out in dm
        t_std = t_logstd.exp()

        t_eps = torch.ones_like(t_std).normal_() 
        t_z = t_mu + t_std * t_eps

        t_z.register_hook(lambda grad: print("t_z_grad: ", grad, ", shape: ", grad.shape))

        klloss = torch.mean(-0.5 - t_logstd + 0.5 * t_mu ** 2 + 0.5 * torch.exp(2 * t_logstd))
        out_pc_batchfull = self.net_geodec(t_z, self.mcvcoeffsdec)
        out_pc_batch = out_pc_batchfull[:,:,0:3]
        dif_pos = out_pc_batch - in_pc_batch

        faces_long = faces.long()
        vet0 = index_selection_nd(dif_pos, faces[:, :, 0], 1)
        vet1 = index_selection_nd(dif_pos, faces[:, :, 1], 1)
        vet2 = index_selection_nd(dif_pos, faces[:, :, 2], 1)
        loss_normal = ((vet1 + vet0 + vet2)/3.0 * t_nor).sum(2).pow(2).mean()

        loss_pose_l1 = self.net_loss.compute_geometric_loss_l1(in_pc_batch[:,:,0:3], out_pc_batch[:,:,0:3])
        loss_laplace_l1 = self.net_loss.compute_laplace_loss_l2(in_pc_batch[:,:,0:3], out_pc_batch[:,:,0:3])

        loss_pose_l1.register_hook(lambda grad: print("Loss_pos_l1_grad: ", grad))
        loss_laplace_l1.register_hook(lambda grad: print("Loss_laplace_l1_grad: ", grad))
        klloss.register_hook(lambda grad: print("KLloss_grad: ", grad))
        loss_normal.register_hook(lambda grad: print("Loss_normal_grad: ", grad))

        loss = loss_pose_l1*self.w_pose + loss_laplace_l1 * self.w_laplace + klloss * self.klweight + loss_normal * self.w_nor

        return loss[None],loss_pose_l1[None],loss_laplace_l1[None], klloss[None], loss_normal[None]
</code></pre>
<p>In this case, it prints right values (weights).<br>
In the opposite case, even if I call <code>retain_grad</code> or call <code>register hook</code> outside <code>forward</code> function it doesn’t call a hook our output <code>None</code>.</p>
<p>And to clarify,</p>
<pre><code class="lang-auto">        loss,loss_pose_l1,loss_laplace_l1, klloss, loss_normal  = net_autoenc(pc_batch, iteration,frame, t_idx, t_nor, faces, True)
</code></pre>
<p><code>net_autoenc</code> is created in such way:</p>
<pre><code class="lang-auto">net_autoenc = Net_autoenc(param, facedata)
net_autoenc = nn.DataParallel(net_autoenc, device_ids=device_ids).to(device)
</code></pre>
<p>If you need to have any extra information, please ask!<br>
Cheers</p>