Added annotation typing to models._utils

[frgfm]

Hello folks 👋

As per #2025, annotation typing are welcome in torchvision. So I started with torchvision.models._utils in this PR. A few points of attention:

• I picked Callable[..., nn.Module] for module constructors (passing BatchNorm2d for instance), I'm unsure whether that is good choice? any recommendation?
• I only used black styling in class constructors for now, the rest are only annotation typing additions. Happy to change everything to black linting if you prefer.

Any feedback is welcome!

[pmeier]

Hey @frgfm, thanks a lot for the PR! Could you please split this PR into multiple as it is very hard to review with that many changes. I think it would be best to do a separate PR for every model / file.

I've reviewed the first three files and found a couple of things.

[frgfm]

@pmeier sure, I updated this PR to focus on models._utils. You got a dozen PRs incoming then haha

[pmeier]

That is alright. I'll go through them one by one.

[pmeier]

mypy CI is failing with

torchvision/models/_utils.py:60: error:
Signature of "forward" incompatible with supertype "ModuleDict" [override]
        def forward(self, x: torch.Tensor) ->...

[pmeier]

@frgfm Removing the annotation is not the way to go here. Either fix it or if it is unfixable (which might be the case here) use the correct annotations and add a #type: ignore[override] to the line.

[pmeier]

The test failures are related. I'll look into it and get back to you.

[pmeier]

@frgfm It seems torchscript is the offender here. I suspect it can't handle the type: ignore comment. Could you remove that to see if my suspicion is right?

If it is I'm not sure to move forward. We can't make torchscript and mypy happy.

[frgfm]

@pmeier Yes, I figured out the same thing yesterday unfortunately.
When I remove the ignore, mypy throws the following:

error: Signature of "forward" incompatible
with supertype "ModuleDict"  [override]

Let me know which way you prefer to proceed with!

[pmeier]

In general torchscript > mypy holds. Thus, while the type hints are nice to have, passing tests for torchscript is a requirement. Could you simply remove the type: ignore comment and push that change? It is not important for now if mypy fails, but the unit tests need to pass.

[fmassa]

Yes.

The situation with torchscript is delicate, and if we need to chose between mypy and torchscript, torchscript wins.

In general, all torchvision models already have type annotation in the forward because they all support torchscript. The difference with mypy is that torchscript assumes that

• if no input type is passed to an argument, it's assumed to be a tensor
• return types are inferred

So if you want to play safe, I would say to only add type annotations to the constructor of the models, or the model builder functions. This way it won't affect torchscript, but mypy won't probably be happy either.

Which means that it might be hard to enable full mypy testing to the models.

[pmeier]

This way it won't affect torchscript, but mypy won't probably be happy either.

Right now, we allow untyped definitions. Thus, if we simply remove the annotations from forward, mypy should be happy.

So if you want to play safe, I would say to only add type annotations to the constructor of the models, or the model builder functions.

@frgfm In that case you could simply revert the last two commits and this PR should be ready.

[pmeier]

LGTM. Thanks a lot @frgfm. Note that the further process will be delayed, since we are currently preparing the next release of torchvision.

[frgfm]

LGTM. Thanks a lot @frgfm. Note that the further process will be delayed, since we are currently preparing the next release of torchvision.

Should I update

from torch.jit.annotations import Dict

to

from typing import Dict

though?

[pmeier]

Yes that would be nice!

[fmassa]

Thanks!