Merge pull request #92 from ayasyrev/dev_0.3.3

0.3.3

Author: ayasyrev

docs/overrides/partials/copyright.html

@@ -1,6 +1,3 @@
-{#-
-  This file was automatically generated - do not edit
--#}
 <div class="md-copyright">
   {% if config.copyright %}
     <div class="md-copyright__highlight">

mkdocs.yaml

@@ -24,3 +24,4 @@ extra:
   analytics:
     provider: google
     property: G-0F3FK713C2
+copyright: Copyright © 2020-2023 Andrei Yasyrev.

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/convmixer.py

@@ -3,37 +3,49 @@
 # Adopted from https://github.com/tmp-iclr/convmixer
 # Home for convmixer: https://github.com/locuslab/convmixer
 from collections import OrderedDict
-from typing import Callable
+from typing import Callable, List, Optional, Union
+
 import torch.nn as nn
+from torch import TensorType
 
 
 class Residual(nn.Module):
-    def __init__(self, fn):
+    def __init__(self, fn: Callable[[TensorType], TensorType]):
         super().__init__()
         self.fn = fn
 
-    def forward(self, x):
+    def forward(self, x: TensorType) -> TensorType:
         return self.fn(x) + x
 
 
 # As original version, act_fn as argument.
-def ConvMixerOriginal(dim, depth,
-                      kernel_size=9, patch_size=7, n_classes=1000,
-                      act_fn=nn.GELU()):
+def ConvMixerOriginal(
+    dim: int,
+    depth: int,
+    kernel_size: int = 9,
+    patch_size: int = 7,
+    n_classes: int = 1000,
+    act_fn: nn.Module = nn.GELU(),
+):
     return nn.Sequential(
         nn.Conv2d(3, dim, kernel_size=patch_size, stride=patch_size),
         act_fn,
         nn.BatchNorm2d(dim),
-        *[nn.Sequential(
-            Residual(nn.Sequential(
-                nn.Conv2d(dim, dim, kernel_size, groups=dim, padding="same"),
+        *[
+            nn.Sequential(
+                Residual(
+                    nn.Sequential(
+                        nn.Conv2d(dim, dim, kernel_size, groups=dim, padding="same"),
+                        act_fn,
+                        nn.BatchNorm2d(dim),
+                    )
+                ),
+                nn.Conv2d(dim, dim, kernel_size=1),
                 act_fn,
-                nn.BatchNorm2d(dim)
-            )),
-            nn.Conv2d(dim, dim, kernel_size=1),
-            act_fn,
-            nn.BatchNorm2d(dim)
-        ) for i in range(depth)],
+                nn.BatchNorm2d(dim),
+            )
+            for _i in range(depth)
+        ],
         nn.AdaptiveAvgPool2d((1, 1)),
         nn.Flatten(),
         nn.Linear(dim, n_classes)
@@ -43,15 +55,35 @@ def ConvMixerOriginal(dim, depth,
 class ConvLayer(nn.Sequential):
     """Basic conv layers block"""
 
-    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
-                 act_fn=nn.GELU(), padding=0, groups=1,
-                 bn_1st=False, pre_act=False):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: Union[int, tuple[int, int]],
+        stride: int = 1,
+        act_fn: nn.Module = nn.GELU(),
+        padding: Union[int, str] = 0,
+        groups: int = 1,
+        bn_1st: bool = False,
+        pre_act: bool = False,
+    ):
 
-        conv_layer = [('conv', nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride,
-                                         padding=padding, groups=groups))]
-        act_bn = [
-            ('act_fn', act_fn),
-            ('bn', nn.BatchNorm2d(out_channels))
+        conv_layer: List[tuple[str, nn.Module]] = [
+            (
+                "conv",
+                nn.Conv2d(
+                    in_channels,
+                    out_channels,
+                    kernel_size,
+                    stride=stride,
+                    padding=padding,
+                    groups=groups,
+                ),
+            )
+        ]
+        act_bn: List[tuple[str, nn.Module]] = [
+            ("act_fn", act_fn),
+            ("bn", nn.BatchNorm2d(out_channels)),
         ]
         if bn_1st:
             act_bn.reverse()
@@ -64,45 +96,79 @@ def __init__(self, in_channels, out_channels, kernel_size, stride=1,
 
 
 class ConvMixer(nn.Sequential):
-
-    def __init__(self, dim: int, depth: int,
-                 kernel_size: int = 9, patch_size: int = 7, n_classes: int = 1000,
-                 act_fn: nn.Module = nn.GELU(),
-                 stem: nn.Module = None,
-                 bn_1st: bool = False, pre_act: bool = False,
-                 init_func: Callable = None):
+    def __init__(
+        self,
+        dim: int,
+        depth: int,
+        kernel_size: int = 9,
+        patch_size: int = 7,
+        n_classes: int = 1000,
+        act_fn: nn.Module = nn.GELU(),
+        stem: Optional[nn.Module] = None,
+        in_chans: int = 3,
+        bn_1st: bool = False,
+        pre_act: bool = False,
+        init_func: Optional[Callable[[nn.Module], None]] = None,
+    ):
         """ConvMixer constructor.
         Adopted from https://github.com/tmp-iclr/convmixer
 
         Args:
-            dim (int): Dimention of model.
+            dim (int): Dimension of model.
             depth (int): Depth of model.
             kernel_size (int, optional): Kernel size. Defaults to 9.
             patch_size (int, optional): Patch size. Defaults to 7.
             n_classes (int, optional): Number of classes. Defaults to 1000.
             act_fn (nn.Module, optional): Activation function. Defaults to nn.GELU().
             stem (nn.Module, optional): You can path different first layer..
-            stem_ks (int, optional): If stem_ch not 0 - kernel size for adittional layer. Defaults to 1.
-            bn_1st (bool, optional): If True - BatchNorm befor activation function. Defaults to False.
-            pre_act (bool, optional): If True - activatin function befor convolution layer. Defaults to False.
+            stem_ks (int, optional): If stem_ch not 0 - kernel size for additional layer. Defaults to 1.
+            bn_1st (bool, optional): If True - BatchNorm before activation function. Defaults to False.
+            pre_act (bool, optional): If True - activation function before convolution layer. Defaults to False.
             init_func (Callable, optional): External function for init model.
 
         """
         if pre_act:
             bn_1st = False
         if stem is None:
-            stem = ConvLayer(3, dim, kernel_size=patch_size, stride=patch_size, act_fn=act_fn, bn_1st=bn_1st)
+            stem = ConvLayer(
+                in_chans,
+                dim,
+                kernel_size=patch_size,
+                stride=patch_size,
+                act_fn=act_fn,
+                bn_1st=bn_1st,
+            )
 
         super().__init__(
             stem,
-            *[nn.Sequential(
-                Residual(
-                    ConvLayer(dim, dim, kernel_size, act_fn=act_fn,
-                              groups=dim, padding="same", bn_1st=bn_1st, pre_act=pre_act)),
-                ConvLayer(dim, dim, kernel_size=1, act_fn=act_fn, bn_1st=bn_1st, pre_act=pre_act))
-              for i in range(depth)],
+            *[
+                nn.Sequential(
+                    Residual(
+                        ConvLayer(
+                            dim,
+                            dim,
+                            kernel_size,
+                            act_fn=act_fn,
+                            groups=dim,
+                            padding="same",
+                            bn_1st=bn_1st,
+                            pre_act=pre_act,
+                        )
+                    ),
+                    ConvLayer(
+                        dim,
+                        dim,
+                        kernel_size=1,
+                        act_fn=act_fn,
+                        bn_1st=bn_1st,
+                        pre_act=pre_act,
+                    ),
+                )
+                for _ in range(depth)
+            ],
             nn.AdaptiveAvgPool2d((1, 1)),
             nn.Flatten(),
-            nn.Linear(dim, n_classes))
+            nn.Linear(dim, n_classes)
+        )
         if init_func is not None:  # pragma: no cover
             init_func(self)

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))  #