[Typing][A-76] Add type annotations for python/paddle/vision/models/inceptionv3.py (#65292)

DrRyanHuang · web-flow · commit 5d634e4c76aa · 2024-06-20T13:09:05.000+08:00
diff --git a/python/paddle/vision/models/inceptionv3.py b/python/paddle/vision/models/inceptionv3.py
@@ -12,7 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import math
+from typing import (
+    TYPE_CHECKING,
+    TypedDict,
+)
+
+from typing_extensions import (
+    NotRequired,
+    Unpack,
+)
 
 import paddle
 from paddle import nn
@@ -23,6 +34,9 @@
 
 from ..ops import ConvNormActivation
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 __all__ = []
 
 model_urls = {
@@ -33,8 +47,13 @@
 }
 
 
+class _InceptionV3Options(TypedDict):
+    num_classes: NotRequired[int]
+    with_pool: NotRequired[bool]
+
+
 class InceptionStem(nn.Layer):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__()
         self.conv_1a_3x3 = ConvNormActivation(
             in_channels=3,
@@ -76,7 +95,7 @@ def __init__(self):
             activation_layer=nn.ReLU,
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.conv_1a_3x3(x)
         x = self.conv_2a_3x3(x)
         x = self.conv_2b_3x3(x)
@@ -88,7 +107,7 @@ def forward(self, x):
 
 
 class InceptionA(nn.Layer):
-    def __init__(self, num_channels, pool_features):
+    def __init__(self, num_channels: int, pool_features: int) -> None:
         super().__init__()
         self.branch1x1 = ConvNormActivation(
             in_channels=num_channels,
@@ -146,7 +165,7 @@ def __init__(self, num_channels, pool_features):
             activation_layer=nn.ReLU,
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         branch1x1 = self.branch1x1(x)
         branch5x5 = self.branch5x5_1(x)
         branch5x5 = self.branch5x5_2(branch5x5)
@@ -164,7 +183,7 @@ def forward(self, x):
 
 
 class InceptionB(nn.Layer):
-    def __init__(self, num_channels):
+    def __init__(self, num_channels: int) -> None:
         super().__init__()
         self.branch3x3 = ConvNormActivation(
             in_channels=num_channels,
@@ -200,7 +219,7 @@ def __init__(self, num_channels):
 
         self.branch_pool = MaxPool2D(kernel_size=3, stride=2)
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         branch3x3 = self.branch3x3(x)
 
         branch3x3dbl = self.branch3x3dbl_1(x)
@@ -215,7 +234,7 @@ def forward(self, x):
 
 
 class InceptionC(nn.Layer):
-    def __init__(self, num_channels, channels_7x7):
+    def __init__(self, num_channels: int, channels_7x7: int) -> None:
         super().__init__()
         self.branch1x1 = ConvNormActivation(
             in_channels=num_channels,
@@ -297,7 +316,7 @@ def __init__(self, num_channels, channels_7x7):
             activation_layer=nn.ReLU,
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         branch1x1 = self.branch1x1(x)
 
         branch7x7 = self.branch7x7_1(x)
@@ -321,7 +340,7 @@ def forward(self, x):
 
 
 class InceptionD(nn.Layer):
-    def __init__(self, num_channels):
+    def __init__(self, num_channels: int) -> None:
         super().__init__()
         self.branch3x3_1 = ConvNormActivation(
             in_channels=num_channels,
@@ -371,7 +390,7 @@ def __init__(self, num_channels):
 
         self.branch_pool = MaxPool2D(kernel_size=3, stride=2)
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         branch3x3 = self.branch3x3_1(x)
         branch3x3 = self.branch3x3_2(branch3x3)
 
@@ -387,7 +406,7 @@ def forward(self, x):
 
 
 class InceptionE(nn.Layer):
-    def __init__(self, num_channels):
+    def __init__(self, num_channels: int) -> None:
         super().__init__()
         self.branch1x1 = ConvNormActivation(
             in_channels=num_channels,
@@ -458,7 +477,7 @@ def __init__(self, num_channels):
             activation_layer=nn.ReLU,
         )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         branch1x1 = self.branch1x1(x)
 
         branch3x3 = self.branch3x3_1(x)
@@ -512,7 +531,7 @@ class InceptionV3(nn.Layer):
             [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
@@ -570,7 +589,7 @@ def __init__(self, num_classes=1000, with_pool=True):
                 bias_attr=ParamAttr(),
             )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.inception_stem(x)
         for inception_block in self.inception_block_list:
             x = inception_block(x)
@@ -585,7 +604,9 @@ def forward(self, x):
         return x
 
 
-def inception_v3(pretrained=False, **kwargs):
+def inception_v3(
+    pretrained: bool = False, **kwargs: Unpack[_InceptionV3Options]
+) -> InceptionV3:
     """Inception v3 model from
     `"Rethinking the Inception Architecture for Computer Vision" <https://arxiv.org/pdf/1512.00567.pdf>`_.
 
