diff --git a/torchvision/models/vgg.py b/torchvision/models/vgg.py index dba534f651d..433b27e9d97 100644 --- a/torchvision/models/vgg.py +++ b/torchvision/models/vgg.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn from .utils import load_state_dict_from_url +from typing import Union, List, Dict, Any, cast __all__ = [ @@ -23,7 +24,12 @@ class VGG(nn.Module): - def __init__(self, features, num_classes=1000, init_weights=True): + def __init__( + self, + features: nn.Module, + num_classes: int = 1000, + init_weights: bool = True + ) -> None: super(VGG, self).__init__() self.features = features self.avgpool = nn.AdaptiveAvgPool2d((7, 7)) @@ -39,14 +45,14 @@ def __init__(self, features, num_classes=1000, init_weights=True): if init_weights: self._initialize_weights() - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.features(x) x = self.avgpool(x) x = torch.flatten(x, 1) x = self.classifier(x) return x - def _initialize_weights(self): + def _initialize_weights(self) -> None: for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu') @@ -60,13 +66,14 @@ def _initialize_weights(self): nn.init.constant_(m.bias, 0) -def make_layers(cfg, batch_norm=False): - layers = [] +def make_layers(cfg: List[Union[str, int]], batch_norm: bool = False) -> nn.Sequential: + layers: List[nn.Module] = [] in_channels = 3 for v in cfg: if v == 'M': layers += [nn.MaxPool2d(kernel_size=2, stride=2)] else: + v = cast(int, v) conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1) if batch_norm: layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)] @@ -76,7 +83,7 @@ def make_layers(cfg, batch_norm=False): return nn.Sequential(*layers) -cfgs = { +cfgs: Dict[str, List[Union[str, int]]] = { 'A': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 'B': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'], 'D': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'], @@ -84,7 +91,7 @@ def make_layers(cfg, batch_norm=False): } -def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs): +def _vgg(arch: str, cfg: str, batch_norm: bool, pretrained: bool, progress: bool, **kwargs: Any) -> VGG: if pretrained: kwargs['init_weights'] = False model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs) @@ -95,7 +102,7 @@ def _vgg(arch, cfg, batch_norm, pretrained, progress, **kwargs): return model -def vgg11(pretrained=False, progress=True, **kwargs): +def vgg11(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 11-layer model (configuration "A") from `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ @@ -106,7 +113,7 @@ def vgg11(pretrained=False, progress=True, **kwargs): return _vgg('vgg11', 'A', False, pretrained, progress, **kwargs) -def vgg11_bn(pretrained=False, progress=True, **kwargs): +def vgg11_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 11-layer model (configuration "A") with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ @@ -117,7 +124,7 @@ def vgg11_bn(pretrained=False, progress=True, **kwargs): return _vgg('vgg11_bn', 'A', True, pretrained, progress, **kwargs) -def vgg13(pretrained=False, progress=True, **kwargs): +def vgg13(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 13-layer model (configuration "B") `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ @@ -128,7 +135,7 @@ def vgg13(pretrained=False, progress=True, **kwargs): return _vgg('vgg13', 'B', False, pretrained, progress, **kwargs) -def vgg13_bn(pretrained=False, progress=True, **kwargs): +def vgg13_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 13-layer model (configuration "B") with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ @@ -139,7 +146,7 @@ def vgg13_bn(pretrained=False, progress=True, **kwargs): return _vgg('vgg13_bn', 'B', True, pretrained, progress, **kwargs) -def vgg16(pretrained=False, progress=True, **kwargs): +def vgg16(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 16-layer model (configuration "D") `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ @@ -150,7 +157,7 @@ def vgg16(pretrained=False, progress=True, **kwargs): return _vgg('vgg16', 'D', False, pretrained, progress, **kwargs) -def vgg16_bn(pretrained=False, progress=True, **kwargs): +def vgg16_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 16-layer model (configuration "D") with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ @@ -161,7 +168,7 @@ def vgg16_bn(pretrained=False, progress=True, **kwargs): return _vgg('vgg16_bn', 'D', True, pretrained, progress, **kwargs) -def vgg19(pretrained=False, progress=True, **kwargs): +def vgg19(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 19-layer model (configuration "E") `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_ @@ -172,7 +179,7 @@ def vgg19(pretrained=False, progress=True, **kwargs): return _vgg('vgg19', 'E', False, pretrained, progress, **kwargs) -def vgg19_bn(pretrained=False, progress=True, **kwargs): +def vgg19_bn(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> VGG: r"""VGG 19-layer model (configuration 'E') with batch normalization `"Very Deep Convolutional Networks For Large-Scale Image Recognition" `_