Merge pull request #49 from ayasyrev/se_sa_fix

fix se, sa settings at MC

Author: ayasyrev

model_constructor/model_constructor.py

@@ -150,12 +150,16 @@ def __init__(self, name='MC', in_chans=3, num_classes=1000,
         self._block_sizes = params['block_sizes']
         if self.stem_sizes[0] != self.in_chans:
             self.stem_sizes = [self.in_chans] + self.stem_sizes
-        if self.se:  # TODO add check issubclass or isinstance of nn.Module
+        if self.se:
             if type(self.se) in (bool, int):  # if se=1 or se=True
                 self.se = SEModule
+            else:
+                self.se = se  # TODO add check issubclass or isinstance of nn.Module
         if self.sa:  # if sa=1 or sa=True
             if type(self.sa) in (bool, int):
                 self.sa = SimpleSelfAttention  # default: ks=1, sym=sym
+            else:
+                self.sa = sa
         if self.se_module or se_reduction:  # pragma: no cover
             print("Deprecated. Pass se_module as se argument, se_reduction as arg to se.")  # add deprecation worning.
 

tests/test_mc.py

@@ -1,7 +1,7 @@
 import torch
 
 from model_constructor import ModelConstructor
-from model_constructor.layers import SEModule, SimpleSelfAttention
+from model_constructor.layers import SEModule, SEModuleConv, SimpleSelfAttention
 
 
 bs_test = 4
@@ -30,3 +30,9 @@ def test_MC():
     model = mc()
     pred = model(xb)
     assert pred.shape == torch.Size([bs_test, num_classes])
+    mc = ModelConstructor(sa=SimpleSelfAttention, se=SEModuleConv, num_classes=num_classes)
+    assert mc.se is SEModuleConv
+    assert mc.sa is SimpleSelfAttention
+    model = mc()
+    pred = model(xb)
+    assert pred.shape == torch.Size([bs_test, num_classes])