Unable to call `super` method with TorchScript

[sharvil]

🐛 Bug

I have a derived nn.Module which calls super.forward(...) in its own implementation. When I try to compile the code to TorchScript, I get:

Tried to access nonexistent attribute or method 'forward' of type 'Tensor'.:
  File "test.py", line 7
  def forward(self, x):
    return super().forward(x)
           ~~~~~~~~~~~~~ <--- HERE

To Reproduce

import torch
import torch.nn as nn


class Foo(nn.Conv1d):
  def forward(self, x):
    return super().forward(x)


model = Foo(128, 128, 1)
torch.jit.script(model)

Expected behavior

I would expect to get a proper ScriptModule out of the torch.jit.script(...) line.

Environment

PyTorch version: 1.6.0
Is debug build: No
CUDA used to build PyTorch: 10.2

OS: Ubuntu 18.04.4 LTS
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
CMake version: Could not collect

Python version: 3.8
Is CUDA available: Yes
CUDA runtime version: 10.2.89
GPU models and configuration:
GPU 0: GeForce GTX 1080 Ti
GPU 1: GeForce GTX 1080 Ti
GPU 2: GeForce GTX 1080 Ti
GPU 3: GeForce GTX 1080 Ti

Nvidia driver version: 440.33.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.5.1

Versions of relevant libraries:
[pip3] numpy==1.19.1
[pip3] torch==1.6.0
[conda] _tflow_select             2.3.0                       mkl
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               10.2.89              hfd86e86_1
[conda] mkl                       2020.1                      217
[conda] mkl-service               2.3.0            py38he904b0f_0
[conda] mkl_fft                   1.1.0            py38h23d657b_0
[conda] mkl_random                1.1.1            py38h0573a6f_0
[conda] numpy                     1.19.1           py38hbc911f0_0
[conda] numpy-base                1.19.1           py38hfa32c7d_0
[conda] pytorch                   1.6.0           py3.8_cuda10.2.89_cudnn7.6.5_0    pytorch
[conda] tensorflow                2.2.0           mkl_py38h6d3daf0_0
[conda] tensorflow-base           2.2.0           mkl_py38h5059a2d_0

cc @suo @gmagogsfm

[gmagogsfm]

Hi sharvil,

As far as I know, torch.nn.Module doesn't really have a "forward" method. So even if calling super() works, it will fail finding forward. Could you help me understand what you are trying to do?

[sharvil]

@gmagogsfm, even though nn.Module doesn't have a forward method defined directly on it, subclasses such as nn.Conv1d do. As shown in the code snippet above, my class derives from nn.Conv1d and attempts to invoke super().forward() in its own forward method. This implementation behaves correctly in Python but does not compile down to TorchScript.

I would like to be able to use modules that augment their superclass' behavior. One could argue that the augmentation should be done through composition instead of inheritance, but composition has implications on serialization which is not always desirable. More concretely:

class MySpecialPaddedConv(nn.Conv1d):
  def forward(self, x):
    x = torch.cat([torch.randn(x.shape[0], x.shape[1], 128), x], dim=2)
    return super().forward(x)

could be written as:

class MySpecialPaddedConv(nn.Module):
  def __init__(self, *args, **kwargs):
    self.inner_conv = nn.Conv1d(*args, **kwargs)

  def forward(self, x):
    x = torch.cat([torch.randn(x.shape[0], x.shape[1], 128), x], dim=2)
    return self.inner_conv(x)

but the latter will introduce one more layer in MySpecialPaddedConv.state_dict for inner_conv.

[gmagogsfm]

Currently inheritance is not supported in TorchScript, we do have discussions about adding it, but doesn't look like it would land soonish. I can see that using composition adds one more layer in state_dict, but I don't quite see why that's undesirable. Could you clarify a bit more? Maybe we could figure out a good work-around for you.

[sharvil]

It's often the case that we have modules that decorate other modules with additional behavior (as above). Inheritance is preferred so we don't have parameters with silly serialized names like transformed_conv.padded_conv.conv. It's much easier to iterate on a project when serialized names are more controllable.

I know there are workarounds like remapping after loading a checkpoint or using tools like https://github.com/lmnt-com/graftr. But it would be nicer to not be in that situation in the first place.