tests models

Author: ayasyrev

src/model_constructor/model_constructor.py

@@ -1,6 +1,6 @@
 from collections import OrderedDict
 from functools import partial
-from typing import Any, Callable, List, Optional, Type, TypeVar, Union
+from typing import Any, Callable, List, Type, TypeVar, Union
 
 import torch.nn as nn
 from pydantic import BaseModel, root_validator
@@ -236,11 +236,11 @@ class ModelCfg(BaseModel):
     stem_sizes: List[int] = [32, 32, 64]
     stem_pool: Union[Callable[[], nn.Module], None] = partial(nn.MaxPool2d, kernel_size=3, stride=2, padding=1)
     stem_bn_end: bool = False
-    init_cnn: Optional[Callable[[nn.Module], None]] = init_cnn
-    make_stem: Optional[Callable[[TModelCfg], Union[nn.Module, nn.Sequential]]] = make_stem
-    make_layer: Optional[Callable[[TModelCfg, int], Union[nn.Module, nn.Sequential]]] = make_layer
-    make_body: Optional[Callable[[TModelCfg], Union[nn.Module, nn.Sequential]]] = make_body
-    make_head: Optional[Callable[[TModelCfg], Union[nn.Module, nn.Sequential]]] = make_head
+    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
 
     class Config:
         arbitrary_types_allowed = True
@@ -262,17 +262,6 @@ class ModelConstructor(ModelCfg):
 
     @root_validator
     def post_init(cls, values):  # pylint: disable=E0213
-        # if values["init_cnn"] is None:
-        #     values["init_cnn"] = init_cnn
-        # if values["make_stem"] is None:
-        #     values["make_stem"] = make_stem
-        # if values["make_layer"] is None:
-        #     values["make_layer"] = make_layer
-        # if values["make_body"] is None:
-        #     values["make_body"] = make_body
-        # if values["make_head"] is None:
-        #     values["make_head"] = make_head
-
         if values["stem_sizes"][0] != values["in_chans"]:
             values["stem_sizes"] = [values["in_chans"]] + values["stem_sizes"]
         if values["se"] and isinstance(values["se"], (bool, int)):  # if se=1 or se=True
@@ -287,15 +276,15 @@ def post_init(cls, values):  # pylint: disable=E0213
 
     @property
     def stem(self):
-        return self.make_stem(self)  # type: ignore
+        return self.make_stem(self)  # pylint: disable=too-many-function-args
 
     @property
     def head(self):
-        return self.make_head(self)  # type: ignore
+        return self.make_head(self)  # pylint: disable=too-many-function-args
 
     @property
     def body(self):
-        return self.make_body(self)  # type: ignore
+        return self.make_body(self)  # pylint: disable=too-many-function-args
 
     @classmethod
     def from_cfg(cls, cfg: ModelCfg):
@@ -305,12 +294,12 @@ def __call__(self):
         model = nn.Sequential(
             OrderedDict([("stem", self.stem), ("body", self.body), ("head", self.head)])
         )
-        self.init_cnn(model)  # type: ignore
+        self.init_cnn(model)  # pylint: disable=too-many-function-args
         model.extra_repr = lambda: f"{self.name}"
         return model
 
     def __repr__(self):
-        se_repr = self.se.__name__ if self.se else "False"
+        se_repr = self.se.__name__ if self.se else "False"  # type: ignore
         return (
             f"{self.name} constructor\n"
             f"  in_chans: {self.in_chans}, num_classes: {self.num_classes}\n"

src/model_constructor/xresnet.py

@@ -1,10 +1,10 @@
-import torch.nn as nn
 from collections import OrderedDict
 
+import torch.nn as nn
+
 from .base_constructor import Net
 from .layers import ConvLayer, Noop, act
 
-
 __all__ = ['DownsampleLayer', 'XResBlock', 'xresnet18', 'xresnet34', 'xresnet50']
 
 

src/model_constructor/yaresnet.py

@@ -129,7 +129,7 @@ class YaResNet34(ModelConstructor):
     block: Type[nn.Module] = YaResBlock
     expansion: int = 1
     layers: List[int] = [3, 4, 6, 3]
-    act_fn: nn.Module = Mish()
+    act_fn: Type[nn.Module] = Mish
 
 
 class YaResNet50(YaResNet34):

tests/test_models.py

@@ -0,0 +1,24 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from model_constructor.model_constructor import XResNet34, XResNet50
+from model_constructor.yaresnet import YaResNet34, YaResNet50
+
+bs_test = 2
+img_size = 16
+xb = torch.rand(bs_test, 3, img_size, img_size)
+
+mc_list = [XResNet34, XResNet50, YaResNet34, YaResNet50]
+act_fn_list = [nn.ReLU, nn.Mish, nn.GELU]
+
+
+@pytest.mark.parametrize("model_constructor", mc_list)
+@pytest.mark.parametrize("act_fn", act_fn_list)
+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

@@ -0,0 +1,36 @@
+import pytest
+import torch
+import torch.nn as nn
+
+from model_constructor.xresnet import xresnet18, xresnet34, xresnet50
+from model_constructor.mxresnet import mxresnet34, mxresnet50
+
+bs_test = 2
+img_size = 16
+xb = torch.rand(bs_test, 3, img_size, img_size)
+
+num_classes = 10
+
+models_list = [xresnet18, xresnet34, xresnet50]
+act_fn_list = [nn.ReLU, nn.Mish, nn.GELU]
+mx_list = [mxresnet34, mxresnet50]
+
+
+@pytest.mark.parametrize("model", models_list)
+def test_model(model):
+    """test models"""
+    mod = model(num_classes=num_classes)
+    pred = mod(xb)
+    assert pred.shape == torch.Size([bs_test, num_classes])
+
+
+@pytest.mark.parametrize("model_constructor", mx_list)
+@pytest.mark.parametrize("act_fn", act_fn_list)
+def test_model_mx(model_constructor, act_fn):
+    """test models"""
+    mc = model_constructor(c_out=num_classes)
+    assert mc.c_out == num_classes
+    mc.act_fn = act_fn()
+    model = mc()
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, num_classes])