TESTS FIX

Author: ayasyrev

src/model_constructor/yaresnet.py

@@ -11,13 +11,12 @@
 from .model_constructor import ModelConstructor
 
 __all__ = [
-    'YaResBlock',
+    "YaResBlock",
+    "YaResNet34",
+    "YaResNet50",
 ]
 
 
-# act_fn = nn.ReLU(inplace=True)
-
-
 class YaResBlock(nn.Module):
     '''YaResBlock. Reduce by pool instead of stride 2'''
 

tests/test_Net.py

@@ -2,6 +2,7 @@
 import torch.nn as nn
 
 from model_constructor.net import Net, NewResBlock, ResBlock
+
 # from model_constructor.layers import SEModule, SimpleSelfAttention
 
 
@@ -19,7 +20,7 @@
     bn_1st=[True, False],
     zero_bn=[True, False],
     stem_bn_end=[True, False],
-    stem_stride_on=[0, 1]
+    stem_stride_on=[0, 1],
 )
 
 
@@ -28,9 +29,8 @@ def value_name(value) -> str:  # pragma: no cover
     if name is not None:
         return name
     if isinstance(value, nn.Module):
-        return value._get_name()
-    else:
-        return value
+        return value._get_name()  # pylint: disable=W0212
+    return value
 
 
 def ids_fn(key, value):
@@ -44,7 +44,8 @@ def pytest_generate_tests(metafunc):
 
 
 def test_Net(
-    block, expansion,
+    block,
+    expansion,
     groups,
 ):
     """test Net"""
@@ -54,15 +55,14 @@ def test_Net(
     name = "Test name"
 
     mc = Net(
-        name, c_in, c_out, block,
+        name,
+        c_in,
+        c_out,
+        block,
         expansion=expansion,
         stem_sizes=[8, 16],
         block_sizes=[16, 32, 64, 128],
         groups=groups,
-        # dw=dw,
-        # div_groups=div_groups,
-        # bn_1st=bn_1st, zero_bn=zero_bn,
-        # stem_bn_end=stem_bn_end,
     )
     assert f"{name} constructor" in str(mc)
     model = mc()
@@ -71,22 +71,23 @@ def test_Net(
     assert pred.shape == torch.Size([bs_test, c_out])
 
 
-def test_Net_SE_SA(
-    block, expansion,
-    se, sa
-):
+def test_Net_SE_SA(block, expansion, se, sa):
     """test Net"""
     c_in = 3
     img_size = 16
     c_out = 8
     name = "Test name"
 
     mc = Net(
-        name, c_in, c_out, block,
+        name,
+        c_in,
+        c_out,
+        block,
         expansion=expansion,
         stem_sizes=[8, 16],
         block_sizes=[16, 32, 64, 128],
-        se=se, sa=sa
+        se=se,
+        sa=sa,
     )
     assert f"{name} constructor" in str(mc)
     model = mc()
@@ -96,7 +97,8 @@ def test_Net_SE_SA(
 
 
 def test_Net_div_gr(
-    block, expansion,
+    block,
+    expansion,
     div_groups,
 ):
     """test Net"""
@@ -106,7 +108,10 @@ def test_Net_div_gr(
     name = "Test name"
 
     mc = Net(
-        name, c_in, c_out, block,
+        name,
+        c_in,
+        c_out,
+        block,
         expansion=expansion,
         stem_sizes=[8, 16],
         block_sizes=[16, 32, 64, 128],
@@ -119,22 +124,22 @@ def test_Net_div_gr(
     assert pred.shape == torch.Size([bs_test, c_out])
 
 
-def test_Net_dw(
-    block, expansion,
-    dw
-):
+def test_Net_dw(block, expansion, dw):
     """test Net"""
     c_in = 3
     img_size = 16
     c_out = 8
     name = "Test name"
 
     mc = Net(
-        name, c_in, c_out, block,
+        name,
+        c_in,
+        c_out,
+        block,
         expansion=expansion,
         stem_sizes=[8, 16],
         block_sizes=[16, 32, 64, 128],
-        dw=dw
+        dw=dw,
     )
     assert f"{name} constructor" in str(mc)
     model = mc()
@@ -144,8 +149,10 @@ def test_Net_dw(
 
 
 def test_Net_2(
-    block, expansion,
-    bn_1st, zero_bn,
+    block,
+    expansion,
+    bn_1st,
+    zero_bn,
 ):
     """test Net"""
     c_in = 3
@@ -154,11 +161,15 @@ def test_Net_2(
     name = "Test name"
 
     mc = Net(
-        name, c_in, c_out, block,
+        name,
+        c_in,
+        c_out,
+        block,
         expansion=expansion,
         stem_sizes=[8, 16],
         block_sizes=[16, 32, 64, 128],
-        bn_1st=bn_1st, zero_bn=zero_bn,
+        bn_1st=bn_1st,
+        zero_bn=zero_bn,
     )
     assert f"{name} constructor" in str(mc)
     model = mc()
@@ -167,25 +178,24 @@ def test_Net_2(
     assert pred.shape == torch.Size([bs_test, c_out])
 
 
-def test_Net_stem(
-    stem_bn_end,
-    stem_stride_on
-):
+def test_Net_stem(stem_bn_end, stem_stride_on):
     """test Net"""
     c_in = 3
     img_size = 16
     c_out = 8
     name = "Test name"
 
     mc = Net(
-        name, c_in, c_out,
+        name,
+        c_in,
+        c_out,
         stem_sizes=[8, 16],
         block_sizes=[16, 32, 64, 128],
         stem_bn_end=stem_bn_end,
-        stem_stride_on=stem_stride_on
+        stem_stride_on=stem_stride_on,
     )
     assert f"{name} constructor" in str(mc)
     model = mc()
     xb = torch.randn(bs_test, c_in, img_size, img_size)
     pred = model(xb)
-    assert pred.shape == torch.Size([bs_test, c_out])
+    assert pred.shape == torch.Size([bs_test, c_out])

tests/test_block.py

@@ -29,8 +29,7 @@ def value_name(value) -> str:
         return name
     if isinstance(value, nn.Module):
         return value._get_name()  # pylint: disable=W0212
-    else:
-        return value
+    return value
 
 
 def ids_fn(key, value):
@@ -48,9 +47,15 @@ def test_block(Block, expansion, mid_channels, stride, div_groups, pool, se, sa)
     in_channels = 8
     out_channels = mid_channels * expansion
     block = Block(
-        expansion, in_channels, mid_channels,
-        stride, div_groups=div_groups,
-        pool=pool, se=se, sa=sa)
+        expansion,
+        in_channels,
+        mid_channels,
+        stride,
+        div_groups=div_groups,
+        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

tests/test_convmixer.py

@@ -4,6 +4,7 @@
 from model_constructor.convmixer import ConvMixer, ConvMixerOriginal
 
 bs_test = 4
+img_size = 16
 
 
 params = dict(
@@ -17,9 +18,8 @@ def value_name(value) -> str:  # pragma: no cover
     if name is not None:
         return name
     if isinstance(value, nn.Module):
-        return value._get_name()
-    else:
-        return value
+        return value._get_name()  # pylint: disable=W0212
+    return value
 
 
 def ids_fn(key, value):
@@ -34,8 +34,6 @@ def pytest_generate_tests(metafunc):
 
 def test_ConvMixer(bn_1st, pre_act):
     """test ConvMixer"""
-    bs_test = 4
-    img_size = 16
     model = ConvMixer(dim=64, depth=4, bn_1st=bn_1st, pre_act=pre_act)
     xb = torch.randn(bs_test, 3, img_size, img_size)
     pred = model(xb)
@@ -44,8 +42,6 @@ def test_ConvMixer(bn_1st, pre_act):
 
 def test_ConvMixerOriginal():
     """test ConvMixerOriginal"""
-    bs_test = 4
-    img_size = 16
     model = ConvMixerOriginal(dim=64, depth=4)
     xb = torch.randn(bs_test, 3, img_size, img_size)
     pred = model(xb)

tests/test_layers.py

@@ -1,7 +1,15 @@
 import torch
 import torch.nn as nn
 
-from model_constructor.layers import ConvBnAct, Flatten, Noop, SEModule, SEModuleConv, SimpleSelfAttention, noop
+from model_constructor.layers import (
+    ConvBnAct,
+    Flatten,
+    Noop,
+    SEModule,
+    SEModuleConv,
+    SimpleSelfAttention,
+    noop,
+)
 
 
 bs_test = 4
@@ -34,9 +42,8 @@ def value_name(value) -> str:
     if name is not None:
         return name
     if isinstance(value, nn.Module):  # pragma: no cover
-        return value._get_name()
-    else:
-        return value
+        return value._get_name()  # pylint: disable=W0212
+    return value
 
 
 def ids_fn(key, value):
@@ -71,14 +78,25 @@ def test_noop():
     assert all(out.eq(xb_copy))
 
 
-def test_ConvBnAct(kernel_size, stride, bias, groups, pre_act, bn_layer, bn_1st, zero_bn):
+def test_ConvBnAct(
+    kernel_size, stride, bias, groups, pre_act, bn_layer, bn_1st, zero_bn
+):
     """test ConvBnAct"""
     in_channels = out_channels = 4
     channel_size = 4
     block = ConvBnAct(
-        in_channels, out_channels, kernel_size, stride,
-        padding=None, bias=bias, groups=groups,
-        pre_act=pre_act, bn_layer=bn_layer, bn_1st=bn_1st, zero_bn=zero_bn)
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride,
+        padding=None,
+        bias=bias,
+        groups=groups,
+        pre_act=pre_act,
+        bn_layer=bn_layer,
+        bn_1st=bn_1st,
+        zero_bn=zero_bn,
+    )
     xb = torch.randn(bs_test, in_channels, channel_size, channel_size)
     out = block(xb)
     out_size = channel_size

tests/test_layers_depr.py

@@ -24,18 +24,17 @@
     bn_1st=[True, False],
     zero_bn=[False, True],
     bias=[False, True],
-    groups=[1, 2]
+    groups=[1, 2],
 )
 
 
-def value_name(value) -> str:
+def value_name(value) -> str:  # pragma: no cover
     name = getattr(value, "__name__", None)
     if name is not None:
         return name
     if isinstance(value, nn.Module):
-        return value._get_name()  # pragma: no cover
-    else:
-        return value
+        return value._get_name()  # pylint: disable=W0212
+    return value
 
 
 def ids_fn(key, value):
@@ -64,9 +63,16 @@ def test_ConvLayer(nf, ks, stride, bn_layer, bn_1st, zero_bn, bias, groups):
     ni = 8
     channel_size = 4
     block = ConvLayer(
-        ni, nf, ks, stride,
-        bn_layer=bn_layer, bn_1st=bn_1st, zero_bn=zero_bn,
-        bias=bias, groups=groups)
+        ni,
+        nf,
+        ks,
+        stride,
+        bn_layer=bn_layer,
+        bn_1st=bn_1st,
+        zero_bn=zero_bn,
+        bias=bias,
+        groups=groups,
+    )
     xb = torch.randn(bs_test, ni, channel_size, channel_size)
     out = block(xb)
     # out_ch = nf