From 90991a01f17db45ecff087a6c33c400ad6039710 Mon Sep 17 00:00:00 2001 From: ayasyrev Date: Thu, 30 Sep 2021 12:18:21 +0300 Subject: [PATCH] rename Net to ModelConstructor --- model_constructor/__init__.py | 5 +- model_constructor/model_constructor.py | 167 +++++++++++++++++++++++++ model_constructor/net.py | 12 +- 3 files changed, 179 insertions(+), 5 deletions(-) create mode 100644 model_constructor/model_constructor.py diff --git a/model_constructor/__init__.py b/model_constructor/__init__.py index f1380ee..b2ab8e6 100644 --- a/model_constructor/__init__.py +++ b/model_constructor/__init__.py @@ -1 +1,4 @@ -__version__ = "0.1.7" +from model_constructor.model_constructor import ModelConstructor # noqa F401 + + +__version__ = "0.1.8" diff --git a/model_constructor/model_constructor.py b/model_constructor/model_constructor.py new file mode 100644 index 0000000..ead16f0 --- /dev/null +++ b/model_constructor/model_constructor.py @@ -0,0 +1,167 @@ +from collections import OrderedDict +from functools import partial + +import torch.nn as nn + +from .layers import ConvLayer, Flatten, SEBlock, SimpleSelfAttention, noop + + +__all__ = ['init_cnn', 'act_fn', 'ResBlock', 'ModelConstructor', 'xresnet34', 'xresnet50'] + + +act_fn = nn.ReLU(inplace=True) + + +def init_cnn(module: nn.Module): + "Init module - kaiming_normal for Conv2d and 0 for biases." + if getattr(module, 'bias', None) is not None: + nn.init.constant_(module.bias, 0) + if isinstance(module, (nn.Conv2d, nn.Linear)): + nn.init.kaiming_normal_(module.weight) + for layer in module.children(): + init_cnn(layer) + + +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): + super().__init__() + nf, ni = nh * expansion, ni * expansion + if div_groups is not None: # check if grops != 1 and div_groups + groups = int(nh / div_groups) + if expansion == 1: + layers = [("conv_0", conv_layer(ni, nh, 3, stride=stride, act_fn=act_fn, bn_1st=bn_1st, + groups=nh if dw else groups)), + ("conv_1", conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False, bn_1st=bn_1st)) + ] + else: + layers = [("conv_0", conv_layer(ni, nh, 1, act_fn=act_fn, bn_1st=bn_1st)), + ("conv_1", conv_layer(nh, nh, 3, stride=stride, act_fn=act_fn, bn_1st=bn_1st, + groups=nh if dw else groups)), + ("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))) + if sa: + layers.append(('sa', SimpleSelfAttention(nf, ks=1, sym=sym))) + self.convs = nn.Sequential(OrderedDict(layers)) + self.pool = noop if stride == 1 else pool + self.idconv = noop if ni == nf else conv_layer(ni, nf, 1, act=False) + self.act_fn = act_fn + + def forward(self, x): + return self.act_fn(self.convs(x) + self.idconv(self.pool(x))) + + +def _make_stem(self): + stem = [(f"conv_{i}", self.conv_layer(self.stem_sizes[i], self.stem_sizes[i + 1], + stride=2 if i == self.stem_stride_on else 1, + bn_layer=(not self.stem_bn_end) if i == (len(self.stem_sizes) - 2) else True, + act_fn=self.act_fn, bn_1st=self.bn_1st)) + for i in range(len(self.stem_sizes) - 1)] + stem.append(('stem_pool', self.stem_pool)) + if self.stem_bn_end: + stem.append(('norm', self.norm(self.stem_sizes[-1]))) + return nn.Sequential(OrderedDict(stem)) + + +def _make_layer(self, expansion, ni, nf, blocks, stride, sa): + layers = [(f"bl_{i}", self.block(expansion, ni if i == 0 else nf, nf, + stride if i == 0 else 1, sa=sa if i == blocks - 1 else False, + 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)) + for i in range(blocks)] + return nn.Sequential(OrderedDict(layers)) + + +def _make_body(self): + blocks = [(f"l_{i}", self._make_layer(self, self.expansion, + ni=self.block_sizes[i], nf=self.block_sizes[i + 1], + blocks=l, stride=1 if i == 0 else 2, + sa=self.sa if i == 0 else False)) + for i, l in enumerate(self.layers)] + return nn.Sequential(OrderedDict(blocks)) + + +def _make_head(self): + head = [('pool', nn.AdaptiveAvgPool2d(1)), + ('flat', Flatten()), + ('fc', nn.Linear(self.block_sizes[-1] * self.expansion, self.c_out))] + return nn.Sequential(OrderedDict(head)) + + +class ModelConstructor(): + """Model constructor. As default - xresnet18""" + def __init__(self, name='MC', c_in=3, c_out=1000, + block=ResBlock, conv_layer=ConvLayer, + block_sizes=[64, 128, 256, 512], layers=[2, 2, 2, 2], + norm=nn.BatchNorm2d, + 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, + bn_1st=True, + zero_bn=True, + stem_stride_on=0, + stem_sizes=[32, 32, 64], + stem_pool=nn.MaxPool2d(kernel_size=3, stride=2, padding=1), + stem_bn_end=False, + _init_cnn=init_cnn, + _make_stem=_make_stem, + _make_layer=_make_layer, + _make_body=_make_body, + _make_head=_make_head, + ): + super().__init__() + + params = locals() + del params['self'] + self.__dict__ = params + self._block_sizes = params['block_sizes'] + if self.stem_sizes[0] != self.c_in: + self.stem_sizes = [self.c_in] + self.stem_sizes + + @property + def block_sizes(self): + return [self.stem_sizes[-1] // self.expansion] + self._block_sizes + [256] * (len(self.layers) - 4) + + @property + def stem(self): + return self._make_stem(self) + + @property + def head(self): + return self._make_head(self) + + @property + def body(self): + return self._make_body(self) + + def __call__(self): + model = nn.Sequential(OrderedDict([ + ('stem', self.stem), + ('body', self.body), + ('head', self.head)])) + self._init_cnn(model) + model.extra_repr = lambda: f"model {self.name}" + return model + + def __repr__(self): + return (f"{self.name} constructor\n" + f" c_in: {self.c_in}, c_out: {self.c_out}\n" + f" expansion: {self.expansion}, groups: {self.groups}, dw: {self.dw}, div_groups: {self.div_groups}\n" + f" sa: {self.sa}, se: {self.se}\n" + f" stem sizes: {self.stem_sizes}, stride on {self.stem_stride_on}\n" + f" body sizes {self._block_sizes}\n" + f" layers: {self.layers}") + + +xresnet34 = partial(ModelConstructor, name='xresnet34', expansion=1, layers=[3, 4, 6, 3]) +xresnet50 = partial(ModelConstructor, name='xresnet34', expansion=4, layers=[3, 4, 6, 3]) diff --git a/model_constructor/net.py b/model_constructor/net.py index 6025739..1daaaa9 100644 --- a/model_constructor/net.py +++ b/model_constructor/net.py @@ -61,8 +61,10 @@ def forward(self, x): # NewResBlock now is YaResBlock - Yet Another ResNet Block! It is now at model_constructor.yaresnet. -class NewResBlock(nn.Module): - '''YaResnet block''' +class NewResBlock(nn.Module): # todo: deprecation worning. + '''YaResnet block. + This is first impl, deprecated, use yaresnet module. + ''' se_block = SEBlock def __init__(self, expansion, ni, nh, stride=1, @@ -137,8 +139,10 @@ def _make_head(self): return nn.Sequential(OrderedDict(head)) -class Net(): - """Model constructor. As default - xresnet18""" +class Net(): # todo: deprecation worning. + """Model constructor. As default - xresnet18. + First version, still here for compatibility. Use ModelConstructor instead. + """ def __init__(self, name='Net', c_in=3, c_out=1000, block=ResBlock, conv_layer=ConvLayer, block_sizes=[64, 128, 256, 512], layers=[2, 2, 2, 2],