Added annotation typing to shufflenet

[frgfm]

Hi there!

As per #2025, annotation typing are welcome in torchvision. So, this PR focuses on torchvision.models.shufflenet this PR!

Any feedback is welcome!

[pmeier]

Minor nitpick plus incremental changes from the other PRs. Otherwise LGTM.

[frgfm]

Still getting a troublesome error with mypy though:

torchvision/models/shufflenetv2.py:148: error: "Tensor" not callable  [operator]
            x = self.stage2(x)
                ^
torchvision/models/shufflenetv2.py:149: error: "Tensor" not callable  [operator]
            x = self.stage3(x)
                ^
torchvision/models/shufflenetv2.py:150: error: "Tensor" not callable  [operator]
            x = self.stage4(x)

Any ideas? I tried a few things in vain unfortunately

[fmassa]

I'll let @pmeier have a look, I don't know why mypy is thinking that stage2 is a Tensor. One thing to note is that stage2 is created programmatically with setattr, which might confuse mypy. Maybe there is a way to annotate stage{n} to tell mypy that it's a nn.Module ?

[pmeier]

I don't know why mypy is thinking that stage2 is a Tensor

Me neither. I think that might be caused by the way nn.Module is typed, but I'm certainly not sure.

One thing to note is that stage2 is created programmatically with setattr, which might confuse mypy. Maybe there is a way to annotate stage{n} to tell mypy that it's a nn.Module ?

We could simply include self.stage2: nn.Module in the constructor. Since mypy only inspects but does not the code we can't use a loop for that. Thus I propose to wrap the loop in a function an call it several times from the constructor.

[fmassa]

Thus I propose to wrap the loop in a function an call it several times from the constructor

I'm not exactly sure how to do it, given that we use setattr. So you propose to make the change that we instead do by hand

self.stage2 = create(...)

?

That is a big big change for just typing, can't we disable mypy somehow?

[pmeier]

can't we disable mypy somehow?

Sure we can. We could also casting in forward

x = cast(nn.Sequential, self.stage2)(x)

or put a block of static annotations above the creation loop

self.stage1: nn.Sequential
self.stage2: nn.Sequential

[facebook-github-bot]

Hi @frgfm!

Thank you for your pull request and welcome to our community. We require contributors to sign our Contributor License Agreement, and we don't seem to have you on file.

In order for us to review and merge your code, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

If you have received this in error or have any questions, please contact us at [email protected]. Thanks!

[fmassa]

Thanks!

[fmassa]

note: it would have been preferable to let stage{} be a nn.Module, as they could potentially be any type of nn.Module, not just Sequential