Merge pull request #89 from ayasyrev/mc_refactor

Mc refactor

Author: ayasyrev

src/model_constructor/__init__.py

@@ -1,7 +1,6 @@
 from model_constructor.convmixer import ConvMixer  # noqa F401
 from model_constructor.model_constructor import (
     ModelConstructor,
-    ResBlock,
     ModelCfg,
 )  # noqa F401
 

src/model_constructor/base_constructor.py

@@ -1,10 +1,31 @@
-import torch.nn as nn
+"""First version of constructor.
+"""
+# Used in examples.
+# first implementation of xresnet - inspired by fastai version.
 from collections import OrderedDict
-from .layers import ConvLayer, Noop, Flatten
+from functools import partial
 
+import torch.nn as nn
 
-__all__ = ['act_fn', 'Stem', 'DownsampleBlock', 'BasicBlock', 'Bottleneck', 'BasicLayer', 'Body', 'Head', 'init_model',
-           'Net']
+from .layers import ConvLayer, Flatten, Noop
+
+__all__ = [
+    "act_fn",
+    "Stem",
+    "DownsampleBlock",
+    "BasicBlock",
+    "Bottleneck",
+    "BasicLayer",
+    "Body",
+    "Head",
+    "init_model",
+    "Net",
+    "DownsampleLayer",
+    "XResBlock",
+    "xresnet18",
+    "xresnet34",
+    "xresnet50",
+]
 
 
 act_fn = nn.ReLU(inplace=True)
@@ -162,3 +183,58 @@ def __init__(self, stem=Stem,
              ('head', head(body_out * expansion, num_classes, **kwargs))
              ]))
         self.init_model(self)
+
+
+# xresnet from fastai
+
+
+class DownsampleLayer(nn.Sequential):
+    """Downsample layer for Xresnet Resblock"""
+
+    def __init__(self, conv_layer, ni, nf, stride, act,
+                 pool=nn.AvgPool2d(2, ceil_mode=True), pool_1st=True,
+                 **kwargs):
+        layers = [] if stride == 1 else [('pool', pool)]
+        layers += [] if ni == nf else [('idconv', conv_layer(ni, nf, 1, act=act, **kwargs))]
+        if not pool_1st:
+            layers.reverse()
+        super().__init__(OrderedDict(layers))
+
+
+class XResBlock(nn.Module):
+    '''XResnet block'''
+
+    def __init__(self, ni, nh, expansion=1, stride=1, zero_bn=True,
+                 conv_layer=ConvLayer, act_fn=act_fn, **kwargs):
+        super().__init__()
+        nf, ni = nh * expansion, ni * expansion
+        layers = [('conv_0', conv_layer(ni, nh, 3, stride=stride, act_fn=act_fn, **kwargs)),
+                  ('conv_1', conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False, act_fn=act_fn, **kwargs))
+                  ] if expansion == 1 else [
+                      ('conv_0', conv_layer(ni, nh, 1, act_fn=act_fn, **kwargs)),
+                      ('conv_1', conv_layer(nh, nh, 3, stride=stride, act_fn=act_fn, **kwargs)),
+                      ('conv_2', conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False, act_fn=act_fn, **kwargs))
+        ]
+        self.convs = nn.Sequential(OrderedDict(layers))
+        self.identity = DownsampleLayer(conv_layer, ni, nf, stride,
+                                        act=False, act_fn=act_fn, **kwargs) if ni != nf or stride == 2 else Noop()
+        self.merge = Noop()
+        self.act_fn = act_fn
+
+    def forward(self, x):
+        return self.act_fn(self.merge(self.convs(x) + self.identity(x)))
+
+
+def xresnet18(**kwargs):
+    """Constructs xresnet18 model. """
+    return Net(stem_sizes=[32, 32], block=XResBlock, blocks=[2, 2, 2, 2], expansion=1, **kwargs)
+
+
+def xresnet34(**kwargs):
+    """Constructs xresnet34 model. """
+    return Net(stem_sizes=[32, 32], block=XResBlock, blocks=[3, 4, 6, 3], expansion=1, **kwargs)
+
+
+def xresnet50(**kwargs):
+    """Constructs xresnet50 model. """
+    return Net(stem_sizes=[32, 32], block=XResBlock, blocks=[3, 4, 6, 3], expansion=4, **kwargs)

src/model_constructor/helpers.py

@@ -0,0 +1,9 @@
+from collections import OrderedDict
+from typing import Iterable
+
+from torch import nn
+
+
+def nn_seq(list_of_tuples: Iterable[tuple[str, nn.Module]]) -> nn.Sequential:
+    """return nn.Sequential from OrderedDict from list of tuples"""
+    return nn.Sequential(OrderedDict(list_of_tuples))  #

src/model_constructor/layers.py

@@ -2,7 +2,7 @@
 from typing import List, Optional, Type, Union
 
 import torch
-import torch.nn as nn
+from torch import nn
 from torch.nn.utils.spectral_norm import spectral_norm
 
 __all__ = [
@@ -21,19 +21,19 @@
 class Flatten(nn.Module):
     """flat x to vector"""
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x.view(x.size(0), -1)
 
 
-def noop(x):
+def noop(x: torch.Tensor) -> torch.Tensor:
     """Dummy func. Return input"""
     return x
 
 
 class Noop(nn.Module):
     """Dummy module"""
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
@@ -114,15 +114,15 @@ def __init__(
         nf,
         ks=3,
         stride=1,
-        act=True,
+        act=True,  # pylint: disable=redefined-outer-name
         act_fn=act,
         bn_layer=True,
         bn_1st=True,
         zero_bn=False,
         padding=None,
         bias=False,
         groups=1,
-        **kwargs
+        **kwargs  # pylint: disable=unused-argument
     ):
 
         if padding is None:
@@ -176,7 +176,7 @@ def __init__(self, n_in: int, ks=1, sym=False, use_bias=False):
         self.sym = sym
         self.n_in = n_in
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.sym:  # check ks=3
             # symmetry hack by https://github.com/mgrankin
             c = self.conv.weight.view(self.n_in, self.n_in)
@@ -195,13 +195,14 @@ def forward(self, x):
         return o.view(*size).contiguous()
 
 
-class SEBlock(nn.Module):  # todo: deprecation warning.
-    "se block"
+class SEBlock(nn.Module):
+    """se block"""
+    # first version
     se_layer = nn.Linear
     act_fn = nn.ReLU(inplace=True)
     use_bias = True
 
-    def __init__(self, c, r=16):
+    def __init__(self, c: int, r: int = 16):
         super().__init__()
         ch = max(c // r, 1)
         self.squeeze = nn.AdaptiveAvgPool2d(1)
@@ -216,15 +217,16 @@ def __init__(self, c, r=16):
             )
         )
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
         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 SEBlockConv(nn.Module):  # todo: deprecation warning.
-    "se block with conv on excitation"
+class SEBlockConv(nn.Module):
+    """se block with conv on excitation"""
+    # first version
     se_layer = nn.Conv2d
     act_fn = nn.ReLU(inplace=True)
     use_bias = True