#2860 regressed torch.jit.script(densenet161).save

[lutzroeder]

🐛 Bug

Traceback (most recent call last):
  File "repro.py", line 7, in <module>
    torch.jit.script(model).save('densenet161.pt')
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_script.py", line 911, in script
    return torch.jit._recursive.create_script_module(
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 370, in create_script_module
    return create_script_module_impl(nn_module, concrete_type, stubs_fn)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 426, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_script.py", line 388, in _construct
    init_fn(script_module)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 406, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 426, in create_script_module_impl
    script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_script.py", line 388, in _construct
    init_fn(script_module)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 406, in init_fn
    scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 382, in create_script_module_impl
    method_stubs = stubs_fn(nn_module)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 618, in infer_methods_to_compile
    stubs.append(make_stub_from_method(nn_module, method))
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 52, in make_stub_from_method
    return make_stub(func, method_name)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/_recursive.py", line 37, in make_stub
    ast = get_jit_def(func, name, self_name="RecursiveScriptModule")
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/frontend.py", line 259, in get_jit_def
    return build_def(ctx, fn_def, type_line, def_name, self_name=self_name)
  File "/usr/local/opt/python@3.8/Frameworks/Python.framework/Versions/3.8/lib/python3.8/site-packages/torch/jit/frontend.py", line 288, in build_def
    type_comment_decl = torch._C.parse_type_comment(type_line)
RuntimeError: expected type comment but found 'def' here:
def forward(self, init_features: Tensor) -> Tensor:  # type: ignore[override]
~~~ <--- HERE

To Reproduce

Steps to reproduce the behavior:

1. Install nightly build pip install --upgrade --pre torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
2. Run:

import torch
import torchvision
model = torchvision.models.densenet161(pretrained=False)
model.eval()
torch.jit.script(model).save('densenet161.pt')

Environment

Collecting environment information...
PyTorch version: 1.8.0.dev20201118
Is debug build: False
OS: Mac OSX 11.0.1 (x86_64)
GCC version: Could not collect
Clang version: 12.0.0 (clang-1200.0.32.27)
CMake version: version 3.18.4
Python version: 3.8 (64-bit runtime)
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Versions of relevant libraries:
[pip] numpy==1.19.0
[pip] torch==1.8.0.dev20201118
[pip] torchtext==0.7.0
[pip] torchvision==0.9.0.dev20201118
[conda] Could not collect

Additional context

@frgfm @fmassa @pmeier #2860

[pmeier]

The offenders are the type: ignore comments in the following three lines:

vision/torchvision/models/densenet.py

Line 74 in 8c28175

@torch.jit._overload_method # type: ignore[no-redef] # noqa: F811

vision/torchvision/models/densenet.py

Line 80 in 8c28175

def forward(self, input: Tensor) -> Tensor: # type: ignore[no-redef] # noqa: F811

vision/torchvision/models/densenet.py

Line 124 in 8c28175

def forward(self, init_features: Tensor) -> Tensor: # type: ignore[override]

They are only needed for the mypy check and can be removed safely. Not sure why our tests didn't catch this. We should check this carefully to avoid something like this in the future.

[frgfm]

@pmeier let me know if you need me to remove or adapt the typing ignore from densenet or the other similar PRs!

[fmassa]

The breakage is directly with the torchscript convertion, not even save.

@frgfm for now, could you please send a PR removing the ignore types for mypy? If you are busy today it's ok, and I'll revert the PR.

I'll investigate why the tests didn't fail here for torchscript, they might be skipped for some reason

[fmassa]

Also, we have type: ignore in many models, we should remove all those that are not in the __init__ of the model.

[pmeier]

Also, we have type: ignore in many models, we should remove all those that are not in the __init__ of the model.

We should only remove these comments on methods that are called from forward. I think for now "not forward" and "init" are the same, but this does not have to be the case.

[fmassa]

@pmeier we should remove all that are not in a torchscript codepath actually.

So __init__ and functions marked with @torch.jit.unused can have it, but not functions called from forward which are not marked with torch.jit.unused or torch.jit.ignore.

So

vision/torchvision/models/googlenet.py

Line 204 in 8c28175

return x # type: ignore[return-value]

and

vision/torchvision/models/googlenet.py

Lines 113 to 114 in 8c28175

	self.aux1 = None # type: ignore[assignment]
	self.aux2 = None # type: ignore[assignment]

are good, but

vision/torchvision/models/densenet.py

Line 74 in 8c28175

@torch.jit._overload_method # type: ignore[no-redef] # noqa: F811

and

vision/torchvision/models/densenet.py

Line 80 in 8c28175

def forward(self, input: Tensor) -> Tensor: # type: ignore[no-redef] # noqa: F811

are not.

So indeed, after a second look, only DenseNet seems to be affected.

[frgfm]

@fmassa @pmeier Alright, I'm on it!
So, for now, the PR should only remove the typing exception (and typing if necessary) in DenseNet (line 74 + all forward with # type: ignore), correct?