Can't use checkpoint_sequential

Xinyuan_Wang: <pre><code class="lang-auto">import torch
from torch import nn
from torch.utils.checkpoint import checkpoint_sequential
from torchvision.models import resnet50

model = resnet50()

modules = [module for k, module in enumerate(model.children())]

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = model.to(device)
inputs = torch.randn(4, 3, 224, 224).to(device)
inputs.requires_grad = True
output = checkpoint_sequential(modules, len(modules), inputs)
criterion = nn.CrossEntropyLoss()
target = torch.empty(1, dtype=torch.long).random_(5).to(device)
loss = criterion(output, target)
loss.backward()
optimizer.step()
</code></pre>
<p>I got <img src="https://discuss.pytorch.org/images/emoji/apple/arrow_down.png?v=12" title=":arrow_down:" class="emoji" alt=":arrow_down:" loading="lazy" width="20" height="20"></p>
<p>/usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py in forward(self, input)<br>
112<br>
113     def forward(self, input: Tensor) → Tensor:<br>
 → 114         return F.linear(input, self.weight, self.bias)<br>
115<br>
116     def extra_repr(self) → str:</p>
<p>RuntimeError: mat1 and mat2 shapes cannot be multiplied (8192x1 and 2048x1000)</p>

Xinyuan_Wang: <p>How to use checkpoint_sequential for resnet50? Any suggestion? <img src="https://discuss.pytorch.org/images/emoji/apple/sob.png?v=12" title=":sob:" class="emoji" alt=":sob:" loading="lazy" width="20" height="20"> <span class="mention">@ptrblck_de</span></p>

Zeyuan.Yin: <p><code>checkpoint_sequential </code> supports <strong>a list of modules/functions in order (sequentially)</strong>.</p>
<p>Please check the resnet50 forward function:</p><aside class="onebox githubblob" data-onebox-src="https://github.com/pytorch/vision/blob/4cba51c5254407ec1e460a8cbb4ea06700e92637/torchvision/models/resnet.py#L266-L285">
  <header class="source">

      <a href="https://github.com/pytorch/vision/blob/4cba51c5254407ec1e460a8cbb4ea06700e92637/torchvision/models/resnet.py#L266-L285" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/vision/blob/4cba51c5254407ec1e460a8cbb4ea06700e92637/torchvision/models/resnet.py#L266-L285" target="_blank" rel="noopener nofollow ugc">pytorch/vision/blob/4cba51c5254407ec1e460a8cbb4ea06700e92637/torchvision/models/resnet.py#L266-L285</a></h4>



    <pre class="onebox"><code class="lang-py">
      <ol class="start lines" start="266" style="counter-reset: li-counter 265 ;">
          <li>def _forward_impl(self, x: Tensor) -&gt; Tensor:</li>
          <li>    # See note [TorchScript super()]</li>
          <li>    x = self.conv1(x)</li>
          <li>    x = self.bn1(x)</li>
          <li>    x = self.relu(x)</li>
          <li>    x = self.maxpool(x)</li>
          <li></li>
          <li>    x = self.layer1(x)</li>
          <li>    x = self.layer2(x)</li>
          <li>    x = self.layer3(x)</li>
          <li>    x = self.layer4(x)</li>
          <li></li>
          <li>    x = self.avgpool(x)</li>
          <li>    x = torch.flatten(x, 1)</li>
          <li>    x = self.fc(x)</li>
          <li></li>
          <li>    return x</li>
          <li></li>
          <li>def forward(self, x: Tensor) -&gt; Tensor:</li>
          <li>    return self._forward_impl(x)</li>
      </ol>
    </code></pre>



  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p><code> x = torch.flatten(x, 1)</code> is not included in your <code>modules = [module for k, module in enumerate(model.children())]</code>.</p>
<p>You can check by printing <code>len(modules)</code>, which is 10, but there are 11 operations in <code>forward</code>. Thus, <code>x = torch.flatten(x, 1)</code> is not considered in your modules.</p>