[Typing][A-80] Add type annotations for vision model MobileNetV3 (#65366)

enkilee · web-flow · commit d07466e3277e · 2024-06-24T19:39:41.000+08:00
diff --git a/python/paddle/vision/models/mobilenetv3.py b/python/paddle/vision/models/mobilenetv3.py
@@ -12,7 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 from functools import partial
+from typing import (
+    TYPE_CHECKING,
+    TypedDict,
+)
+
+from typing_extensions import NotRequired, Unpack
 
 import paddle
 from paddle import nn
@@ -21,6 +29,9 @@
 from ..ops import ConvNormActivation
 from ._utils import _make_divisible
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 __all__ = []
 
 model_urls = {
@@ -35,6 +46,12 @@
 }
 
 
+class _MobileNetV3Options(TypedDict):
+    scale: NotRequired[float]
+    num_classes: NotRequired[int]
+    with_pool: NotRequired[bool]
+
+
 class SqueezeExcitation(nn.Layer):
     """
     This block implements the Squeeze-and-Excitation block from https://arxiv.org/abs/1709.01507 (see Fig. 1).
@@ -51,26 +68,26 @@ class SqueezeExcitation(nn.Layer):
 
     def __init__(
         self,
-        input_channels,
-        squeeze_channels,
-        activation=nn.ReLU,
-        scale_activation=nn.Sigmoid,
-    ):
+        input_channels: int,
+        squeeze_channels: int,
+        activation: nn.Layer = nn.ReLU,
+        scale_activation: nn.Layer = nn.Sigmoid,
+    ) -> None:
         super().__init__()
         self.avgpool = nn.AdaptiveAvgPool2D(1)
         self.fc1 = nn.Conv2D(input_channels, squeeze_channels, 1)
         self.fc2 = nn.Conv2D(squeeze_channels, input_channels, 1)
         self.activation = activation()
         self.scale_activation = scale_activation()
 
-    def _scale(self, input):
+    def _scale(self, input: Tensor) -> Tensor:
         scale = self.avgpool(input)
         scale = self.fc1(scale)
         scale = self.activation(scale)
         scale = self.fc2(scale)
         return self.scale_activation(scale)
 
-    def forward(self, input):
+    def forward(self, input: Tensor) -> Tensor:
         scale = self._scale(input)
         return scale * input
 
@@ -114,15 +131,15 @@ def adjust_channels(channels, scale=1.0):
 class InvertedResidual(nn.Layer):
     def __init__(
         self,
-        in_channels,
-        expanded_channels,
-        out_channels,
-        filter_size,
-        stride,
-        use_se,
-        activation_layer,
-        norm_layer,
-    ):
+        in_channels: int,
+        expanded_channels: int,
+        out_channels: int,
+        filter_size: int,
+        stride: int,
+        use_se: bool,
+        activation_layer: nn.Layer,
+        norm_layer: nn.Layer,
+    ) -> None:
         super().__init__()
         self.use_res_connect = stride == 1 and in_channels == out_channels
         self.use_se = use_se
@@ -167,7 +184,7 @@ def __init__(
             activation_layer=None,
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         identity = x
         if self.expand:
             x = self.expand_conv(x)
@@ -194,8 +211,13 @@ class MobileNetV3(nn.Layer):
     """
 
     def __init__(
-        self, config, last_channel, scale=1.0, num_classes=1000, with_pool=True
-    ):
+        self,
+        config: list[InvertedResidualConfig],
+        last_channel: int,
+        scale: float = 1.0,
+        num_classes: int = 1000,
+        with_pool: bool = True,
+    ) -> None:
         super().__init__()
 
         self.config = config
@@ -257,7 +279,7 @@ def __init__(
                 nn.Linear(self.last_channel, num_classes),
             )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.conv(x)
         x = self.blocks(x)
         x = self.lastconv(x)
@@ -301,7 +323,12 @@ class MobileNetV3Small(MobileNetV3):
             [1, 1000]
     """
 
-    def __init__(self, scale=1.0, num_classes=1000, with_pool=True):
+    def __init__(
+        self,
+        scale: float = 1.0,
+        num_classes: int = 1000,
+        with_pool: bool = True,
+    ) -> None:
         config = [
             InvertedResidualConfig(16, 3, 16, 16, True, "relu", 2, scale),
             InvertedResidualConfig(16, 3, 72, 24, False, "relu", 2, scale),
@@ -354,7 +381,12 @@ class MobileNetV3Large(MobileNetV3):
             [1, 1000]
     """
 
-    def __init__(self, scale=1.0, num_classes=1000, with_pool=True):
+    def __init__(
+        self,
+        scale: float = 1.0,
+        num_classes: int = 1000,
+        with_pool: bool = True,
+    ) -> None:
         config = [
             InvertedResidualConfig(16, 3, 16, 16, False, "relu", 1, scale),
             InvertedResidualConfig(16, 3, 64, 24, False, "relu", 2, scale),
@@ -400,7 +432,12 @@ def __init__(self, scale=1.0, num_classes=1000, with_pool=True):
         )
 
 
-def _mobilenet_v3(arch, pretrained=False, scale=1.0, **kwargs):
+def _mobilenet_v3(
+    arch: str,
+    pretrained: bool = False,
+    scale: float = 1.0,
+    **kwargs: Unpack[_MobileNetV3Options],
+) -> MobileNetV3:
     if arch == "mobilenet_v3_large":
         model = MobileNetV3Large(scale=scale, **kwargs)
     else:
@@ -419,7 +456,11 @@ def _mobilenet_v3(arch, pretrained=False, scale=1.0, **kwargs):
     return model
 
 
-def mobilenet_v3_small(pretrained=False, scale=1.0, **kwargs):
+def mobilenet_v3_small(
+    pretrained: bool = False,
+    scale: float = 1.0,
+    **kwargs: Unpack[_MobileNetV3Options],
+) -> MobileNetV3Small:
     """MobileNetV3 Small architecture model from
     `"Searching for MobileNetV3" <https://arxiv.org/abs/1905.02244>`_.
 
@@ -458,7 +499,11 @@ def mobilenet_v3_small(pretrained=False, scale=1.0, **kwargs):
     return model
 
 
-def mobilenet_v3_large(pretrained=False, scale=1.0, **kwargs):
+def mobilenet_v3_large(
+    pretrained: bool = False,
+    scale: float = 1.0,
+    **kwargs: Unpack[_MobileNetV3Options],
+) -> MobileNetV3Large:
     """MobileNetV3 Large architecture model from
     `"Searching for MobileNetV3" <https://arxiv.org/abs/1905.02244>`_.
 
