diff --git a/torchvision/models/inception.py b/torchvision/models/inception.py index a131c8c754f..9e49e446849 100644 --- a/torchvision/models/inception.py +++ b/torchvision/models/inception.py @@ -3,9 +3,9 @@ import torch import torch.nn as nn import torch.nn.functional as F -from torch.jit.annotations import Optional from torch import Tensor from .utils import load_state_dict_from_url +from typing import Callable, Any, Optional, Tuple, List __all__ = ['Inception3', 'inception_v3', 'InceptionOutputs', '_InceptionOutputs'] @@ -24,7 +24,7 @@ _InceptionOutputs = InceptionOutputs -def inception_v3(pretrained=False, progress=True, **kwargs): +def inception_v3(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> "Inception3": r"""Inception v3 model architecture from `"Rethinking the Inception Architecture for Computer Vision" `_. @@ -63,8 +63,14 @@ def inception_v3(pretrained=False, progress=True, **kwargs): class Inception3(nn.Module): - def __init__(self, num_classes=1000, aux_logits=True, transform_input=False, - inception_blocks=None, init_weights=None): + def __init__( + self, + num_classes: int = 1000, + aux_logits: bool = True, + transform_input: bool = False, + inception_blocks: Optional[List[Callable[..., nn.Module]]] = None, + init_weights: Optional[bool] = None + ) -> None: super(Inception3, self).__init__() if inception_blocks is None: inception_blocks = [ @@ -124,7 +130,7 @@ def __init__(self, num_classes=1000, aux_logits=True, transform_input=False, nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) - def _transform_input(self, x): + def _transform_input(self, x: Tensor) -> Tensor: if self.transform_input: x_ch0 = torch.unsqueeze(x[:, 0], 1) * (0.229 / 0.5) + (0.485 - 0.5) / 0.5 x_ch1 = torch.unsqueeze(x[:, 1], 1) * (0.224 / 0.5) + (0.456 - 0.5) / 0.5 @@ -132,7 +138,7 @@ def _transform_input(self, x): x = torch.cat((x_ch0, x_ch1, x_ch2), 1) return x - def _forward(self, x): + def _forward(self, x: Tensor) -> Tuple[Tensor, Optional[Tensor]]: # N x 3 x 299 x 299 x = self.Conv2d_1a_3x3(x) # N x 32 x 149 x 149 @@ -188,13 +194,13 @@ def _forward(self, x): return x, aux @torch.jit.unused - def eager_outputs(self, x: torch.Tensor, aux: Optional[Tensor]) -> InceptionOutputs: + def eager_outputs(self, x: Tensor, aux: Optional[Tensor]) -> InceptionOutputs: if self.training and self.aux_logits: return InceptionOutputs(x, aux) else: return x # type: ignore[return-value] - def forward(self, x): + def forward(self, x: Tensor) -> InceptionOutputs: x = self._transform_input(x) x, aux = self._forward(x) aux_defined = self.training and self.aux_logits @@ -208,7 +214,12 @@ def forward(self, x): class InceptionA(nn.Module): - def __init__(self, in_channels, pool_features, conv_block=None): + def __init__( + self, + in_channels: int, + pool_features: int, + conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: super(InceptionA, self).__init__() if conv_block is None: conv_block = BasicConv2d @@ -223,7 +234,7 @@ def __init__(self, in_channels, pool_features, conv_block=None): self.branch_pool = conv_block(in_channels, pool_features, kernel_size=1) - def _forward(self, x): + def _forward(self, x: Tensor) -> List[Tensor]: branch1x1 = self.branch1x1(x) branch5x5 = self.branch5x5_1(x) @@ -239,14 +250,18 @@ def _forward(self, x): outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool] return outputs - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: outputs = self._forward(x) return torch.cat(outputs, 1) class InceptionB(nn.Module): - def __init__(self, in_channels, conv_block=None): + def __init__( + self, + in_channels: int, + conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: super(InceptionB, self).__init__() if conv_block is None: conv_block = BasicConv2d @@ -256,7 +271,7 @@ def __init__(self, in_channels, conv_block=None): self.branch3x3dbl_2 = conv_block(64, 96, kernel_size=3, padding=1) self.branch3x3dbl_3 = conv_block(96, 96, kernel_size=3, stride=2) - def _forward(self, x): + def _forward(self, x: Tensor) -> List[Tensor]: branch3x3 = self.branch3x3(x) branch3x3dbl = self.branch3x3dbl_1(x) @@ -268,14 +283,19 @@ def _forward(self, x): outputs = [branch3x3, branch3x3dbl, branch_pool] return outputs - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: outputs = self._forward(x) return torch.cat(outputs, 1) class InceptionC(nn.Module): - def __init__(self, in_channels, channels_7x7, conv_block=None): + def __init__( + self, + in_channels: int, + channels_7x7: int, + conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: super(InceptionC, self).__init__() if conv_block is None: conv_block = BasicConv2d @@ -294,7 +314,7 @@ def __init__(self, in_channels, channels_7x7, conv_block=None): self.branch_pool = conv_block(in_channels, 192, kernel_size=1) - def _forward(self, x): + def _forward(self, x: Tensor) -> List[Tensor]: branch1x1 = self.branch1x1(x) branch7x7 = self.branch7x7_1(x) @@ -313,14 +333,18 @@ def _forward(self, x): outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool] return outputs - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: outputs = self._forward(x) return torch.cat(outputs, 1) class InceptionD(nn.Module): - def __init__(self, in_channels, conv_block=None): + def __init__( + self, + in_channels: int, + conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: super(InceptionD, self).__init__() if conv_block is None: conv_block = BasicConv2d @@ -332,7 +356,7 @@ def __init__(self, in_channels, conv_block=None): self.branch7x7x3_3 = conv_block(192, 192, kernel_size=(7, 1), padding=(3, 0)) self.branch7x7x3_4 = conv_block(192, 192, kernel_size=3, stride=2) - def _forward(self, x): + def _forward(self, x: Tensor) -> List[Tensor]: branch3x3 = self.branch3x3_1(x) branch3x3 = self.branch3x3_2(branch3x3) @@ -345,14 +369,18 @@ def _forward(self, x): outputs = [branch3x3, branch7x7x3, branch_pool] return outputs - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: outputs = self._forward(x) return torch.cat(outputs, 1) class InceptionE(nn.Module): - def __init__(self, in_channels, conv_block=None): + def __init__( + self, + in_channels: int, + conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: super(InceptionE, self).__init__() if conv_block is None: conv_block = BasicConv2d @@ -369,7 +397,7 @@ def __init__(self, in_channels, conv_block=None): self.branch_pool = conv_block(in_channels, 192, kernel_size=1) - def _forward(self, x): + def _forward(self, x: Tensor) -> List[Tensor]: branch1x1 = self.branch1x1(x) branch3x3 = self.branch3x3_1(x) @@ -393,24 +421,29 @@ def _forward(self, x): outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool] return outputs - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: outputs = self._forward(x) return torch.cat(outputs, 1) class InceptionAux(nn.Module): - def __init__(self, in_channels, num_classes, conv_block=None): + def __init__( + self, + in_channels: int, + num_classes: int, + conv_block: Optional[Callable[..., nn.Module]] = None + ) -> None: super(InceptionAux, self).__init__() if conv_block is None: conv_block = BasicConv2d self.conv0 = conv_block(in_channels, 128, kernel_size=1) self.conv1 = conv_block(128, 768, kernel_size=5) - self.conv1.stddev = 0.01 + self.conv1.stddev = 0.01 # type: ignore[assignment] self.fc = nn.Linear(768, num_classes) - self.fc.stddev = 0.001 + self.fc.stddev = 0.001 # type: ignore[assignment] - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: # N x 768 x 17 x 17 x = F.avg_pool2d(x, kernel_size=5, stride=3) # N x 768 x 5 x 5 @@ -430,12 +463,17 @@ def forward(self, x): class BasicConv2d(nn.Module): - def __init__(self, in_channels, out_channels, **kwargs): + def __init__( + self, + in_channels: int, + out_channels: int, + **kwargs: Any + ) -> None: super(BasicConv2d, self).__init__() self.conv = nn.Conv2d(in_channels, out_channels, bias=False, **kwargs) self.bn = nn.BatchNorm2d(out_channels, eps=0.001) - def forward(self, x): + def forward(self, x: Tensor) -> Tensor: x = self.conv(x) x = self.bn(x) return F.relu(x, inplace=True)