fix blocks, tests blocks

Author: ayasyrev

model_constructor/layers.py

@@ -196,7 +196,7 @@ def __init__(self,
                  gate=nn.Sigmoid
                  ):
         super().__init__()
-        reducted = channels // reduction
+        reducted = max(channels // reduction, 1)  # preserve zero-element tensors
         if rd_channels is None:
             rd_channels = reducted
         else:

model_constructor/model_constructor.py

@@ -53,13 +53,18 @@ def __init__(self, expansion, in_channels, mid_channels, stride=1,
         if sa:
             layers.append(('sa', sa(out_channels)))
         self.convs = nn.Sequential(OrderedDict(layers))
-        id_layers = []
-        if stride != 1 and pool:
-            id_layers.append(("pool", pool))
-        id_layers += [] if in_channels == out_channels else [("id_conv", conv_layer(in_channels, out_channels, 1,
-                                                                                    stride=1 if pool else stride,
-                                                                                    act_fn=False))]
-        self.id_conv = None if id_layers == [] else nn.Sequential(OrderedDict(id_layers))
+        if stride != 1 or in_channels != out_channels:
+            id_layers = []
+            if stride != 1 and pool is not None:
+                id_layers.append(("pool", pool))
+            if in_channels != out_channels or (stride != 1 and pool is None):
+                id_layers += [("id_conv", conv_layer(
+                    in_channels, out_channels, 1,
+                    stride=1 if pool else stride,
+                    act_fn=False))]
+            self.id_conv = nn.Sequential(OrderedDict(id_layers))
+        else:
+            self.id_conv = None
         self.act_fn = act_fn
 
     def forward(self, x):

model_constructor/yaresnet.py

@@ -30,8 +30,10 @@ def __init__(self, expansion, in_channels, mid_channels, stride=1,
             groups = int(mid_channels / div_groups)
         if stride != 1:
             if pool is None:
-                raise Exception("pool not passed")
-            self.reduce = pool
+                self.reduce = conv_layer(in_channels, in_channels, 1, stride=2)
+                # warnings.warn("pool not passed")  # need to warn?
+            else:
+                self.reduce = pool
         else:
             self.reduce = None
         layers = [("conv_0", conv_layer(in_channels, mid_channels, 3, stride=1,
@@ -42,7 +44,8 @@ def __init__(self, expansion, in_channels, mid_channels, stride=1,
                       ("conv_0", conv_layer(in_channels, mid_channels, 1, act_fn=act_fn, bn_1st=bn_1st)),
                       ("conv_1", conv_layer(mid_channels, mid_channels, 3, stride=1, act_fn=act_fn, bn_1st=bn_1st,
                                             groups=mid_channels if dw else groups)),
-                      ("conv_2", conv_layer(mid_channels, out_channels, 1, zero_bn=zero_bn, act_fn=False, bn_1st=bn_1st))
+                      ("conv_2", conv_layer(
+                          mid_channels, out_channels, 1, zero_bn=zero_bn, act_fn=False, bn_1st=bn_1st))
         ]
         if se:
             layers.append(('se', se(out_channels)))

tests/__init__.py

@@ -0,0 +1,52 @@
+# import pytest
+import torch
+import torch.nn as nn
+from model_constructor.layers import SEModule, SimpleSelfAttention
+
+from model_constructor.model_constructor import ResBlock
+from model_constructor.yaresnet import YaResBlock
+
+bs_test = 4
+img_size = 16
+
+
+params = {
+    "Block": [ResBlock, YaResBlock],
+    "expansion": [1, 2],
+    "mid_channels": [8, 16],
+    "stride": [1, 2],
+    "pool": [None, nn.AvgPool2d(2, ceil_mode=True)],
+    "se": [None, SEModule],
+    "sa": [None, SimpleSelfAttention],
+}
+
+
+def value_name(value) -> str:
+    name = getattr(value, "__name__", None)
+    if name is not None:
+        return name
+    if isinstance(value, nn.Module):
+        return value._get_name()
+    else:
+        return value
+
+
+def ids_fn(key, value):
+    return [f"{key[:2]}_{value_name(v)}" for v in value]
+
+
+def pytest_generate_tests(metafunc):
+    for key, value in params.items():
+        if key in metafunc.fixturenames:
+            metafunc.parametrize(key, value, ids=ids_fn(key, value))
+
+
+def test_block(Block, expansion, mid_channels, stride, pool, se, sa):
+    """test block"""
+    in_channels = 8
+    out_channels = mid_channels * expansion
+    block = Block(expansion, in_channels, mid_channels, stride, pool=pool, se=se, sa=sa)
+    xb = torch.randn(bs_test, in_channels * expansion, img_size, img_size)
+    y = block(xb)
+    out_size = img_size if stride == 1 else img_size // stride
+    assert y.shape == torch.Size([bs_test, out_channels, out_size, out_size])