Upd typing, 2

pytorch/metrics/det_metrics.py

@@ -19,7 +19,7 @@ class CocoDetMApMetric(mx.metric.EvalMetric):
         Processed image height.
     coco_annotations_file_path : str
         COCO anotation file path.
-    contiguous_id_to_json : list of int
+    contiguous_id_to_json : list(int)
         Processed IDs.
     validation_ids : bool, default False
         Whether to use temporary file for estimation.

pytorch/pytorchcv/models/bagnet.py

@@ -29,11 +29,11 @@ class BagNetBottleneck(nn.Module):
         Bottleneck factor.
     """
     def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride,
-                 bottleneck_factor=4):
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: int | tuple[int, int],
+                 stride: int | tuple[int, int],
+                 bottleneck_factor: int = 4):
         super(BagNetBottleneck, self).__init__()
         mid_channels = out_channels // bottleneck_factor
 
@@ -74,10 +74,10 @@ class BagNetUnit(nn.Module):
         Strides of the second body convolution.
     """
     def __init__(self,
-                 in_channels,
-                 out_channels,
-                 kernel_size,
-                 stride):
+                 in_channels: int,
+                 out_channels: int,
+                 kernel_size: int | tuple[int, int],
+                 stride: int | tuple[int, int]):
         super(BagNetUnit, self).__init__()
         self.resize_identity = (in_channels != out_channels) or (stride != 1)
 
@@ -120,8 +120,8 @@ class BagNetInitBlock(nn.Module):
         Number of output channels.
     """
     def __init__(self,
-                 in_channels,
-                 out_channels):
+                 in_channels: int,
+                 out_channels: int):
         super(BagNetInitBlock, self).__init__()
         self.conv1 = conv1x1(
             in_channels=in_channels,
@@ -150,7 +150,7 @@ class BagNet(nn.Module):
         Number of output channels for the initial unit.
     final_pool_size : int
         Size of the pooling windows for final pool.
-    normal_kernel_sizes : list of int
+    normal_kernel_sizes : list(int)
         Count of the first units with 3x3 convolution window size for each stage.
     in_channels : int, default 3
         Number of input channels.
@@ -160,10 +160,10 @@ class BagNet(nn.Module):
         Number of classification classes.
     """
     def __init__(self,
-                 channels,
-                 init_block_channels,
-                 final_pool_size,
-                 normal_kernel_sizes,
+                 channels: list[list[int]],
+                 init_block_channels: int,
+                 final_pool_size: int,
+                 normal_kernel_sizes: list[int],
                  in_channels: int = 3,
                  in_size: tuple[int, int] = (224, 224),
                  num_classes: int = 1000):
@@ -212,7 +212,7 @@ def forward(self, x):
         return x
 
 
-def get_bagnet(field,
+def get_bagnet(field: int,
                model_name: str | None = None,
                pretrained: bool = False,
                root: str = os.path.join("~", ".torch", "models"),
@@ -236,7 +236,6 @@ def get_bagnet(field,
     nn.Module
         Desired module.
     """
-
     layers = [3, 4, 6, 3]
 
     if field == 9:
@@ -292,7 +291,10 @@ def bagnet9(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_bagnet(field=9, model_name="bagnet9", **kwargs)
+    return get_bagnet(
+        field=9,
+        model_name="bagnet9",
+        **kwargs)
 
 
 def bagnet17(**kwargs) -> nn.Module:
@@ -312,7 +314,10 @@ def bagnet17(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_bagnet(field=17, model_name="bagnet17", **kwargs)
+    return get_bagnet(
+        field=17,
+        model_name="bagnet17",
+        **kwargs)
 
 
 def bagnet33(**kwargs) -> nn.Module:
@@ -332,7 +337,10 @@ def bagnet33(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_bagnet(field=33, model_name="bagnet33", **kwargs)
+    return get_bagnet(
+        field=33,
+        model_name="bagnet33",
+        **kwargs)
 
 
 def _test():

pytorch/pytorchcv/models/bamresnet.py

@@ -23,8 +23,8 @@ class DenseBlock(nn.Module):
         Number of output features.
     """
     def __init__(self,
-                 in_features,
-                 out_features):
+                 in_features: int,
+                 out_features: int):
         super(DenseBlock, self).__init__()
         self.fc = nn.Linear(
             in_features=in_features,
@@ -53,9 +53,9 @@ class ChannelGate(nn.Module):
         Number of dense blocks.
     """
     def __init__(self,
-                 channels,
-                 reduction_ratio=16,
-                 num_layers=1):
+                 channels: int,
+                 reduction_ratio: int = 16,
+                 num_layers: int = 1):
         super(ChannelGate, self).__init__()
         mid_channels = channels // reduction_ratio
 
@@ -99,10 +99,10 @@ class SpatialGate(nn.Module):
         Dilation/padding value for corresponding convolutions.
     """
     def __init__(self,
-                 channels,
-                 reduction_ratio=16,
-                 num_dil_convs=2,
-                 dilation=4):
+                 channels: int,
+                 reduction_ratio: int = 16,
+                 num_dil_convs: int = 2,
+                 dilation: int = 4):
         super(SpatialGate, self).__init__()
         mid_channels = channels // reduction_ratio
 
@@ -145,7 +145,7 @@ class BamBlock(nn.Module):
         Number of input/output channels.
     """
     def __init__(self,
-                 channels):
+                 channels: int):
         super(BamBlock, self).__init__()
         self.ch_att = ChannelGate(channels=channels)
         self.sp_att = SpatialGate(channels=channels)
@@ -173,10 +173,10 @@ class BamResUnit(nn.Module):
         Whether to use a bottleneck or simple block in units.
     """
     def __init__(self,
-                 in_channels,
-                 out_channels,
-                 stride,
-                 bottleneck):
+                 in_channels: int,
+                 out_channels: int,
+                 stride: int | tuple[int, int],
+                 bottleneck: bool):
         super(BamResUnit, self).__init__()
         self.use_bam = (stride != 1)
 
@@ -216,9 +216,9 @@ class BamResNet(nn.Module):
         Number of classification classes.
     """
     def __init__(self,
-                 channels,
-                 init_block_channels,
-                 bottleneck,
+                 channels: list[list[int]],
+                 init_block_channels: int,
+                 bottleneck: bool,
                  in_channels: int = 3,
                  in_size: tuple[int, int] = (224, 224),
                  num_classes: int = 1000):
@@ -355,7 +355,10 @@ def bam_resnet18(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_resnet(blocks=18, model_name="bam_resnet18", **kwargs)
+    return get_resnet(
+        blocks=18,
+        model_name="bam_resnet18",
+        **kwargs)
 
 
 def bam_resnet34(**kwargs) -> nn.Module:
@@ -374,7 +377,10 @@ def bam_resnet34(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_resnet(blocks=34, model_name="bam_resnet34", **kwargs)
+    return get_resnet(
+        blocks=34,
+        model_name="bam_resnet34",
+        **kwargs)
 
 
 def bam_resnet50(**kwargs) -> nn.Module:
@@ -393,7 +399,10 @@ def bam_resnet50(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_resnet(blocks=50, model_name="bam_resnet50", **kwargs)
+    return get_resnet(
+        blocks=50,
+        model_name="bam_resnet50",
+        **kwargs)
 
 
 def bam_resnet101(**kwargs) -> nn.Module:
@@ -412,7 +421,10 @@ def bam_resnet101(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_resnet(blocks=101, model_name="bam_resnet101", **kwargs)
+    return get_resnet(
+        blocks=101,
+        model_name="bam_resnet101",
+        **kwargs)
 
 
 def bam_resnet152(**kwargs) -> nn.Module:
@@ -431,7 +443,10 @@ def bam_resnet152(**kwargs) -> nn.Module:
     nn.Module
         Desired module.
     """
-    return get_resnet(blocks=152, model_name="bam_resnet152", **kwargs)
+    return get_resnet(
+        blocks=152,
+        model_name="bam_resnet152",
+        **kwargs)
 
 
 def _test():