yaresnet

Author: ayasyrev

src/model_constructor/yaresnet.py

@@ -1,15 +1,22 @@
 # YaResBlock - former NewResBlock.
 # Yet another ResNet.
 
-import torch.nn as nn
-from functools import partial
 from collections import OrderedDict
-from .layers import ConvBnAct
-from .net import Net
+from functools import partial
+from typing import Union
+
+import torch.nn as nn
 from torch.nn import Mish
 
+from .layers import ConvBnAct
+from .net import Net
 
-__all__ = ['YaResBlock', 'yaresnet_parameters', 'yaresnet34', 'yaresnet50']
+__all__ = [
+    'YaResBlock',
+    # 'yaresnet_parameters',
+    # 'yaresnet34',
+    # 'yaresnet50',
+]
 
 
 act_fn = nn.ReLU(inplace=True)
@@ -18,16 +25,29 @@
 class YaResBlock(nn.Module):
     '''YaResBlock. Reduce by pool instead of stride 2'''
 
-    def __init__(self, expansion, in_channels, mid_channels, stride=1,
-                 conv_layer=ConvBnAct, act_fn=act_fn, zero_bn=True, bn_1st=True,
-                 groups=1, dw=False, div_groups=None,
-                 pool=None,
-                 se=None, sa=None,
-                 ):
+    def __init__(
+        self,
+        expansion: int,
+        in_channels: int,
+        mid_channels: int,
+        stride: int = 1,
+        conv_layer=ConvBnAct,
+        act_fn: nn.Module = act_fn,
+        zero_bn: bool = True,
+        bn_1st: bool = True,
+        groups: int = 1,
+        dw: bool = False,
+        div_groups: Union[None, int] = None,
+        pool: Union[nn.Module, None] = None,
+        se: Union[nn.Module, None] = None,
+        sa: Union[nn.Module, None] = None,
+    ):
         super().__init__()
+        # pool defined at ModelConstructor.
         out_channels, in_channels = mid_channels * expansion, in_channels * expansion
         if div_groups is not None:  # check if groups != 1 and div_groups
             groups = int(mid_channels / div_groups)
+
         if stride != 1:
             if pool is None:
                 self.reduce = conv_layer(in_channels, in_channels, 1, stride=2)
@@ -36,23 +56,69 @@ def __init__(self, expansion, in_channels, mid_channels, stride=1,
                 self.reduce = pool
         else:
             self.reduce = None
-        layers = [("conv_0", conv_layer(in_channels, mid_channels, 3, stride=1,
-                                        act_fn=act_fn, bn_1st=bn_1st, groups=in_channels if dw else groups)),
-                  ("conv_1", conv_layer(mid_channels, out_channels, 3, zero_bn=zero_bn,
-                                        act_fn=False, bn_1st=bn_1st, groups=mid_channels if dw else groups))
-                  ] if expansion == 1 else [
-                      ("conv_0", conv_layer(in_channels, mid_channels, 1, act_fn=act_fn, bn_1st=bn_1st)),
-                      ("conv_1", conv_layer(mid_channels, mid_channels, 3, stride=1, act_fn=act_fn, bn_1st=bn_1st,
-                                            groups=mid_channels if dw else groups)),
-                      ("conv_2", conv_layer(
-                          mid_channels, out_channels, 1, zero_bn=zero_bn, act_fn=False, bn_1st=bn_1st))
-        ]
+        if expansion == 1:
+            layers = [
+                ("conv_0", conv_layer(
+                    in_channels,
+                    mid_channels,
+                    3,
+                    stride=1,
+                    act_fn=act_fn,
+                    bn_1st=bn_1st,
+                    groups=in_channels if dw else groups,
+                ),),
+                ("conv_1", conv_layer(
+                    mid_channels,
+                    out_channels,
+                    3,
+                    zero_bn=zero_bn,
+                    act_fn=False,
+                    bn_1st=bn_1st,
+                    groups=mid_channels if dw else groups,
+                ),),
+            ]
+        else:
+            layers = [
+                ("conv_0", conv_layer(
+                    in_channels,
+                    mid_channels,
+                    1,
+                    act_fn=act_fn,
+                    bn_1st=bn_1st,
+                ),),
+                ("conv_1", conv_layer(
+                    mid_channels,
+                    mid_channels,
+                    3,
+                    stride=1,
+                    act_fn=act_fn,
+                    bn_1st=bn_1st,
+                    groups=mid_channels if dw else groups,
+                ),),
+                ("conv_2", conv_layer(
+                    mid_channels,
+                    out_channels,
+                    1,
+                    zero_bn=zero_bn,
+                    act_fn=False,
+                    bn_1st=bn_1st,
+                ),),  # noqa E501
+            ]
         if se:
-            layers.append(('se', se(out_channels)))
+            layers.append(("se", se(out_channels)))
         if sa:
-            layers.append(('sa', sa(out_channels)))
+            layers.append(("sa", sa(out_channels)))
         self.convs = nn.Sequential(OrderedDict(layers))
-        self.id_conv = None if in_channels == out_channels else conv_layer(in_channels, out_channels, 1, act_fn=False)
+        if in_channels != out_channels:
+            self.id_conv = conv_layer(
+                in_channels,
+                out_channels,
+                1,
+                stride=1,
+                act_fn=False,
+            )
+        else:
+            self.id_conv = None
         self.merge = act_fn
 
     def forward(self, x):
@@ -62,6 +128,6 @@ def forward(self, x):
         return self.merge(self.convs(x) + identity)
 
 
-yaresnet_parameters = {'block': YaResBlock, 'stem_sizes': [3, 32, 64, 64], 'act_fn': Mish(), 'stem_stride_on': 1}
-yaresnet34 = partial(Net, name='YaResnet34', expansion=1, layers=[3, 4, 6, 3], **yaresnet_parameters)
-yaresnet50 = partial(Net, name='YaResnet50', expansion=4, layers=[3, 4, 6, 3], **yaresnet_parameters)
+# yaresnet_parameters = {'block': YaResBlock, 'stem_sizes': [3, 32, 64, 64], 'act_fn': Mish(), 'stem_stride_on': 1}
+# yaresnet34 = partial(Net, name='YaResnet34', expansion=1, layers=[3, 4, 6, 3], **yaresnet_parameters)
+# yaresnet50 = partial(Net, name='YaResnet50', expansion=4, layers=[3, 4, 6, 3], **yaresnet_parameters)