Training with gradient checkpoints (torch.utils.checkpoint) appears to reduce performance of model

JWL: <p>I have a snippet of code that uses gradient checkpoints from <code>torch.utils.checkpoint</code> to reduce GPU memory:</p>
<pre><code>    if use_checkpointing:
        res2, res3, res4, res5 = checkpoint.checkpoint(self.resnet_backbone, data['data'])
        fpn_p2, fpn_p3, fpn_p4, fpn_p5, fpn_p6 = checkpoint.checkpoint(self.fpn, res2, res3, res4, res5)
    else:
        res2, res3, res4, res5 = self.resnet_backbone(data['data'])
        fpn_p2, fpn_p3, fpn_p4, fpn_p5, fpn_p6 = self.fpn(res2, res3, res4, res5)
</code></pre>
<p>When I use gradient checkpointing, after training my model, the performance is worse. Does anyone know why this might be?</p>

albanD: <p>Hi,</p>
<p>This most likely happens because the first part of your model doesn’t get gradient because of some quirks of how checkpointing works.<br>
Can you try making <code>data['data']</code> require gradients before giving it to the checkpoint? (you can ignore the computed gradient, just add a <code>data['data'].requires_grad_()</code> if it is a Tensor.</p>

JWL: <p>When I do that I appear to get a new error:</p>
<pre><code class="lang-auto">[4]&lt;stderr&gt;:    loss.backward()
[4]&lt;stderr&gt;:  File "/pythonhome_pypi/lib/python3.6/site-packages/torch/tensor.py", line 195, in backward
[4]&lt;stderr&gt;:    torch.autograd.backward(self, gradient, retain_graph, create_graph)
[4]&lt;stderr&gt;:  File "/pythonhome_pypi/lib/python3.6/site-packages/torch/autograd/__init__.py", line 99, in backward
[4]&lt;stderr&gt;:    allow_unreachable=True)  # allow_unreachable flag
[4]&lt;stderr&gt;:  File "/pythonhome_pypi/lib/python3.6/site-packages/torch/autograd/function.py", line 77, in apply
[4]&lt;stderr&gt;:    return self._forward_cls.backward(self, *args)
[4]&lt;stderr&gt;:  File "/pythonhome_pypi/lib/python3.6/site-packages/torch/utils/checkpoint.py", line 99, in backward
[4]&lt;stderr&gt;:    torch.autograd.backward(outputs, args)
[4]&lt;stderr&gt;:  File "/pythonhome_pypi/lib/python3.6/site-packages/torch/autograd/__init__.py", line 99, in backward
[4]&lt;stderr&gt;:    allow_unreachable=True)  # allow_unreachable flag
[4]&lt;stderr&gt;:RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>
<p>To be specific I added <code>.requires_grad_()</code>:</p>
<pre><code class="lang-auto">res2, res3, res4, res5 = checkpoint.checkpoint(self.resnet_backbone, data['data'].requires_grad_())
</code></pre>

albanD: <p>Your <code>loss</code> does not requires gradients anymore because you added an extra <code>requires_grad_()</code>? That looks weird. Did you change anything else by any chance?</p>

JWL: <p>Yeah that is all I added, I just double checked and this is the case D:</p>

albanD: <p>Can you double check if the outputs of the checkpoint require gradient or not? In both cases.</p>

JWL: <p>I tested all cases:</p>
<pre><code class="lang-auto">res2, res3, res4, res5 = checkpoint.checkpoint(self.resnet_backbone, data['data'])
print(res2.requires_grad, res3.requires_grad, res4.requires_grad, res5.requires_grad)
res2, res3, res4, res5 = checkpoint.checkpoint(self.resnet_backbone, data['data'].requires_grad_())
print(res2.requires_grad, res3.requires_grad, res4.requires_grad, res5.requires_grad)
res2, res3, res4, res5 = self.resnet_backbone(data['data'])
print(res2.requires_grad, res3.requires_grad, res4.requires_grad, res5.requires_grad)
</code></pre>
<p>This outputs:</p>
<pre><code class="lang-auto">False False False False
True True True True
False True True True
</code></pre>
<p>In the original resnet backbone architecture of my code I have frozen res2 layer, could this be the issue?</p>
<p>Edit: it appears this is the culprit as I can train when I unfreeze the res2 layer. I suppose to fix this I can just <code>.detach()</code> it, ie.</p>
<pre><code class="lang-auto">res2, res3, res4, res5 = checkpoint.checkpoint(self.resnet_backbone, data['data'].requires_grad_())
res2 = res2.detach()
</code></pre>

albanD: <p>Ok,</p>
<p>So adding the requires grad on the data does fix the issue that the backbone would not be learning.</p>
<p>For the res2 freezing I’m not sure what you mean here.<br>
The checkpoint has this behavior that it make all outputs require gradient, because it does not know which elements will actually require it yet.<br>
Note that in the final computation during the backward, that gradient (should) will be discarded and not used, so the frozen part should remain frozen. Even though you don’t see it in the forward pass.</p>
<p>Adding the detach() after the checkpoint will work as well <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

JWL: <p>Yeah so in my model, I have code that detaches the gradients for res2, ie:</p>
<pre><code class="lang-auto">        with torch.no_grad():
            out = self.mod1(img)
            out = self.mod2(self.pool2(out))
            out = self.mod3(self.pool3(out)).detach()
        res2 = out
</code></pre>
<p>And since checkpoint has the behaviour that it makes all outputs require a gradient, when it tries to do it here I think it gets conflicted. Since in the model definition a gradient is not required but then the flag <code>requires_grad</code> is true. Hence it outputs the error <code>[4]&lt;stderr&gt;:RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn</code>. Or so I think is the case…</p>

albanD: <p>Ho good catch !<br>
Another quirk of the checkpoint module…<br>
So, here indeed, it might be simpler to remove this special code and detach() outside of the checkpoint.</p>
<p>Note that in this code, the <code>detach()</code> is not needed as the <code>no_grad()</code> block already prevent any graph from being created.</p>

Jennsi: <p>Hey guys <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>It seems that I have a similar problem. My model is made up of a backbone CNN (namely ResNet50) and some individual layers at the end. I added checkpoint.checkpoint_sequential() for the resnet and checkpoint.checkpoint() for the individual layers at the end. The performance of my model drops heavily.</p>
<p>So I checked requires_grad on the outputs of the checkpoints as proposed in this post and in the first epoch everything was fine, but in the second epoch it turned to False. I tried to add<br>
<code> requires_grad_ ()</code><br>
which resolved the problem for checkpoint.checkpoint().</p>
<p>However, for checkpoint.sequential_checkpoint it still gives False in the second epoch. The weird thing is, that during evaluation, i.e. with torch.no_grad(), requires_grad on the output of the checkpoints resolved to Ture <img src="https://discuss.pytorch.org/images/emoji/apple/face_with_raised_eyebrow.png?v=9" title=":face_with_raised_eyebrow:" class="emoji" alt=":face_with_raised_eyebrow:"></p>
<p>If anyone has any idea, I would be pretty thankful <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Edit: changed the checkpointing inside resnet to checkpoint.checkpoint too and not requires_grad is set to false for both after the checkpointing after the first epoch…</p>