typing

Author: ayasyrev

src/model_constructor/layers.py

@@ -258,14 +258,14 @@ class SEModule(nn.Module):
 
     def __init__(
         self,
-        channels,
-        reduction=16,
-        rd_channels=None,
-        rd_max=False,
-        se_layer=nn.Linear,
-        act_fn=nn.ReLU(inplace=True),
-        use_bias=True,
-        gate=nn.Sigmoid,
+        channels: int,
+        reduction: int = 16,
+        rd_channels: Optional[int] = None,
+        rd_max: bool = False,
+        se_layer: type[nn.Module] = nn.Linear,
+        act_fn: nn.Module = nn.ReLU(inplace=True),
+        use_bias: bool = True,
+        gate: type[nn.Module] = nn.Sigmoid,
     ):
         super().__init__()
         reducted = max(channels // reduction, 1)  # preserve zero-element tensors
@@ -286,7 +286,7 @@ def __init__(
             )
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         bs, c, _, _ = x.shape
         y = self.squeeze(x).view(bs, c)
         y = self.excitation(y).view(bs, c, 1, 1)
@@ -298,14 +298,14 @@ class SEModuleConv(nn.Module):
 
     def __init__(
         self,
-        channels,
-        reduction=16,
-        rd_channels=None,
-        rd_max=False,
-        se_layer=nn.Conv2d,
-        act_fn=nn.ReLU(inplace=True),
-        use_bias=True,
-        gate=nn.Sigmoid,
+        channels: int,
+        reduction: int = 16,
+        rd_channels: Optional[int] = None,
+        rd_max: bool = False,
+        se_layer: type[nn.Module] = nn.Conv2d,
+        act_fn: nn.Module = nn.ReLU(inplace=True),
+        use_bias: bool = True,
+        gate: type[nn.Module] = nn.Sigmoid,
     ):
         super().__init__()
         #       rd_channels = math.ceil(channels//reduction/8)*8
@@ -327,7 +327,7 @@ def __init__(
             )
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         y = self.squeeze(x)
         y = self.excitation(y)
         return x * y.expand_as(x)

src/model_constructor/model_constructor.py

@@ -42,7 +42,7 @@ class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
     zero_bn: bool = True
     stem_stride_on: int = 0
     stem_sizes: list[int] = [64]
-    stem_pool: Union[Callable[[], nn.Module], None] = partial(
+    stem_pool: Optional[Callable[[], nn.Module]] = partial(
         nn.MaxPool2d, kernel_size=3, stride=2, padding=1
     )
     stem_bn_end: bool = False
@@ -61,7 +61,7 @@ def __repr__(self) -> str:
         )
 
 
-def make_stem(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+def make_stem(cfg: ModelCfg) -> nn.Sequential:
     """Create Resnet stem."""
     stem: ListStrMod = [
         (
@@ -116,15 +116,15 @@ def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:  # type: ignore
     )
 
 
-def make_body(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+def make_body(cfg: ModelCfg) -> nn.Sequential:
     """Create model body."""
     return nn_seq(
         (f"l_{layer_num}", cfg.make_layer(cfg, layer_num))  # type: ignore
         for layer_num in range(len(cfg.layers))
     )
 
 
-def make_head(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+def make_head(cfg: ModelCfg) -> nn.Sequential:
     """Create head."""
     head = [
         ("pool", nn.AdaptiveAvgPool2d(1)),
@@ -138,10 +138,10 @@ class ModelConstructor(ModelCfg):
     """Model constructor. As default - resnet18"""
 
     init_cnn: Callable[[nn.Module], None] = init_cnn
-    make_stem: Callable[[ModelCfg], ModSeq] = make_stem  # type: ignore
-    make_layer: Callable[[ModelCfg, int], ModSeq] = make_layer  # type: ignore
-    make_body: Callable[[ModelCfg], ModSeq] = make_body  # type: ignore
-    make_head: Callable[[ModelCfg], ModSeq] = make_head  # type: ignore
+    make_stem: Callable[[ModelCfg], ModSeq] = make_stem
+    make_layer: Callable[[ModelCfg, int], ModSeq] = make_layer
+    make_body: Callable[[ModelCfg], ModSeq] = make_body
+    make_head: Callable[[ModelCfg], ModSeq] = make_head
 
     @field_validator("se")
     def set_se(  # pylint: disable=no-self-argument

src/model_constructor/universal_blocks.py

@@ -134,7 +134,7 @@ def __init__(
             self.id_conv = None
         self.act_fn = get_act(act_fn)
 
-    def forward(self, x: torch.Tensor):  # type: ignore
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         identity = self.id_conv(x) if self.id_conv is not None else x
         return self.act_fn(self.convs(x) + identity)
 
@@ -255,14 +255,14 @@ def __init__(
             self.id_conv = None
         self.merge = get_act(act_fn)
 
-    def forward(self, x: torch.Tensor):  # type: ignore
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.reduce:
             x = self.reduce(x)
         identity = self.id_conv(x) if self.id_conv is not None else x
         return self.merge(self.convs(x) + identity)
 
 
-def make_stem(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+def make_stem(cfg: ModelCfg) -> nn.Sequential:
     """Create xResnet stem -> 3 conv 3*3 instead of 1 conv 7*7"""
     len_stem = len(cfg.stem_sizes)
     stem: list[tuple[str, nn.Module]] = [
@@ -286,7 +286,7 @@ def make_stem(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
     return nn_seq(stem)
 
 
-def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:  # type: ignore
+def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:
     """Create layer (stage)"""
     # if no pool on stem - stride = 2 for first layer block in body
     stride = 1 if cfg.stem_pool and layer_num == 0 else 2
@@ -316,15 +316,15 @@ def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:  # type: ignore
     )
 
 
-def make_body(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+def make_body(cfg: ModelCfg) -> nn.Sequential:
     """Create model body."""
     return nn_seq(
         (f"l_{layer_num}", cfg.make_layer(cfg, layer_num))  # type: ignore
         for layer_num in range(len(cfg.layers))
     )
 
 
-def make_head(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+def make_head(cfg: ModelCfg) -> nn.Sequential:
     """Create head."""
     head = [
         ("pool", nn.AdaptiveAvgPool2d(1)),

src/model_constructor/xresnet.py

@@ -14,7 +14,7 @@
 ]
 
 
-def xresnet_stem(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+def xresnet_stem(cfg: ModelCfg) -> nn.Sequential:
     """Create xResnet stem -> 3 conv 3*3 instead 1 conv 7*7"""
     len_stem = len(cfg.stem_sizes)
     stem: ListStrMod = [

src/model_constructor/yaresnet.py

@@ -100,7 +100,7 @@ def __init__(
             self.id_conv = None
         self.merge = get_act(act_fn)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.reduce:
             x = self.reduce(x)
         identity = self.id_conv(x) if self.id_conv is not None else x
@@ -195,7 +195,7 @@ def __init__(
             self.id_conv = None
         self.merge = get_act(act_fn)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.reduce:
             x = self.reduce(x)
         identity = self.id_conv(x) if self.id_conv is not None else x