disable type warning

Author: ayasyrev

src/model_constructor/model_constructor.py

@@ -129,7 +129,7 @@ def forward(self, x):
         return self.act_fn(self.convs(x) + identity)
 
 
-def make_stem(cfg: TModelCfg) -> nn.Sequential:
+def make_stem(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
     stem: List[tuple[str, nn.Module]] = [
         (f"conv_{i}", cfg.conv_layer(
             cfg.stem_sizes[i],  # type: ignore
@@ -150,7 +150,7 @@ def make_stem(cfg: TModelCfg) -> nn.Sequential:
     return nn.Sequential(OrderedDict(stem))
 
 
-def make_layer(cfg: TModelCfg, layer_num: int) -> nn.Sequential:
+def make_layer(cfg: TModelCfg, layer_num: int) -> nn.Sequential:  # type: ignore
     #  expansion, in_channels, out_channels, blocks, stride, sa):
     # 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
@@ -186,7 +186,7 @@ def make_layer(cfg: TModelCfg, layer_num: int) -> nn.Sequential:
     )
 
 
-def make_body(cfg: TModelCfg) -> nn.Sequential:
+def make_body(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
     return nn.Sequential(
         OrderedDict(
             [
@@ -200,7 +200,7 @@ def make_body(cfg: TModelCfg) -> nn.Sequential:
     )
 
 
-def make_head(cfg: TModelCfg) -> nn.Sequential:
+def make_head(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
     head = [
         ("pool", nn.AdaptiveAvgPool2d(1)),
         ("flat", nn.Flatten()),
@@ -237,10 +237,10 @@ class ModelCfg(BaseModel):
     stem_pool: Union[Callable[[], nn.Module], None] = partial(nn.MaxPool2d, kernel_size=3, stride=2, padding=1)
     stem_bn_end: bool = False
     init_cnn: Callable[[nn.Module], None] = init_cnn
-    make_stem: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_stem
-    make_layer: Callable[[TModelCfg, int], Union[nn.Module, nn.Sequential]] = make_layer
-    make_body: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_body
-    make_head: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_head
+    make_stem: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_stem  # type: ignore
+    make_layer: Callable[[TModelCfg, int], Union[nn.Module, nn.Sequential]] = make_layer  # type: ignore
+    make_body: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_body  # type: ignore
+    make_head: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_head  # type: ignore
 
     class Config:
         arbitrary_types_allowed = True

tests/test_block.py

@@ -28,7 +28,7 @@ def value_name(value) -> str:
     if name is not None:
         return name
     if isinstance(value, nn.Module):
-        return value._get_name()
+        return value._get_name()  # pylint: disable=W0212
     else:
         return value
 

tests/test_models.py

@@ -18,7 +18,6 @@
 def test_mc(model_constructor, act_fn):
     """test models"""
     mc = model_constructor(act_fn=act_fn)
-    # assert "name='MC'" in str()
     model = mc()
     pred = model(xb)
     assert pred.shape == torch.Size([bs_test, 1000])

tests/test_models_old.py

@@ -2,8 +2,8 @@
 import torch
 import torch.nn as nn
 
-from model_constructor.xresnet import xresnet18, xresnet34, xresnet50
 from model_constructor.mxresnet import mxresnet34, mxresnet50
+from model_constructor.xresnet import xresnet18, xresnet34, xresnet50
 
 bs_test = 2
 img_size = 16