Use of torch.utils.checkpoint.checkpoint causes simple model to diverge

ResidentMario: <p>I have defined a simple model that trains on the CIFAR10 dataset included in <code>torchvision</code>. This model converges to a loss of ~0.72 after 10 epochs of training. Here is the model definition:</p>
<pre><code class="lang-python">class CIFAR10Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn_block_1 = nn.Sequential(*[
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Dropout(0.25)
        ])
        self.cnn_block_2 = nn.Sequential(*[
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Dropout(0.25)
        ])
        self.flatten = lambda inp: torch.flatten(inp, 1)
        self.head = nn.Sequential(*[
            nn.Linear(64 * 8 * 8, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, 10)
        ])
    
    def forward(self, X):
        X = self.cnn_block_1(X)
        X = self.cnn_block_2(X)
        X = self.flatten(X)
        X = self.head(X)
        return X
</code></pre>
<p>I then defined and trained a new version of this model, which has been updated to use <a href="https://pytorch.org/docs/stable/checkpoint.html" rel="noopener nofollow ugc"><code>torch.utils.checkpoint.checkpoint</code></a>. This model has the following definition (updates were made to exclude <code>nn.Dropout</code> from checkpointing, <a href="https://github.com/prigoyal/pytorch_memonger/blob/master/tutorial/Checkpointing_for_PyTorch_models.ipynb" rel="noopener nofollow ugc">since Dropout is incompatible with checkpointing</a>):</p>
<pre><code class="lang-python">class CIFAR10Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.cnn_block_1 = nn.Sequential(*[
            nn.Conv2d(3, 32, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        ])
        self.dropout_1 = nn.Dropout(0.25)
        self.cnn_block_2 = nn.Sequential(*[
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)
        ])
        self.dropout_2 = nn.Dropout(0.25)
        self.flatten = lambda inp: torch.flatten(inp, 1)
        self.linearize = nn.Sequential(*[
            nn.Linear(64 * 8 * 8, 512),
            nn.ReLU()
        ])
        self.dropout_3 = nn.Dropout(0.5)
        self.out = nn.Linear(512, 10)
    
    def forward(self, X):
        X = torch.utils.checkpoint.checkpoint(self.cnn_block_1, X)
        X = self.dropout_1(X)
        X = torch.utils.checkpoint.checkpoint(self.cnn_block_2, X)
        X = self.dropout_2(X)
        X = self.flatten(X)
        X = self.linearize(X)
        X = self.dropout_3(X)
        X = self.out(X)
        return X
</code></pre>
<p>However, this model achieves radically worse performance than its uncheckpointed peer: just ~1.59 loss after 10 epochs!</p>
<p>I don’t understand why this is the case. This is my first time using this API; perhaps someone here who understands it better than I can point out where my error is?</p>
<p>For reference, here is the full code and logs: <a href="https://gist.github.com/ResidentMario/e3254172b4706191089bb63ecd610e21" rel="noopener nofollow ugc">uncheckpointed version</a>, <a href="https://gist.github.com/ResidentMario/9c3a90504d1a027aab926fd65ae08139" rel="noopener nofollow ugc">checkpointed version</a>.</p>

ResidentMario: <p>The <a href="https://pytorch.org/docs/stable/notes/autograd.html" rel="noopener nofollow ugc">PyTorch autograd docs state</a>:</p>
<blockquote>
<p>If there’s a single input to an operation that requires gradient, its output will also require gradient. Conversely, only if all inputs don’t require gradient, the output also won’t require it.</p>
</blockquote>
<p>The input to a model like this one will be a vector with <code>no_grad=False</code> (since we are performing gradient descent relative to the model weights, not relative to the values of the input sample itself). Unless you manually freeze it, any module you apply to that vector containing learnable parameters (e.g. <code>nn.Conv2d</code> or <code>nn.Linear</code>, but not <code>nn.Dropout</code> or <code>nn.MaxPool2d</code>) will upgrade it to vector with <code>no_grad=True</code> automatically.</p>
<p><strong>Applying checkpointing will not do this for you</strong>. Instead, the way that checkpointing is implemented, whether the output vector produced by the checkpointed module is <code>no_grad=True</code> (and thus subject to gradient updates) or <code>no_grad=False</code> (and thus frozen) is solely determined by whether the <em>input</em> vector is <code>no_grad=True</code> or <code>no_grad=False</code>.</p>
<p>This has the side effect (?) that the first block in the model <em>cannot</em> be checkpointed. E.g. in this model the input to <code>self.cnn_block_1</code> is the input vector, which has <code>no_grad=False</code>; so setting <code>X = torch.utils.checkpoint.checkpoint(self.cnn_block_1, X)</code> causes every gradient in <code>cnn_block_1</code> to stay fixed.</p>
<p>I don’t think I would have been able to figure this out on my own; I got to this conclusion by Googling around and stumbling across <a href="https://realmyang.com/squeeze-most-out-of-your-gpu-checkpoint-in-pytorch/" rel="noopener nofollow ugc">this article</a>, which mentions this issue in the footnotes. Seems like something that should really be added to the docs.</p>
<p>Removing the checkpointing on <code>self.cnn_block_1</code> fixed this problem and produced a model that converged correctly again.</p>

ResidentMario: <p>Hopefully the next person that runs into this issue will find this thread and not be so confused. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

raining_day513: <p>It seems that this problem has been fixed. See this <a href="https://discuss.pytorch.org/t/checkpoint-with-no-grad-requiring-inputs-problem/19117/19">reply</a> and maybe the entire thread.</p>