pytorch / pytorch/vision

[proposal] Use self.flatten instead of torch.flatten and when becomes possible derive ResNet from nn.Sequential (scripting+quantization is blocker), would simplify model surgery in the most frequent cases

Open
#3,331 23 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

module: models needs discussion
Dominant language
Python
Stars
17.9k
Forks
7.3k
Avg merge
1d 15h
Merged PRs (30d)
13

Description

Currently In https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py#L243:

        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)

If it instead used x = self.flatten(x), then it would simplify model surgery: del model.avgpool, model.flatten, model.fc. Also in this case the class can just derive from Sequential and use OrderedDict to pass submodules (like in https://discuss.pytorch.org/t/ux-mix-of-nn-sequential-and-nn-moduledict/104724/2?u=vadimkantorov), this would preserve checkpoint compat as well. The method forward could then be removed

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start in torchvision/models/resnet.py around line 243 and review the linked PyTorch discussion about mixing nn.Sequential and nn.ModuleDict. Investigate how scripting and quantization constrain replacing torch.flatten and deriving ResNet from Sequential, including checkpoint compatibility. Done means the proposed model surgery is supported without breaking those constraints.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
computer-vision, machine-learning
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.