[Typing][A-78] Add type annotations for python/paddle/vision/models/mobilenetv1.py (#65323)

DrRyanHuang · web-flow · commit 9456a126b546 · 2024-06-21T10:00:28.000+08:00
diff --git a/python/paddle/vision/models/mobilenetv1.py b/python/paddle/vision/models/mobilenetv1.py
@@ -12,14 +12,34 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    TypedDict,
+)
+
+from typing_extensions import NotRequired, Unpack
+
 import paddle
 from paddle import nn
+from paddle._typing import Size2
 from paddle.utils.download import get_weights_path_from_url
 
 from ..ops import ConvNormActivation
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 __all__ = []
 
+
+class _MobileNetV1Options(TypedDict):
+    scale: NotRequired[float]
+    num_classes: NotRequired[int]
+    with_pool: NotRequired[bool]
+
+
 model_urls = {
     'mobilenetv1_1.0': (
         'https://paddle-hapi.bj.bcebos.com/models/mobilenetv1_1.0.pdparams',
@@ -31,13 +51,13 @@
 class DepthwiseSeparable(nn.Layer):
     def __init__(
         self,
-        in_channels,
-        out_channels1,
-        out_channels2,
-        num_groups,
-        stride,
-        scale,
-    ):
+        in_channels: int,
+        out_channels1: int,
+        out_channels2: int,
+        num_groups: int,
+        stride: Size2,
+        scale: float,
+    ) -> None:
         super().__init__()
 
         self._depthwise_conv = ConvNormActivation(
@@ -57,7 +77,7 @@ def __init__(
             padding=0,
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         x = self._depthwise_conv(x)
         x = self._pointwise_conv(x)
         return x
@@ -91,7 +111,12 @@ class MobileNetV1(nn.Layer):
             [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:
         super().__init__()
         self.scale = scale
         self.dwsl = []
@@ -230,7 +255,7 @@ def __init__(self, scale=1.0, num_classes=1000, with_pool=True):
         if num_classes > 0:
             self.fc = nn.Linear(int(1024 * scale), num_classes)
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.conv1(x)
         for dws in self.dwsl:
             x = dws(x)
@@ -244,7 +269,9 @@ def forward(self, x):
         return x
 
 
-def _mobilenet(arch, pretrained=False, **kwargs):
+def _mobilenet(
+    arch: str, pretrained: bool = False, **kwargs: Unpack[_MobileNetV1Options]
+) -> MobileNetV1:
     model = MobileNetV1(**kwargs)
     if pretrained:
         assert (
@@ -260,7 +287,11 @@ def _mobilenet(arch, pretrained=False, **kwargs):
     return model
 
 
-def mobilenet_v1(pretrained=False, scale=1.0, **kwargs):
+def mobilenet_v1(
+    pretrained: bool = False,
+    scale: float = 1.0,
+    **kwargs: Unpack[_MobileNetV1Options],
+) -> MobileNetV1:
     """MobileNetV1 from
     `"MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications" <https://arxiv.org/abs/1704.04861>`_.
 
