tests for constructors - Mc & Net

Author: ayasyrev

tests/test_Net.py

@@ -46,7 +46,6 @@ def pytest_generate_tests(metafunc):
 def test_Net(
     block, expansion,
     groups,
-    # dw, div_groups,
 ):
     """test Net"""
     c_in = 3
@@ -72,6 +71,30 @@ def test_Net(
     assert pred.shape == torch.Size([bs_test, c_out])
 
 
+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,
+        expansion=expansion,
+        stem_sizes=[8, 16],
+        block_sizes=[16, 32, 64, 128],
+        se=se, sa=sa
+    )
+    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])
+
+
 def test_Net_div_gr(
     block, expansion,
     div_groups,

tests/test_mc.py

@@ -24,6 +24,9 @@ def test_MC():
     model = mc()
     pred = model(xb)
     assert pred.shape == torch.Size([bs_test, num_classes])
-    mc = ModelConstructor(sa=1, se=1)
+    mc = ModelConstructor(sa=1, se=1, num_classes=num_classes)
     assert mc.se is SEModule
     assert mc.sa is SimpleSelfAttention
+    model = mc()
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, num_classes])