fix twist

Author: ayasyrev

model_constructor/twist.py

@@ -1,17 +1,15 @@
-__all__ = ['ConvTwist', 'ConvLayerTwist', 'NewResBlockTwist', 'ResBlockTwist']
-
-# Cell
-# from functools import partial
 from collections import OrderedDict
 from .layers import ConvLayer, noop, act_fn, SimpleSelfAttention
 
-# Cell
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
 
-# Cell
+
+__all__ = ['ConvTwist', 'ConvLayerTwist', 'NewResBlockTwist', 'ResBlockTwist']
+
+
 class ConvTwist(nn.Module):
     '''Replacement for Conv2d (kernelsize 3x3)'''
     permute = True
@@ -102,12 +100,15 @@ def forward(self, inpt):
     def extra_repr(self):
         return f"twist: {self.twist}, permute: {self.permute}, same: {self.same}, groups: {self.groups}"
 
-# Cell
+
 class ConvLayerTwist(ConvLayer):  # replace Conv2d by Twist
+    '''Conv layer with ConvTwist'''
     Conv2d = ConvTwist
 
-# Cell
+
 class NewResBlockTwist(nn.Module):
+    '''Resnet block with ConvTwist'''
+
     def __init__(self, expansion, ni, nh, stride=1,
                  conv_layer=ConvLayer, act_fn=act_fn, bn_1st=True,
                  pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False, zero_bn=True, **kvargs):
@@ -131,8 +132,10 @@ def forward(self, x):
         o = self.reduce(x)
         return self.merge(self.convs(o) + self.idconv(o))
 
-# Cell
+
 class ResBlockTwist(nn.Module):
+    '''Resnet block with ConvTwist'''
+
     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, **kvargs):
@@ -153,4 +156,4 @@ def __init__(self, expansion, ni, nh, stride=1,
         self.act_fn = act_fn
 
     def forward(self, x):
-        return self.act_fn(self.convs(x) + self.idconv(self.pool(x)))
+        return self.act_fn(self.convs(x) + self.idconv(self.pool(x)))