Can't use checkpoint_sequential

Xinyuan_Wang · June 15, 2023, 12:55am

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()

I got

/usr/local/lib/python3.10/dist-packages/torch/nn/modules/linear.py in forward(self, input)
112
113 def forward(self, input: Tensor) → Tensor:
→ 114 return F.linear(input, self.weight, self.bias)
115
116 def extra_repr(self) → str:

RuntimeError: mat1 and mat2 shapes cannot be multiplied (8192x1 and 2048x1000)

Xinyuan_Wang · June 15, 2023, 1:01am

How to use checkpoint_sequential for resnet50? Any suggestion? @ptrblck_de

Zeyuan.Yin · August 16, 2023, 11:13pm

checkpoint_sequential supports a list of modules/functions in order (sequentially).

Please check the resnet50 forward function:

github.com

pytorch/vision/blob/4cba51c5254407ec1e460a8cbb4ea06700e92637/torchvision/models/resnet.py#L266-L285


      
          def _forward_impl(self, x: Tensor) -> Tensor:
              # See note [TorchScript super()]
              x = self.conv1(x)
              x = self.bn1(x)
              x = self.relu(x)
              x = self.maxpool(x)
          
              x = self.layer1(x)
              x = self.layer2(x)
              x = self.layer3(x)
              x = self.layer4(x)
          
              x = self.avgpool(x)
              x = torch.flatten(x, 1)
              x = self.fc(x)
          
              return x
          
          def forward(self, x: Tensor) -> Tensor:
              return self._forward_impl(x)

x = torch.flatten(x, 1) is not included in your modules = [module for k, module in enumerate(model.children())].

You can check by printing len(modules), which is 10, but there are 11 operations in forward. Thus, x = torch.flatten(x, 1) is not considered in your modules.