Vision: #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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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/[email protected]/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:

https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/models/densenet.py#L74

https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/models/densenet.py#L80

https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/models/densenet.py#L124

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 https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/models/googlenet.py#L204 and https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/models/googlenet.py#L113-L114 are good, but
https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/models/densenet.py#L74 and https://github.com/pytorch/vision/blob/8c281757a0daf3a8e92cbb4bded0e5e6b389a375/torchvision/models/densenet.py#L80 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?