Merge pull request #103 from ayasyrev/tests

Tests

Author: ayasyrev

src/model_constructor/blocks.py

@@ -66,7 +66,7 @@ def __init__(
             id_layers: ListStrMod = []
             if (
                 stride != 1 and pool is not None
-            ):  # if pool - reduce by pool else stride 2 art id_conv
+            ):  # if pool - reduce by pool else stride 2 at id_conv
                 id_layers.append(("pool", pool()))
             if in_channels != out_channels or (stride != 1 and pool is None):
                 id_layers.append(

src/model_constructor/helpers.py

@@ -103,7 +103,7 @@ def __repr_changed_args__(self) -> list[str]:
 
     def print_cfg(self) -> None:
         """Print full config"""
-        print(f"{self.__repr_name__()}(\n  {self.__repr_str__(chr(10) + '  ')})")
+        print(self.__repr__())
 
     def print_changed(self) -> None:
         """Print changed fields."""

src/model_constructor/model_constructor.py

@@ -67,7 +67,7 @@ def set_modules(  # pylint: disable=no-self-argument
             return value
         if isinstance(value, str):
             return instantiate_module(value)
-        raise ValueError(f"{info.field_name} must be str or nn.Module")
+        # raise ValueError(f"{info.field_name} must be str or nn.Module")
 
     @field_validator("se", "sa")
     def set_se(  # pylint: disable=no-self-argument
@@ -77,7 +77,7 @@ def set_se(  # pylint: disable=no-self-argument
             return DEFAULT_SE_SA[info.field_name]
         if is_module(value):
             return value
-        raise ValueError(f"{info.field_name} must be bool or nn.Module")
+        # raise ValueError(f"{info.field_name} must be bool or nn.Module")  # no need - check at init
 
     @field_validator("se_module", "se_reduction")  # pragma: no cover
     def deprecation_warning(  # pylint: disable=no-self-argument

tests/test_helpers.py

@@ -0,0 +1,71 @@
+from functools import partial
+
+from pytest import CaptureFixture
+from torch import nn
+
+from model_constructor.helpers import Cfg, instantiate_module, is_module
+
+
+class Cfg2(Cfg):
+    int_value: int = 10
+
+
+def test_is_module():
+    """test is_module"""
+    assert not is_module("some string")
+    assert is_module(nn.Module)
+    assert is_module(nn.ReLU)
+    assert not is_module(nn)
+    assert is_module(partial(nn.ReLU, inplace=True))
+    assert not is_module(partial(int, "10"))
+
+
+def test_instantiate_module():
+    """test instantiate_module"""
+    mod = instantiate_module("ReLU")
+    assert mod is nn.ReLU
+    mod = instantiate_module("nn.Tanh")
+    assert mod is nn.Tanh
+    mod = instantiate_module("torch.nn.SELU")
+    assert mod is nn.SELU
+    # wrong name
+    try:
+        mod = instantiate_module("wrong_name")
+    except ImportError as err:
+        assert str(err) == "Module wrong_name not found at torch.nn"
+    # wrong module
+    try:
+        mod = instantiate_module("wrong_module.some_name")
+    except ImportError as err:
+        assert str(err) == "Module wrong_module not found"
+    # not nn.Module
+    try:
+        mod = instantiate_module("model_constructor.helpers.instantiate_module")
+    except ImportError as err:
+        assert str(err) == "Module instantiate_module is not a nn.Module"
+
+
+def test_cfg_repr_print(capsys: CaptureFixture[str]):
+    """test repr and print results"""
+    cfg = Cfg()
+    repr_res = cfg.__repr__()
+    assert repr_res == "Cfg(\n  )"
+    cfg.print_changed()
+    out = capsys.readouterr().out
+    assert out == "Nothing changed\n"
+    cfg.name = "cfg_name"
+    repr_res = cfg.__repr__()
+    assert repr_res == "Cfg(\n  name='cfg_name')"
+    cfg.print_cfg()
+    out = capsys.readouterr().out
+    assert out == "Cfg(\n  name='cfg_name')\n"
+    # changed fields. default - name is not in changed
+    cfg = Cfg2(name="cfg_name")
+    cfg.print_changed()
+    out = capsys.readouterr().out
+    assert out == "Nothing changed\n"
+    assert "name" in cfg.model_fields_set
+    cfg = Cfg2(int_value=0)
+    cfg.print_changed()
+    out = capsys.readouterr().out
+    assert out == "Changed fields:\n   int_value: 0\n"

tests/test_mc.py

@@ -2,9 +2,9 @@
 
 import torch
 
-from model_constructor.blocks import BottleneckBlock
+from model_constructor.blocks import BasicBlock, BottleneckBlock
 from model_constructor.layers import SEModule, SEModuleConv, SimpleSelfAttention
-from model_constructor.model_constructor import ModelConstructor
+from model_constructor.model_constructor import ModelCfg, ModelConstructor
 
 bs_test = 4
 in_chans = 3
@@ -94,3 +94,46 @@ def test_MC_bottleneck():
     assert model.body.l_0.bl_0.convs.conv_0.conv.in_channels == 64
     assert model.body.l_0.bl_0.convs.conv_0.conv.out_channels == 128
     assert model.body.l_0.bl_1.convs.conv_0.conv.in_channels == 256
+
+
+def test_ModelCfg():
+    """test ModelCfg"""
+    # default - just create config with custom name
+    cfg = ModelCfg(name="custom_name")
+    repr_str = cfg.__repr__()
+    assert repr_str.startswith("custom_name")
+    # initiate from string
+    cfg = ModelCfg(act_fn="torch.nn.Mish")
+    assert cfg.act_fn is torch.nn.Mish
+    # wrong name
+    try:
+        cfg = ModelCfg(act_fn="wrong_name")
+    except ImportError as err:
+        assert str(err) == "Module wrong_name not found at torch.nn"
+    cfg = ModelCfg(act_fn="nn.Tanh")
+    assert cfg.act_fn is torch.nn.Tanh
+    cfg = ModelCfg(block="model_constructor.blocks.BottleneckBlock")
+    assert cfg.block is BottleneckBlock
+
+
+def test_create_model_class_methods():
+    """test class methods ModelConstructor"""
+    # create model
+    model = ModelConstructor.create_model(act_fn="Mish", num_classes=10)
+    assert str(model.body.l_0.bl_0.convs.conv_0.act_fn) == "Mish(inplace=True)"
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, 10])
+    # from cfg
+    cfg = ModelCfg(block=BottleneckBlock, num_classes=10)
+    mc = ModelConstructor.from_cfg(cfg)
+    model = mc()
+    assert isinstance(model.body.l_0.bl_0, BottleneckBlock)
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, 10])
+
+    cfg.block = BasicBlock
+    cfg.num_classes = 2
+    model = ModelConstructor.create_model(cfg)
+    assert isinstance(model.body.l_0.bl_0, BasicBlock)
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, 2])

tests/test_models.py

@@ -40,3 +40,14 @@ def test_mc(model_constructor: type[ModelConstructor], act_fn: type[nn.Module]):
     model = mc()
     pred = model(xb)
     assert pred.shape == torch.Size([bs_test, 1000])
+
+
+def test_xresnet_stem():
+    """test xresnet stem"""
+    mc = XResNet()
+    assert mc.stem_bn_end == False
+    mc.stem_bn_end = True
+    stem = mc.stem
+    assert isinstance(stem[-1], nn.BatchNorm2d)
+    stem_out = stem(xb)
+    assert stem_out.shape == torch.Size([bs_test, 64, 4, 4])

tests/test_models_universal_blocks.py

@@ -35,3 +35,14 @@ def test_mc(model_constructor: type[ModelConstructor], act_fn: type[nn.Module]):
     model = mc()
     pred = model(xb)
     assert pred.shape == torch.Size([bs_test, 1000])
+
+
+def test_stem_bn_end():
+    """test stem"""
+    mc = XResNet()
+    assert mc.stem_bn_end == False
+    mc.stem_bn_end = True
+    stem = mc.stem
+    assert isinstance(stem[-1], nn.BatchNorm2d)
+    stem_out = stem(xb)
+    assert stem_out.shape == torch.Size([bs_test, 64, 4, 4])