tests stem etc

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/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_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_bnend():
+    """test stem"""
+    mc = ModelConstructor()
+    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])