[Typing][A-81] Add type annotations for paddle/vision/models/resnet.py (#65487)

Asthestarsfalll · web-flow · commit 10a4dcc31c63 · 2024-06-28T03:00:00.000+08:00
diff --git a/python/paddle/vision/models/mobilenetv1.py b/python/paddle/vision/models/mobilenetv1.py
@@ -35,7 +35,6 @@
 
 
 class _MobileNetV1Options(TypedDict):
-    scale: NotRequired[float]
     num_classes: NotRequired[int]
     with_pool: NotRequired[bool]
 
diff --git a/python/paddle/vision/models/mobilenetv2.py b/python/paddle/vision/models/mobilenetv2.py
@@ -42,7 +42,6 @@
 
 
 class _MobileNetV2Options(TypedDict):
-    scale: NotRequired[float]
     num_classes: NotRequired[int]
     with_pool: NotRequired[bool]
 
@@ -54,7 +53,7 @@ def __init__(
         oup: int,
         stride: int,
         expand_ratio: float,
-        norm_layer: nn.Layer = nn.BatchNorm2D,
+        norm_layer: type[nn.Layer] = nn.BatchNorm2D,
     ) -> None:
         super().__init__()
         self.stride = stride
diff --git a/python/paddle/vision/models/mobilenetv3.py b/python/paddle/vision/models/mobilenetv3.py
@@ -47,7 +47,6 @@
 
 
 class _MobileNetV3Options(TypedDict):
-    scale: NotRequired[float]
     num_classes: NotRequired[int]
     with_pool: NotRequired[bool]
 
@@ -70,8 +69,8 @@ def __init__(
         self,
         input_channels: int,
         squeeze_channels: int,
-        activation: nn.Layer = nn.ReLU,
-        scale_activation: nn.Layer = nn.Sigmoid,
+        activation: type[nn.Layer] = nn.ReLU,
+        scale_activation: type[nn.Layer] = nn.Sigmoid,
     ) -> None:
         super().__init__()
         self.avgpool = nn.AdaptiveAvgPool2D(1)
diff --git a/python/paddle/vision/models/resnet.py b/python/paddle/vision/models/resnet.py
@@ -11,14 +11,49 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
 
 import paddle
 from paddle import nn
+from paddle._typing import Size2
 from paddle.utils.download import get_weights_path_from_url
 
+if TYPE_CHECKING:
+    from typing import Literal, TypedDict
+
+    from typing_extensions import NotRequired, Unpack
+
+    from paddle import Tensor
+
+    _ResNetArch = Literal[
+        'resnet18',
+        'resnet34',
+        'resnet50',
+        'resnet101',
+        'resnet152',
+        'resnext50_32x4d',
+        'resnext50_64x4d',
+        'resnext50_64x4d',
+        'resnext101_32x4d',
+        'resnext101_64x4d',
+        'resnext152_32x4d',
+        'resnext152_64x4d',
+        'wide_resnet50_2',
+        'wide_resnet101_2',
+    ]
+
+    class _ResNetOptions(TypedDict):
+        width: NotRequired[int]
+        num_classes: NotRequired[int]
+        with_pool: NotRequired[bool]
+        groups: NotRequired[int]
+
+
 __all__ = []
 
-model_urls = {
+model_urls: dict[str, tuple[str, str]] = {
     'resnet18': (
         'https://paddle-hapi.bj.bcebos.com/models/resnet18.pdparams',
         'cf548f46534aa3560945be4b95cd11c4',
@@ -75,22 +110,22 @@
 
 
 class BasicBlock(nn.Layer):
-    expansion = 1
+    expansion: int = 1
 
     def __init__(
         self,
-        inplanes,
-        planes,
-        stride=1,
-        downsample=None,
-        groups=1,
-        base_width=64,
-        dilation=1,
-        norm_layer=None,
-    ):
+        inplanes: int,
+        planes: int,
+        stride: Size2 = 1,
+        downsample: nn.Layer | None = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: type[nn.Layer] | None = None,
+    ) -> None:
         super().__init__()
         if norm_layer is None:
-            norm_layer = nn.BatchNorm2D
+            norm_layer: type[nn.BatchNorm2D] = nn.BatchNorm2D
 
         if dilation > 1:
             raise NotImplementedError(
@@ -107,7 +142,7 @@ def __init__(
         self.downsample = downsample
         self.stride = stride
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         identity = x
 
         out = self.conv1(x)
@@ -127,22 +162,22 @@ def forward(self, x):
 
 
 class BottleneckBlock(nn.Layer):
-    expansion = 4
+    expansion: int = 4
 
     def __init__(
         self,
-        inplanes,
-        planes,
-        stride=1,
-        downsample=None,
-        groups=1,
-        base_width=64,
-        dilation=1,
-        norm_layer=None,
-    ):
+        inplanes: int,
+        planes: int,
+        stride: Size2 = 1,
+        downsample: nn.Layer | None = None,
+        groups: int = 1,
+        base_width: int = 64,
+        dilation: int = 1,
+        norm_layer: type[nn.Layer] | None = None,
+    ) -> None:
         super().__init__()
         if norm_layer is None:
-            norm_layer = nn.BatchNorm2D
+            norm_layer: type[nn.BatchNorm2D] = nn.BatchNorm2D
         width = int(planes * (base_width / 64.0)) * groups
 
         self.conv1 = nn.Conv2D(inplanes, width, 1, bias_attr=False)
@@ -168,7 +203,7 @@ def __init__(
         self.downsample = downsample
         self.stride = stride
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         identity = x
 
         out = self.conv1(x)
@@ -233,15 +268,22 @@ class ResNet(nn.Layer):
             [1, 1000]
     """
 
+    groups: int
+    base_width: int
+    num_classes: int
+    with_pool: bool
+    inplanes: int
+    dilation: int
+
     def __init__(
         self,
-        block,
-        depth=50,
-        width=64,
-        num_classes=1000,
-        with_pool=True,
-        groups=1,
-    ):
+        block: type[BasicBlock | BottleneckBlock],
+        depth: int = 50,
+        width: int = 64,
+        num_classes: int = 1000,
+        with_pool: bool = True,
+        groups: int = 1,
+    ) -> None:
         super().__init__()
         layer_cfg = {
             18: [2, 2, 2, 2],
@@ -281,7 +323,14 @@ def __init__(
         if num_classes > 0:
             self.fc = nn.Linear(512 * block.expansion, num_classes)
 
-    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+    def _make_layer(
+        self,
+        block: type[BasicBlock | BottleneckBlock],
+        planes: int,
+        blocks: int,
+        stride: int = 1,
+        dilate: bool = False,
+    ) -> nn.Sequential:
         norm_layer = self._norm_layer
         downsample = None
         previous_dilation = self.dilation
@@ -327,7 +376,7 @@ def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
 
         return nn.Sequential(*layers)
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.conv1(x)
         x = self.bn1(x)
         x = self.relu(x)
@@ -347,7 +396,13 @@ def forward(self, x):
         return x
 
 
-def _resnet(arch, Block, depth, pretrained, **kwargs):
+def _resnet(
+    arch: _ResNetArch,
+    Block: type[BasicBlock | BottleneckBlock],
+    depth: int,
+    pretrained: bool,
+    **kwargs: Unpack[_ResNetOptions],
+) -> ResNet:
     model = ResNet(Block, depth, **kwargs)
     if pretrained:
         assert (
@@ -363,7 +418,7 @@ def _resnet(arch, Block, depth, pretrained, **kwargs):
     return model
 
 
-def resnet18(pretrained=False, **kwargs):
+def resnet18(pretrained=False, **kwargs: Unpack[_ResNetOptions]) -> ResNet:
     """ResNet 18-layer model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
 
@@ -396,7 +451,9 @@ def resnet18(pretrained=False, **kwargs):
     return _resnet('resnet18', BasicBlock, 18, pretrained, **kwargs)
 
 
-def resnet34(pretrained=False, **kwargs):
+def resnet34(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNet 34-layer model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
 
@@ -429,7 +486,9 @@ def resnet34(pretrained=False, **kwargs):
     return _resnet('resnet34', BasicBlock, 34, pretrained, **kwargs)
 
 
-def resnet50(pretrained=False, **kwargs):
+def resnet50(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNet 50-layer model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
 
@@ -462,7 +521,9 @@ def resnet50(pretrained=False, **kwargs):
     return _resnet('resnet50', BottleneckBlock, 50, pretrained, **kwargs)
 
 
-def resnet101(pretrained=False, **kwargs):
+def resnet101(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNet 101-layer model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
 
@@ -495,7 +556,9 @@ def resnet101(pretrained=False, **kwargs):
     return _resnet('resnet101', BottleneckBlock, 101, pretrained, **kwargs)
 
 
-def resnet152(pretrained=False, **kwargs):
+def resnet152(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNet 152-layer model from
     `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
 
@@ -528,7 +591,9 @@ def resnet152(pretrained=False, **kwargs):
     return _resnet('resnet152', BottleneckBlock, 152, pretrained, **kwargs)
 
 
-def resnext50_32x4d(pretrained=False, **kwargs):
+def resnext50_32x4d(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNeXt-50 32x4d model from
     `"Aggregated Residual Transformations for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
 
@@ -563,7 +628,9 @@ def resnext50_32x4d(pretrained=False, **kwargs):
     return _resnet('resnext50_32x4d', BottleneckBlock, 50, pretrained, **kwargs)
 
 
-def resnext50_64x4d(pretrained=False, **kwargs):
+def resnext50_64x4d(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNeXt-50 64x4d model from
     `"Aggregated Residual Transformations for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
 
@@ -598,7 +665,9 @@ def resnext50_64x4d(pretrained=False, **kwargs):
     return _resnet('resnext50_64x4d', BottleneckBlock, 50, pretrained, **kwargs)
 
 
-def resnext101_32x4d(pretrained=False, **kwargs):
+def resnext101_32x4d(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNeXt-101 32x4d model from
     `"Aggregated Residual Transformations for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
 
@@ -635,7 +704,9 @@ def resnext101_32x4d(pretrained=False, **kwargs):
     )
 
 
-def resnext101_64x4d(pretrained=False, **kwargs):
+def resnext101_64x4d(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNeXt-101 64x4d model from
     `"Aggregated Residual Transformations for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
 
@@ -672,7 +743,9 @@ def resnext101_64x4d(pretrained=False, **kwargs):
     )
 
 
-def resnext152_32x4d(pretrained=False, **kwargs):
+def resnext152_32x4d(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNeXt-152 32x4d model from
     `"Aggregated Residual Transformations for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
 
@@ -709,7 +782,9 @@ def resnext152_32x4d(pretrained=False, **kwargs):
     )
 
 
-def resnext152_64x4d(pretrained=False, **kwargs):
+def resnext152_64x4d(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """ResNeXt-152 64x4d model from
     `"Aggregated Residual Transformations for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
 
@@ -746,7 +821,9 @@ def resnext152_64x4d(pretrained=False, **kwargs):
     )
 
 
-def wide_resnet50_2(pretrained=False, **kwargs):
+def wide_resnet50_2(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """Wide ResNet-50-2 model from
     `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
 
@@ -780,7 +857,9 @@ def wide_resnet50_2(pretrained=False, **kwargs):
     return _resnet('wide_resnet50_2', BottleneckBlock, 50, pretrained, **kwargs)
 
 
-def wide_resnet101_2(pretrained=False, **kwargs):
+def wide_resnet101_2(
+    pretrained: bool = False, **kwargs: Unpack[_ResNetOptions]
+) -> ResNet:
     """Wide ResNet-101-2 model from
     `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
 
diff --git a/python/paddle/vision/models/squeezenet.py b/python/paddle/vision/models/squeezenet.py
@@ -34,7 +34,6 @@
 
 
 class _SqueezeNetOptions(TypedDict):
-    version: NotRequired[str]
     num_classes: NotRequired[int]
     with_pool: NotRequired[bool]
 
