[Typing][A-75] Add type annotations for python/paddle/vision/models/googlenet.py (#65290)

DrRyanHuang · SigureMo · web-flow · commit cf87663f6a83 · 2024-06-20T03:29:27.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/vision/models/googlenet.py b/python/paddle/vision/models/googlenet.py
@@ -12,9 +12,19 @@
 # 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
 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,
@@ -27,6 +37,8 @@
 from paddle.nn.initializer import Uniform
 from paddle.utils.download import get_weights_path_from_url
 
+if TYPE_CHECKING:
+    from paddle import Tensor
 __all__ = []
 
 model_urls = {
@@ -37,15 +49,20 @@
 }
 
 
-def xavier(channels, filter_size):
+def xavier(channels: int, filter_size: int) -> ParamAttr:
     stdv = (3.0 / (filter_size**2 * channels)) ** 0.5
     param_attr = ParamAttr(initializer=Uniform(-stdv, stdv))
     return param_attr
 
 
 class ConvLayer(nn.Layer):
     def __init__(
-        self, num_channels, num_filters, filter_size, stride=1, groups=1
+        self,
+        num_channels: int,
+        num_filters: int,
+        filter_size: int,
+        stride: Size2 = 1,
+        groups: int = 1,
     ):
         super().__init__()
 
@@ -59,22 +76,22 @@ def __init__(
             bias_attr=False,
         )
 
-    def forward(self, inputs):
+    def forward(self, inputs: Tensor) -> Tensor:
         y = self._conv(inputs)
         return y
 
 
 class Inception(nn.Layer):
     def __init__(
         self,
-        input_channels,
-        output_channels,
-        filter1,
-        filter3R,
-        filter3,
-        filter5R,
-        filter5,
-        proj,
+        input_channels: int,
+        output_channels: int,
+        filter1: int,
+        filter3R: int,
+        filter3: int,
+        filter5R: int,
+        filter5: int,
+        proj: int,
     ):
         super().__init__()
 
@@ -87,7 +104,7 @@ def __init__(
 
         self._convprj = ConvLayer(input_channels, proj, 1)
 
-    def forward(self, inputs):
+    def forward(self, inputs: Tensor) -> Tensor:
         conv1 = self._conv1(inputs)
 
         conv3r = self._conv3r(inputs)
@@ -132,7 +149,7 @@ class GoogLeNet(nn.Layer):
             [1, 1000] [1, 1000] [1, 1000]
     """
 
-    def __init__(self, num_classes=1000, with_pool=True):
+    def __init__(self, num_classes: int = 1000, with_pool: bool = True) -> None:
         super().__init__()
         self.num_classes = num_classes
         self.with_pool = with_pool
@@ -181,7 +198,7 @@ def __init__(self, num_classes=1000, with_pool=True):
             self._drop_o2 = Dropout(p=0.7, mode="downscale_in_infer")
             self._out2 = Linear(1024, num_classes, weight_attr=xavier(1024, 1))
 
-    def forward(self, inputs):
+    def forward(self, inputs: Tensor) -> tuple[Tensor, Tensor, Tensor]:
         x = self._conv(inputs)
         x = self._pool(x)
         x = self._conv_1(x)
@@ -227,10 +244,17 @@ def forward(self, inputs):
             out2 = self._drop_o2(out2)
             out2 = self._out2(out2)
 
-        return [out, out1, out2]
+        return out, out1, out2
+
+
+class _GoogLeNetOptions(TypedDict):
+    num_classes: NotRequired[int]
+    with_pool: NotRequired[bool]
 
 
-def googlenet(pretrained=False, **kwargs):
+def googlenet(
+    pretrained: bool = False, **kwargs: Unpack[_GoogLeNetOptions]
+) -> GoogLeNet:
     """GoogLeNet (Inception v1) model architecture from
     `"Going Deeper with Convolutions" <https://arxiv.org/pdf/1409.4842.pdf>`_.
 
