From 5d7ebbe312ad028b2e0c094ad63df77d3b81e808 Mon Sep 17 00:00:00 2001 From: frgfm Date: Wed, 21 Oct 2020 17:49:08 +0200 Subject: [PATCH 1/6] style: Added annotation typing for shufflenet --- torchvision/models/shufflenetv2.py | 39 +++++++++++++++++++++--------- 1 file changed, 27 insertions(+), 12 deletions(-) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index 14f9521886c..2ccbfd142e6 100644 --- a/torchvision/models/shufflenetv2.py +++ b/torchvision/models/shufflenetv2.py @@ -1,6 +1,9 @@ import torch +from torch import Tensor import torch.nn as nn from .utils import load_state_dict_from_url +from typing import Callable, Any +from torch.jit.annotations import List __all__ = [ @@ -16,7 +19,7 @@ } -def channel_shuffle(x, groups): +def channel_shuffle(x: Tensor, groups: int) -> Tensor: # type: (torch.Tensor, int) -> torch.Tensor batchsize, num_channels, height, width = x.data.size() channels_per_group = num_channels // groups @@ -34,7 +37,12 @@ def channel_shuffle(x, groups): class InvertedResidual(nn.Module): - def __init__(self, inp, oup, stride): + def __init__( + self, + inp: int, + oup: int, + stride: int + ): super(InvertedResidual, self).__init__() if not (1 <= stride <= 3): @@ -68,10 +76,11 @@ def __init__(self, inp, oup, stride): ) @staticmethod - def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False): + def depthwise_conv(i: int, o: int, kernel_size: int, stride: int = 1, padding: int = 0, + bias: bool = False) -> nn.Conv2d: return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i) - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: if self.stride == 1: x1, x2 = x.chunk(2, dim=1) out = torch.cat((x1, self.branch2(x2)), dim=1) @@ -84,7 +93,13 @@ def forward(self, x): class ShuffleNetV2(nn.Module): - def __init__(self, stages_repeats, stages_out_channels, num_classes=1000, inverted_residual=InvertedResidual): + def __init__( + self, + stages_repeats: List[int], + stages_out_channels: List[int], + num_classes: int = 1000, + inverted_residual: Callable[..., nn.Module] = InvertedResidual + ): super(ShuffleNetV2, self).__init__() if len(stages_repeats) != 3: @@ -122,7 +137,7 @@ def __init__(self, stages_repeats, stages_out_channels, num_classes=1000, invert self.fc = nn.Linear(output_channels, num_classes) - def _forward_impl(self, x): + def _forward_impl(self, x: Tensor) -> Tensor: # See note [TorchScript super()] x = self.conv1(x) x = self.maxpool(x) @@ -134,11 +149,11 @@ def _forward_impl(self, x): x = self.fc(x) return x - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: return self._forward_impl(x) -def _shufflenetv2(arch, pretrained, progress, *args, **kwargs): +def _shufflenetv2(arch: str, pretrained: bool, progress: bool, *args: Any, **kwargs: Any) -> ShuffleNetV2: model = ShuffleNetV2(*args, **kwargs) if pretrained: @@ -152,7 +167,7 @@ def _shufflenetv2(arch, pretrained, progress, *args, **kwargs): return model -def shufflenet_v2_x0_5(pretrained=False, progress=True, **kwargs): +def shufflenet_v2_x0_5(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2: """ Constructs a ShuffleNetV2 with 0.5x output channels, as described in `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design" @@ -166,7 +181,7 @@ def shufflenet_v2_x0_5(pretrained=False, progress=True, **kwargs): [4, 8, 4], [24, 48, 96, 192, 1024], **kwargs) -def shufflenet_v2_x1_0(pretrained=False, progress=True, **kwargs): +def shufflenet_v2_x1_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2: """ Constructs a ShuffleNetV2 with 1.0x output channels, as described in `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design" @@ -180,7 +195,7 @@ def shufflenet_v2_x1_0(pretrained=False, progress=True, **kwargs): [4, 8, 4], [24, 116, 232, 464, 1024], **kwargs) -def shufflenet_v2_x1_5(pretrained=False, progress=True, **kwargs): +def shufflenet_v2_x1_5(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2: """ Constructs a ShuffleNetV2 with 1.5x output channels, as described in `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design" @@ -194,7 +209,7 @@ def shufflenet_v2_x1_5(pretrained=False, progress=True, **kwargs): [4, 8, 4], [24, 176, 352, 704, 1024], **kwargs) -def shufflenet_v2_x2_0(pretrained=False, progress=True, **kwargs): +def shufflenet_v2_x2_0(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ShuffleNetV2: """ Constructs a ShuffleNetV2 with 2.0x output channels, as described in `"ShuffleNet V2: Practical Guidelines for Efficient CNN Architecture Design" From 13b8d472ece3a946feb08d2c2e7cf2b3eaf87d5e Mon Sep 17 00:00:00 2001 From: frgfm Date: Wed, 21 Oct 2020 18:09:16 +0200 Subject: [PATCH 2/6] fix: Removed duplicate type hint --- torchvision/models/shufflenetv2.py | 1 - 1 file changed, 1 deletion(-) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index 2ccbfd142e6..1b424982ddb 100644 --- a/torchvision/models/shufflenetv2.py +++ b/torchvision/models/shufflenetv2.py @@ -20,7 +20,6 @@ def channel_shuffle(x: Tensor, groups: int) -> Tensor: - # type: (torch.Tensor, int) -> torch.Tensor batchsize, num_channels, height, width = x.data.size() channels_per_group = num_channels // groups From b37b1aa21b01a826d8fcf310cfc908653a4ffbc6 Mon Sep 17 00:00:00 2001 From: frgfm Date: Thu, 22 Oct 2020 12:08:30 +0200 Subject: [PATCH 3/6] refactor: Removed un-necessary import --- torchvision/models/shufflenetv2.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index 1b424982ddb..5f50bd6ddfb 100644 --- a/torchvision/models/shufflenetv2.py +++ b/torchvision/models/shufflenetv2.py @@ -2,8 +2,7 @@ from torch import Tensor import torch.nn as nn from .utils import load_state_dict_from_url -from typing import Callable, Any -from torch.jit.annotations import List +from typing import Callable, Any, List __all__ = [ From 959212ab148654e5f5df55157847680cb471f46a Mon Sep 17 00:00:00 2001 From: frgfm Date: Thu, 22 Oct 2020 21:20:29 +0200 Subject: [PATCH 4/6] fix: Fixed constructor typing --- torchvision/models/shufflenetv2.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index 5f50bd6ddfb..d83788f9dda 100644 --- a/torchvision/models/shufflenetv2.py +++ b/torchvision/models/shufflenetv2.py @@ -40,7 +40,7 @@ def __init__( inp: int, oup: int, stride: int - ): + ) -> None: super(InvertedResidual, self).__init__() if not (1 <= stride <= 3): @@ -97,7 +97,7 @@ def __init__( stages_out_channels: List[int], num_classes: int = 1000, inverted_residual: Callable[..., nn.Module] = InvertedResidual - ): + ) -> None: super(ShuffleNetV2, self).__init__() if len(stages_repeats) != 3: From b1d78c6a06d1f29624b61ab68a3ddccbc45cbbb7 Mon Sep 17 00:00:00 2001 From: frgfm Date: Thu, 22 Oct 2020 21:21:15 +0200 Subject: [PATCH 5/6] style: Added black formatting on depthwise_conv --- torchvision/models/shufflenetv2.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index d83788f9dda..d16457eba24 100644 --- a/torchvision/models/shufflenetv2.py +++ b/torchvision/models/shufflenetv2.py @@ -74,8 +74,14 @@ def __init__( ) @staticmethod - def depthwise_conv(i: int, o: int, kernel_size: int, stride: int = 1, padding: int = 0, - bias: bool = False) -> nn.Conv2d: + def depthwise_conv( + i: int, + o: int, + kernel_size: int, + stride: int = 1, + padding: int = 0, + bias: bool = False + ) -> nn.Conv2d: return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i) def forward(self, x: Tensor) -> Tensor: From ae76c5f9a9c36cacab3302819c58cbf9ae2e3d07 Mon Sep 17 00:00:00 2001 From: frgfm Date: Fri, 6 Nov 2020 10:42:04 +0100 Subject: [PATCH 6/6] style: Fixed stage typing in shufflenet --- torchvision/models/shufflenetv2.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/torchvision/models/shufflenetv2.py b/torchvision/models/shufflenetv2.py index d16457eba24..9ba090ad09b 100644 --- a/torchvision/models/shufflenetv2.py +++ b/torchvision/models/shufflenetv2.py @@ -123,6 +123,10 @@ def __init__( self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + # Static annotations for mypy + self.stage2: nn.Sequential + self.stage3: nn.Sequential + self.stage4: nn.Sequential stage_names = ['stage{}'.format(i) for i in [2, 3, 4]] for name, repeats, output_channels in zip( stage_names, stages_repeats, self._stage_out_channels[1:]):