[Typing][A-82] Add type annotations for paddle/vision/models/shufflenetv2.py (#65559)

Asthestarsfalll · SigureMo · web-flow · commit ca9800173f8e · 2024-06-29T12:34:37.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/vision/models/alexnet.py b/python/paddle/vision/models/alexnet.py
@@ -25,7 +25,6 @@
 import paddle
 import paddle.nn.functional as F
 from paddle import nn
-from paddle._typing import Size2
 from paddle.base.param_attr import ParamAttr
 from paddle.nn import Conv2D, Dropout, Linear, MaxPool2D, ReLU
 from paddle.nn.initializer import Uniform
@@ -42,10 +41,10 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
+    from paddle._typing import Size2
 
-
-class _AlexNetOptions(TypedDict):
-    num_classes: NotRequired[int]
+    class _AlexNetOptions(TypedDict):
+        num_classes: NotRequired[int]
 
 
 class ConvPoolLayer(nn.Layer):
@@ -109,6 +108,8 @@ class AlexNet(nn.Layer):
             [1, 1000]
     """
 
+    num_classes: int
+
     def __init__(self, num_classes: int = 1000) -> None:
         super().__init__()
         self.num_classes = num_classes
diff --git a/python/paddle/vision/models/densenet.py b/python/paddle/vision/models/densenet.py
@@ -20,7 +20,6 @@
 
 import paddle
 from paddle import nn
-from paddle._typing import Size2
 from paddle.base.param_attr import ParamAttr
 from paddle.nn import (
     AdaptiveAvgPool2D,
@@ -40,6 +39,7 @@
     from typing_extensions import NotRequired
 
     from paddle import Tensor
+    from paddle._typing import Size2
 
     _DenseNetArch = Literal[
         "densenet121",
@@ -270,6 +270,9 @@ class DenseNet(nn.Layer):
             [1, 1000]
     """
 
+    num_classes: int
+    with_pool: bool
+
     def __init__(
         self,
         layers: int = 121,
diff --git a/python/paddle/vision/models/googlenet.py b/python/paddle/vision/models/googlenet.py
@@ -24,7 +24,6 @@
 import paddle
 import paddle.nn.functional as F
 from paddle import nn
-from paddle._typing import Size2
 from paddle.base.param_attr import ParamAttr
 from paddle.nn import (
     AdaptiveAvgPool2D,
@@ -39,6 +38,13 @@
 
 if TYPE_CHECKING:
     from paddle import Tensor
+    from paddle._typing import Size2
+
+    class _GoogLeNetOptions(TypedDict):
+        num_classes: NotRequired[int]
+        with_pool: NotRequired[bool]
+
+
 __all__ = []
 
 model_urls = {
@@ -149,6 +155,9 @@ class GoogLeNet(nn.Layer):
             [1, 1000] [1, 1000] [1, 1000]
     """
 
+    num_classes: int
+    with_pool: bool
+
     def __init__(self, num_classes: int = 1000, with_pool: bool = True) -> None:
         super().__init__()
         self.num_classes = num_classes
@@ -247,11 +256,6 @@ def forward(self, inputs: Tensor) -> tuple[Tensor, Tensor, Tensor]:
         return out, out1, out2
 
 
-class _GoogLeNetOptions(TypedDict):
-    num_classes: NotRequired[int]
-    with_pool: NotRequired[bool]
-
-
 def googlenet(
     pretrained: bool = False, **kwargs: Unpack[_GoogLeNetOptions]
 ) -> GoogLeNet:
diff --git a/python/paddle/vision/models/inceptionv3.py b/python/paddle/vision/models/inceptionv3.py
@@ -37,6 +37,11 @@
 if TYPE_CHECKING:
     from paddle import Tensor
 
+    class _InceptionV3Options(TypedDict):
+        num_classes: NotRequired[int]
+        with_pool: NotRequired[bool]
+
+
 __all__ = []
 
 model_urls = {
@@ -47,11 +52,6 @@
 }
 
 
-class _InceptionV3Options(TypedDict):
-    num_classes: NotRequired[int]
-    with_pool: NotRequired[bool]
-
-
 class InceptionStem(nn.Layer):
     def __init__(self) -> None:
         super().__init__()
@@ -531,6 +531,9 @@ class InceptionV3(nn.Layer):
             [1, 1000]
     """
 
+    num_classes: int
+    with_pool: bool
+
     def __init__(self, num_classes: int = 1000, with_pool: bool = True) -> None:
         super().__init__()
         self.num_classes = num_classes
diff --git a/python/paddle/vision/models/lenet.py b/python/paddle/vision/models/lenet.py
@@ -53,6 +53,8 @@ class LeNet(nn.Layer):
             [1, 10]
     """
 
+    num_classes: int
+
     def __init__(self, num_classes: int = 10) -> None:
         super().__init__()
         self.num_classes = num_classes
diff --git a/python/paddle/vision/models/mobilenetv1.py b/python/paddle/vision/models/mobilenetv1.py
@@ -23,20 +23,20 @@
 
 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
+    from paddle._typing import Size2
 
-__all__ = []
+    class _MobileNetV1Options(TypedDict):
+        num_classes: NotRequired[int]
+        with_pool: NotRequired[bool]
 
 
-class _MobileNetV1Options(TypedDict):
-    num_classes: NotRequired[int]
-    with_pool: NotRequired[bool]
+__all__ = []
 
 
 model_urls = {
@@ -110,6 +110,10 @@ class MobileNetV1(nn.Layer):
             [1, 1000]
     """
 
+    scale: float
+    num_classes: int
+    with_pool: bool
+
     def __init__(
         self,
         scale: float = 1.0,
diff --git a/python/paddle/vision/models/mobilenetv2.py b/python/paddle/vision/models/mobilenetv2.py
@@ -31,6 +31,11 @@
 if TYPE_CHECKING:
     from paddle import Tensor
 
+    class _MobileNetV2Options(TypedDict):
+        num_classes: NotRequired[int]
+        with_pool: NotRequired[bool]
+
+
 __all__ = []
 
 model_urls = {
@@ -41,11 +46,6 @@
 }
 
 
-class _MobileNetV2Options(TypedDict):
-    num_classes: NotRequired[int]
-    with_pool: NotRequired[bool]
-
-
 class InvertedResidual(nn.Layer):
     def __init__(
         self,
@@ -124,6 +124,9 @@ class MobileNetV2(nn.Layer):
             [1, 1000]
     """
 
+    num_classes: int
+    with_pool: bool
+
     def __init__(
         self,
         scale: float = 1.0,
diff --git a/python/paddle/vision/models/mobilenetv3.py b/python/paddle/vision/models/mobilenetv3.py
@@ -32,6 +32,11 @@
 if TYPE_CHECKING:
     from paddle import Tensor
 
+    class _MobileNetV3Options(TypedDict):
+        num_classes: NotRequired[int]
+        with_pool: NotRequired[bool]
+
+
 __all__ = []
 
 model_urls = {
@@ -46,11 +51,6 @@
 }
 
 
-class _MobileNetV3Options(TypedDict):
-    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).
@@ -209,6 +209,10 @@ class MobileNetV3(nn.Layer):
         with_pool (bool, optional): Use pool before the last fc layer or not. Default: True.
     """
 
+    scale: float
+    num_classes: int
+    with_pool: bool
+
     def __init__(
         self,
         config: list[InvertedResidualConfig],
diff --git a/python/paddle/vision/models/resnet.py b/python/paddle/vision/models/resnet.py
@@ -17,7 +17,6 @@
 
 import paddle
 from paddle import nn
-from paddle._typing import Size2
 from paddle.utils.download import get_weights_path_from_url
 
 if TYPE_CHECKING:
@@ -26,6 +25,7 @@
     from typing_extensions import NotRequired, Unpack
 
     from paddle import Tensor
+    from paddle._typing import Size2
 
     _ResNetArch = Literal[
         'resnet18',
@@ -177,7 +177,7 @@ def __init__(
     ) -> None:
         super().__init__()
         if norm_layer is None:
-            norm_layer: type[nn.BatchNorm2D] = nn.BatchNorm2D
+            norm_layer = nn.BatchNorm2D
         width = int(planes * (base_width / 64.0)) * groups
 
         self.conv1 = nn.Conv2D(inplanes, width, 1, bias_attr=False)
diff --git a/python/paddle/vision/models/shufflenetv2.py b/python/paddle/vision/models/shufflenetv2.py
diff --git a/python/paddle/vision/models/squeezenet.py b/python/paddle/vision/models/squeezenet.py
diff --git a/python/paddle/vision/models/vgg.py b/python/paddle/vision/models/vgg.py
