fix base constructor

Author: ayasyrev

model_constructor/base_constructor.py

@@ -1,15 +1,15 @@
-__all__ = ['act_fn', 'Stem', 'DownsampleBlock', 'BasicBlock', 'Bottleneck', 'BasicLayer', 'Body', 'Head', 'init_model',
-           'Net']
-
-# Cell
 import torch.nn as nn
 from collections import OrderedDict
 from .layers import ConvLayer, Noop, Flatten
 
-# Cell
+
+__all__ = ['act_fn', 'Stem', 'DownsampleBlock', 'BasicBlock', 'Bottleneck', 'BasicLayer', 'Body', 'Head', 'init_model',
+           'Net']
+
+
 act_fn = nn.ReLU(inplace=True)
 
-# Cell
+
 class Stem(nn.Sequential):
     """Base stem"""
 
@@ -35,12 +35,12 @@ def __init__(self, c_in=3, stem_sizes=[], stem_out=64,
     def extra_repr(self):
         return f"sizes: {self.sizes}"
 
-# Cell
+
 def DownsampleBlock(conv_layer, ni, nf, ks, stride, act=False, **kwargs):
     '''Base downsample for res-like blocks'''
     return conv_layer(ni, nf, ks, stride, act, **kwargs)
 
-# Cell
+
 class BasicBlock(nn.Module):
     """Basic block (simplified) as in pytorch resnet"""
     def __init__(self, ni, nf, expansion=1, stride=1, zero_bn=False,
@@ -63,7 +63,7 @@ def forward(self, x):
             identity = self.downsample(x)
         return self.act_conn(self.merge(out + identity))
 
-# Cell
+
 class Bottleneck(nn.Module):
     '''Bottlneck block for resnet models'''
     def __init__(self, ni, nh, expansion=4, stride=1, zero_bn=False,
@@ -90,7 +90,7 @@ def forward(self, x):
             identity = self.downsample(x)
         return self.act_conn(self.merge(out + identity))
 
-# Cell
+
 class BasicLayer(nn.Sequential):
     '''Layer from blocks'''
     def __init__(self, block, blocks, ni, nf, expansion, stride, sa=False, **kwargs):
@@ -108,7 +108,7 @@ def __init__(self, block, blocks, ni, nf, expansion, stride, sa=False, **kwargs)
     def extra_repr(self):
         return f'from {self.ni * self.expansion} to {self.nf}, {self.blocks} blocks, expansion {self.expansion}.'
 
-# Cell
+
 class Body(nn.Sequential):
     '''Constructor for body'''
     def __init__(self, block,
@@ -126,7 +126,7 @@ def __init__(self, block,
                   for i in range(num_layers)]
         super().__init__(OrderedDict(layers))
 
-# Cell
+
 class Head(nn.Sequential):
     '''base head'''
     def __init__(self, ni, nf, **kwargs):
@@ -136,14 +136,14 @@ def __init__(self, ni, nf, **kwargs):
              ('fc', nn.Linear(ni, nf)),
              ]))
 
-# Cell
+
 def init_model(model, nonlinearity='leaky_relu'):
     '''Init model'''
     for m in model.modules():
         if isinstance(m, (nn.Conv2d, nn.Linear)):
             nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity=nonlinearity)
 
-# Cell
+
 class Net(nn.Sequential):
     '''Constructor for model'''
     def __init__(self, stem=Stem,
@@ -161,4 +161,4 @@ def __init__(self, stem=Stem,
                            sa=sa, **kwargs)),
              ('head', head(body_out * expansion, num_classes, **kwargs))
              ]))
-        self.init_model(self)
+        self.init_model(self)