SEBlocks refactor.

Fixes #28

Author: ayasyrev

model_constructor/__init__.py

@@ -1,4 +1,4 @@
-from model_constructor.model_constructor import ModelConstructor  # noqa F401
+from model_constructor.model_constructor import ModelConstructor, ResBlock  # noqa F401
 
 
 __version__ = "0.1.8"

model_constructor/layers.py

@@ -102,7 +102,7 @@ def forward(self, x):
         return o.view(*size).contiguous()
 
 
-class SEBlock(nn.Module):
+class SEBlock(nn.Module):  # todo: deprecation worning.
     "se block"
     se_layer = nn.Linear
     act_fn = nn.ReLU(inplace=True)
@@ -126,7 +126,7 @@ def forward(self, x):
         return x * y.expand_as(x)
 
 
-class SEBlockConv(nn.Module):
+class SEBlockConv(nn.Module):  # todo: deprecation worning.
     "se block with conv on excitation"
     se_layer = nn.Conv2d
     act_fn = nn.ReLU(inplace=True)
@@ -149,3 +149,60 @@ def forward(self, x):
         y = self.squeeze(x)
         y = self.excitation(y)
         return x * y.expand_as(x)
+
+
+class SEModule(nn.Module):
+    "se block"
+
+    def __init__(self,
+                 channels,
+                 reduction=16,
+                 se_layer=nn.Linear,
+                 act_fn=nn.ReLU(inplace=True),  # ? obj or class?
+                 use_bias=True,
+                 gate=nn.Sigmoid
+                 ):
+        super().__init__()
+        rd_channels = channels // reduction
+        self.squeeze = nn.AdaptiveAvgPool2d(1)
+        self.excitation = nn.Sequential(
+            OrderedDict([('fc_reduce', se_layer(channels, rd_channels, bias=use_bias)),
+                         ('se_act', act_fn),
+                         ('fc_expand', se_layer(rd_channels, channels, bias=use_bias)),
+                         ('se_gate', gate())
+                         ]))
+
+    def forward(self, x):
+        bs, c, _, _ = x.shape
+        y = self.squeeze(x).view(bs, c)
+        y = self.excitation(y).view(bs, c, 1, 1)
+        return x * y.expand_as(x)
+
+
+class SEModuleConv(nn.Module):
+    "se block with conv on excitation"
+
+    def __init__(self,
+                 channels,
+                 reduction=16,
+                 se_layer=nn.Conv2d,
+                 act_fn=nn.ReLU(inplace=True),
+                 use_bias=True,
+                 gate=nn.Sigmoid
+                 ):
+        super().__init__()
+#       rd_channels = math.ceil(channels//reduction/8)*8
+        rd_channels = channels // reduction
+        self.squeeze = nn.AdaptiveAvgPool2d(1)
+        self.excitation = nn.Sequential(
+            OrderedDict([
+                ('conv_reduce', se_layer(channels, rd_channels, 1, bias=use_bias)),
+                ('se_act', act_fn),
+                ('conv_expand', se_layer(rd_channels, channels, 1, bias=use_bias)),
+                ('gate', gate())
+            ]))
+
+    def forward(self, x):
+        y = self.squeeze(x)
+        y = self.excitation(y)
+        return x * y.expand_as(x)

model_constructor/model_constructor.py

@@ -3,7 +3,7 @@
 
 import torch.nn as nn
 
-from .layers import ConvLayer, Flatten, SEBlock, SimpleSelfAttention, noop
+from .layers import ConvLayer, Flatten, SEModule, SimpleSelfAttention, noop
 
 
 __all__ = ['init_cnn', 'act_fn', 'ResBlock', 'ModelConstructor', 'xresnet34', 'xresnet50']
@@ -24,12 +24,13 @@ def init_cnn(module: nn.Module):
 
 class ResBlock(nn.Module):
     '''Resnet block'''
-    se_block = SEBlock
 
     def __init__(self, expansion, ni, nh, stride=1,
                  conv_layer=ConvLayer, act_fn=act_fn, zero_bn=True, bn_1st=True,
-                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False, se=False, se_reduction=16,
-                 groups=1, dw=False, div_groups=None):
+                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False,
+                 groups=1, dw=False, div_groups=None,
+                 se_module=SEModule, se=False, se_reduction=16
+                 ):
         super().__init__()
         nf, ni = nh * expansion, ni * expansion
         if div_groups is not None:  # check if grops != 1 and div_groups
@@ -46,7 +47,7 @@ def __init__(self, expansion, ni, nh, stride=1,
                       ("conv_2", conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False, bn_1st=bn_1st))
                       ]
         if se:
-            layers.append(('se', self.se_block(nf, se_reduction)))
+            layers.append(('se', se_module(nf, se_reduction)))
         if sa:
             layers.append(('sa', SimpleSelfAttention(nf, ks=1, sym=sym)))
         self.convs = nn.Sequential(OrderedDict(layers))
@@ -76,7 +77,7 @@ def _make_layer(self, expansion, ni, nf, blocks, stride, sa):
                                      conv_layer=self.conv_layer, act_fn=self.act_fn, pool=self.pool,
                                      zero_bn=self.zero_bn, bn_1st=self.bn_1st,
                                      groups=self.groups, div_groups=self.div_groups,
-                                     dw=self.dw, se=self.se))
+                                     dw=self.dw, se_module=self.se_module, se=self.se, se_reduction=self.se_reduction))
               for i in range(blocks)]
     return nn.Sequential(OrderedDict(layers))
 
@@ -106,7 +107,8 @@ def __init__(self, name='MC', c_in=3, c_out=1000,
                  act_fn=nn.ReLU(inplace=True),
                  pool=nn.AvgPool2d(2, ceil_mode=True),
                  expansion=1, groups=1, dw=False, div_groups=None,
-                 sa=False, se=False, se_reduction=16,
+                 sa=False,
+                 se=False, se_module=SEModule, se_reduction=16,
                  bn_1st=True,
                  zero_bn=True,
                  stem_stride_on=0,
@@ -150,7 +152,7 @@ def __call__(self):
             ('body', self.body),
             ('head', self.head)]))
         self._init_cnn(model)
-        model.extra_repr = lambda: f"model {self.name}"
+        model.extra_repr = lambda: f"{self.name}"
         return model
 
     def __repr__(self):

model_constructor/yaresnet.py

@@ -4,7 +4,7 @@
 import torch.nn as nn
 from functools import partial
 from collections import OrderedDict
-from .layers import SEBlock, ConvLayer, act_fn, noop, SimpleSelfAttention
+from .layers import SEModule, ConvLayer, act_fn, noop, SimpleSelfAttention
 from .net import Net
 from torch.nn import Mish
 
@@ -14,12 +14,13 @@
 
 class YaResBlock(nn.Module):
     '''YaResBlock. Reduce by pool instead of stride 2'''
-    se_block = SEBlock
 
     def __init__(self, expansion, ni, nh, stride=1,
                  conv_layer=ConvLayer, act_fn=act_fn, zero_bn=True, bn_1st=True,
-                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False, se=False,
-                 groups=1, dw=False, div_groups=None):
+                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False,
+                 groups=1, dw=False, div_groups=None,
+                 se_module=SEModule, se=False, se_reduction=16
+                 ):
         super().__init__()
         nf, ni = nh * expansion, ni * expansion
         if div_groups is not None:  # check if grops != 1 and div_groups
@@ -35,7 +36,7 @@ def __init__(self, expansion, ni, nh, stride=1,
                       ("conv_2", conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False, bn_1st=bn_1st))
         ]
         if se:
-            layers.append(('se', self.se_block(nf)))
+            layers.append(('se', se_module(nf, se_reduction)))
         if sa:
             layers.append(('sa', SimpleSelfAttention(nf, ks=1, sym=sym)))
         self.convs = nn.Sequential(OrderedDict(layers))