[Typing][A-85] Add type annotations for python/paddle/vision/models/vgg.py (#65381)

86kkd · SigureMo · web-flow · commit e78c557f7681 · 2024-06-23T16:17:19.000+08:00
---------

Co-authored-by: Nyakku Shigure &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/vision/models/vgg.py b/python/paddle/vision/models/vgg.py
@@ -12,10 +12,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    TypedDict,
+)
+
+from typing_extensions import NotRequired, Unpack
+
 import paddle
 from paddle import nn
 from paddle.utils.download import get_weights_path_from_url
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.nn import Layer, Sequential
+
+
+class _VGGOptions(TypedDict):
+    num_classes: NotRequired[int]
+    with_pool: NotRequired[bool]
+
+
 __all__ = []
 
 model_urls = {
@@ -63,7 +84,9 @@ class VGG(nn.Layer):
             [1, 1000]
     """
 
-    def __init__(self, features, num_classes=1000, with_pool=True):
+    def __init__(
+        self, features: Layer, num_classes: int = 1000, with_pool: bool = True
+    ) -> None:
         super().__init__()
         self.features = features
         self.num_classes = num_classes
@@ -83,7 +106,7 @@ def __init__(self, features, num_classes=1000, with_pool=True):
                 nn.Linear(4096, num_classes),
             )
 
-    def forward(self, x):
+    def forward(self, x: Tensor) -> Tensor:
         x = self.features(x)
 
         if self.with_pool:
@@ -96,7 +119,9 @@ def forward(self, x):
         return x
 
 
-def make_layers(cfg, batch_norm=False):
+def make_layers(
+    cfg: dict[str, list[Any]], batch_norm: bool = False
+) -> Sequential:
     layers = []
     in_channels = 3
     for v in cfg:
@@ -177,7 +202,13 @@ def make_layers(cfg, batch_norm=False):
 }
 
 
-def _vgg(arch, cfg, batch_norm, pretrained, **kwargs):
+def _vgg(
+    arch: str,
+    cfg: Literal["A", "B", "D", "E"],
+    batch_norm: bool,
+    pretrained: bool,
+    **kwargs: Unpack[_VGGOptions],
+) -> VGG:
     model = VGG(make_layers(cfgs[cfg], batch_norm=batch_norm), **kwargs)
 
     if pretrained:
@@ -194,7 +225,11 @@ def _vgg(arch, cfg, batch_norm, pretrained, **kwargs):
     return model
 
 
-def vgg11(pretrained=False, batch_norm=False, **kwargs):
+def vgg11(
+    pretrained: bool = False,
+    batch_norm: bool = False,
+    **kwargs: Unpack[_VGGOptions],
+) -> VGG:
     """VGG 11-layer model from
     `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_.
 
@@ -231,7 +266,11 @@ def vgg11(pretrained=False, batch_norm=False, **kwargs):
     return _vgg(model_name, 'A', batch_norm, pretrained, **kwargs)
 
 
-def vgg13(pretrained=False, batch_norm=False, **kwargs):
+def vgg13(
+    pretrained: bool = False,
+    batch_norm: bool = False,
+    **kwargs: Unpack[_VGGOptions],
+) -> VGG:
     """VGG 13-layer model from
     `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_.
 
@@ -268,7 +307,11 @@ def vgg13(pretrained=False, batch_norm=False, **kwargs):
     return _vgg(model_name, 'B', batch_norm, pretrained, **kwargs)
 
 
-def vgg16(pretrained=False, batch_norm=False, **kwargs):
+def vgg16(
+    pretrained: bool = False,
+    batch_norm: bool = False,
+    **kwargs: Unpack[_VGGOptions],
+) -> VGG:
     """VGG 16-layer model from
     `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_.
 
@@ -305,7 +348,11 @@ def vgg16(pretrained=False, batch_norm=False, **kwargs):
     return _vgg(model_name, 'D', batch_norm, pretrained, **kwargs)
 
 
-def vgg19(pretrained=False, batch_norm=False, **kwargs):
+def vgg19(
+    pretrained: bool = False,
+    batch_norm: bool = False,
+    **kwargs: Unpack[_VGGOptions],
+) -> VGG:
     """VGG 19-layer model from
     `"Very Deep Convolutional Networks For Large-Scale Image Recognition" <https://arxiv.org/pdf/1409.1556.pdf>`_.
 
