[Typing][A-74] Add type annotations for paddle/vision/models/densenet.py (#65486)

Asthestarsfalll · SigureMo · web-flow · commit 841e254d9665 · 2024-06-28T08:58:33.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/vision/models/densenet.py b/python/paddle/vision/models/densenet.py
@@ -11,11 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # 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
+
+from typing_extensions import Unpack
 
 import paddle
 from paddle import nn
+from paddle._typing import Size2
 from paddle.base.param_attr import ParamAttr
 from paddle.nn import (
     AdaptiveAvgPool2D,
@@ -29,9 +34,31 @@
 from paddle.nn.initializer import Uniform
 from paddle.utils.download import get_weights_path_from_url
 
+if TYPE_CHECKING:
+    from typing import Literal, TypedDict
+
+    from typing_extensions import NotRequired
+
+    from paddle import Tensor
+
+    _DenseNetArch = Literal[
+        "densenet121",
+        "densenet161",
+        "densenet169",
+        "densenet201",
+        "densenet264",
+    ]
+
+    class _DenseNetOptions(TypedDict):
+        bn_size: NotRequired[int]
+        dropout: NotRequired[float]
+        num_classes: NotRequired[int]
+        with_pool: NotRequired[bool]
+
+
 __all__ = []
 
-model_urls = {
+model_urls: dict[str, tuple[str, str]] = {
     'densenet121': (
         'https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/DenseNet121_pretrained.pdparams',
         'db1b239ed80a905290fd8b01d3af08e4',
@@ -58,14 +85,14 @@
 class BNACConvLayer(nn.Layer):
     def __init__(
         self,
-        num_channels,
-        num_filters,
-        filter_size,
-        stride=1,
-        pad=0,
-        groups=1,
-        act="relu",
-    ):
+        num_channels: int,
+        num_filters: int,
+        filter_size: Size2,
+        stride: Size2 = 1,
+        pad: Size2 = 0,
+        groups: int = 1,
+        act: str = "relu",
+    ) -> None:
         super().__init__()
         self._batch_norm = BatchNorm(num_channels, act=act)
 
@@ -80,14 +107,18 @@ def __init__(
             bias_attr=False,
         )
 
-    def forward(self, input):
+    def forward(self, input: Tensor) -> Tensor:
         y = self._batch_norm(input)
         y = self._conv(y)
         return y
 
 
 class DenseLayer(nn.Layer):
-    def __init__(self, num_channels, growth_rate, bn_size, dropout):
+    dropout: float
+
+    def __init__(
+        self, num_channels: int, growth_rate: int, bn_size: int, dropout: float
+    ) -> None:
         super().__init__()
         self.dropout = dropout
 
@@ -110,7 +141,7 @@ def __init__(self, num_channels, growth_rate, bn_size, dropout):
         if dropout:
             self.dropout_func = Dropout(p=dropout, mode="downscale_in_infer")
 
-    def forward(self, input):
+    def forward(self, input: Tensor) -> Tensor:
         conv = self.bn_ac_func1(input)
         conv = self.bn_ac_func2(conv)
         if self.dropout:
@@ -120,9 +151,17 @@ def forward(self, input):
 
 
 class DenseBlock(nn.Layer):
+    dropout: float
+
     def __init__(
-        self, num_channels, num_layers, bn_size, growth_rate, dropout, name=None
-    ):
+        self,
+        num_channels: int,
+        num_layers: int,
+        bn_size: int,
+        growth_rate: int,
+        dropout: float,
+        name: str | None = None,
+    ) -> None:
         super().__init__()
         self.dropout = dropout
         self.dense_layer_func = []
@@ -142,15 +181,15 @@ def __init__(
             )
             pre_channel = pre_channel + growth_rate
 
-    def forward(self, input):
+    def forward(self, input: Tensor) -> Tensor:
         conv = input
         for func in self.dense_layer_func:
             conv = func(conv)
         return conv
 
 
 class TransitionLayer(nn.Layer):
-    def __init__(self, num_channels, num_output_features):
+    def __init__(self, num_channels: int, num_output_features: int) -> None:
         super().__init__()
 
         self.conv_ac_func = BNACConvLayer(
@@ -163,7 +202,7 @@ def __init__(self, num_channels, num_output_features):
 
         self.pool2d_avg = AvgPool2D(kernel_size=2, stride=2, padding=0)
 
-    def forward(self, input):
+    def forward(self, input: Tensor) -> Tensor:
         y = self.conv_ac_func(input)
         y = self.pool2d_avg(y)
         return y
@@ -172,14 +211,14 @@ def forward(self, input):
 class ConvBNLayer(nn.Layer):
     def __init__(
         self,
-        num_channels,
-        num_filters,
-        filter_size,
-        stride=1,
-        pad=0,
-        groups=1,
-        act="relu",
-    ):
+        num_channels: int,
+        num_filters: int,
+        filter_size: Size2,
+        stride: Size2 = 1,
+        pad: Size2 = 0,
+        groups: int = 1,
+        act: str = "relu",
+    ) -> None:
         super().__init__()
 
         self._conv = Conv2D(
@@ -194,7 +233,7 @@ def __init__(
         )
         self._batch_norm = BatchNorm(num_filters, act=act)
 
-    def forward(self, input):
+    def forward(self, input: Tensor) -> Tensor:
         y = self._conv(input)
         y = self._batch_norm(y)
         return y
@@ -233,12 +272,12 @@ class DenseNet(nn.Layer):
 
     def __init__(
         self,
-        layers=121,
-        bn_size=4,
-        dropout=0.0,
-        num_classes=1000,
-        with_pool=True,
-    ):
+        layers: int = 121,
+        bn_size: int = 4,
+        dropout: float = 0.0,
+        num_classes: int = 1000,
+        with_pool: bool = True,
+    ) -> None:
         super().__init__()
         self.num_classes = num_classes
         self.with_pool = with_pool
@@ -313,7 +352,7 @@ def __init__(
                 bias_attr=ParamAttr(),
             )
 
-    def forward(self, input):
+    def forward(self, input: Tensor) -> Tensor:
         conv = self.conv1_func(input)
         conv = self.pool2d_max(conv)
 
@@ -334,7 +373,12 @@ def forward(self, input):
         return y
 
 
-def _densenet(arch, layers, pretrained, **kwargs):
+def _densenet(
+    arch: _DenseNetArch,
+    layers: int,
+    pretrained: bool,
+    **kwargs: Unpack[_DenseNetOptions],
+) -> DenseNet:
     model = DenseNet(layers=layers, **kwargs)
     if pretrained:
         assert (
@@ -350,7 +394,9 @@ def _densenet(arch, layers, pretrained, **kwargs):
     return model
 
 
-def densenet121(pretrained=False, **kwargs):
+def densenet121(
+    pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
+) -> DenseNet:
     """DenseNet 121-layer model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
 
@@ -383,7 +429,9 @@ def densenet121(pretrained=False, **kwargs):
     return _densenet('densenet121', 121, pretrained, **kwargs)
 
 
-def densenet161(pretrained=False, **kwargs):
+def densenet161(
+    pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
+) -> DenseNet:
     """DenseNet 161-layer model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
 
@@ -416,7 +464,9 @@ def densenet161(pretrained=False, **kwargs):
     return _densenet('densenet161', 161, pretrained, **kwargs)
 
 
-def densenet169(pretrained=False, **kwargs):
+def densenet169(
+    pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
+) -> DenseNet:
     """DenseNet 169-layer model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
 
@@ -449,7 +499,9 @@ def densenet169(pretrained=False, **kwargs):
     return _densenet('densenet169', 169, pretrained, **kwargs)
 
 
-def densenet201(pretrained=False, **kwargs):
+def densenet201(
+    pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
+) -> DenseNet:
     """DenseNet 201-layer model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
 
@@ -481,7 +533,9 @@ def densenet201(pretrained=False, **kwargs):
     return _densenet('densenet201', 201, pretrained, **kwargs)
 
 
-def densenet264(pretrained=False, **kwargs):
+def densenet264(
+    pretrained: bool = False, **kwargs: Unpack[_DenseNetOptions]
+) -> DenseNet:
     """DenseNet 264-layer model from
     `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.
 
