[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

[vadimkantorov]

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

[datumbox]

Thanks for the proposal. Replacing torch.flatten with nn.Flatten seems it won't break anything but could you describe a bit more your use-case and how it will help you achieve what you want?

Concerning doing a major refactoring and inheriting from Sequential, this might have side-effects on all the other models that depend on resnet (segmentation, object detection etc), so I'm not sure if it can be done in a backward-compatible manner. Note also that replacing the forward is not something we can do due to how Quantization works:

vision/torchvision/models/quantization/resnet.py

Lines 93 to 100 in 6116812

	def forward(self, x):
	x = self.quant(x)
	# Ensure scriptability
	# super(QuantizableResNet,self).forward(x)
	# is not scriptable
	x = self._forward_impl(x)
	x = self.dequant(x)
	return x

@fmassa Let me know if you see any problems about replacing flatten.

[fmassa]

I would be fine replacing torch.flatten with nn.Flatten, although it would only simplify model surgery if we were to make further modifications to the ResNet model as you suggested (making it inherit from nn.Sequential). Without this, only del model.fc etc wouldn't be enough.

So from this perspective, there is limited value in replacing torch.flatten with nn.Flatten, it being only syntactic sugar from one another.

[vadimkantorov]

Oh, I didn't realize:

# Ensure scriptability 
# super(QuantizableResNet,self).forward(x) 
# is not scriptable 

If this becomes fixed, I guess inheriting from nn.Sequential should be fine.

My real usecase: I'm working with a codebase https://github.com/ajabri/videowalk/blob/master/code/resnet.py#L41 that had to reimplement resnet forward because it wants to do simple model surgery

[fmassa]

@vadimkantorov model surgery in PyTorch generally requires re-writing forward or other parts of the model, except for nn.Sequential, so I would say it's a valid requirement to ask the users to be a bit more verbose

[vadimkantorov]

model surgery in PyTorch generally requires re-writing forward or other parts of the model, except for nn.Sequential

Of course. I understand that in general it's the case, but if it can be made simpler, I think it should be made simpler even if in general case full re-writing is required. In this particular case, it seems that because of quantization scripting limitations deriving from nn.Sequential is a no-go, but when that's solved it could be nice to derive ResNet from nn.Sequential.

I'll rename thie issue accordingly

[vadimkantorov]

It seems that deriving from nn.Sequential would also be good from DeepSpeed-compat standpoint: pytorch/pytorch#51574

For quantization case, couldn't it just insert quant/dequant-calling modules in the beginning and the end of sequential? then everything should work and no modification of forward would be needed

[vadimkantorov]

They are already modules:

        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()

One would just be required to modify the constructor and insert them before conv1 and after layer4

[vadimkantorov]

@fmassa Even if we just replace torch.flatten by self.flatten, it already simplifies model surgery by model.avgpool = model.flatten = model.fc = nn.Identity() (instead of del model.avgpool, model.flatten, model.fc)