act func inplace

Author: ayasyrev

src/model_constructor/layers.py

@@ -10,7 +10,7 @@
     "noop",
     "Noop",
     "ConvLayer",
-    "act_fn",
+    "act",
     "conv1d",
     "SimpleSelfAttention",
     "SEBlock",
@@ -43,7 +43,16 @@ def forward(self, x):
         return x
 
 
-act_fn = nn.ReLU(inplace=True)
+act = nn.ReLU(inplace=True)
+
+
+def get_act(act_fn: Type[nn.Module], inplace: bool = True) -> nn.Module:
+    """Return obj of act_fn, inplace if possible."""
+    try:
+        res = act_fn(inplace=inplace)  # type: ignore
+    except TypeError:
+        res = act_fn()
+    return res
 
 
 class ConvBnAct(nn.Sequential):
@@ -95,7 +104,7 @@ def __init__(
                 act_position = 1
             else:
                 act_position = len(layers)
-            layers.insert(act_position, ("act_fn", act_fn(inplace=True)))  # type: ignore
+            layers.insert(act_position, ("act_fn", get_act(act_fn)))  # type: ignore
         super().__init__(OrderedDict(layers))
 
 
@@ -112,7 +121,7 @@ def __init__(
         ks=3,
         stride=1,
         act=True,
-        act_fn=act_fn,
+        act_fn=act,
         bn_layer=True,
         bn_1st=True,
         zero_bn=False,

src/model_constructor/model_constructor.py

@@ -5,7 +5,7 @@
 import torch.nn as nn
 from pydantic import BaseModel, root_validator
 
-from .layers import ConvBnAct, SEModule, SimpleSelfAttention
+from .layers import ConvBnAct, SEModule, SimpleSelfAttention, get_act
 
 __all__ = [
     "init_cnn",
@@ -32,7 +32,7 @@ def __init__(
         groups: int = 1,
         dw: bool = False,
         div_groups: Union[None, int] = None,
-        pool: Union[Callable[[Any], nn.Module], None] = None,
+        pool: Union[Callable[[], nn.Module], None] = None,
         se: Union[nn.Module, None] = None,
         sa: Union[nn.Module, None] = None,
     ):
@@ -109,7 +109,7 @@ def __init__(
             self.id_conv = nn.Sequential(OrderedDict(id_layers))
         else:
             self.id_conv = None
-        self.act_fn = act_fn(inplace=True)  # type: ignore
+        self.act_fn = get_act(act_fn)  # type: ignore
 
     def forward(self, x):
         identity = self.id_conv(x) if self.id_conv is not None else x
@@ -141,13 +141,13 @@ class ModelCfg(BaseModel):
     zero_bn: bool = True
     stem_stride_on: int = 0
     stem_sizes: List[int] = [32, 32, 64]
-    stem_pool: Union[Callable[[Any], nn.Module], None] = partial(nn.MaxPool2d, kernel_size=3, stride=2, padding=1)
+    stem_pool: Union[Callable[[], nn.Module], None] = partial(nn.MaxPool2d, kernel_size=3, stride=2, padding=1)
     stem_bn_end: bool = False
     init_cnn: Optional[Callable[[nn.Module], None]] = None
-    make_stem: Optional[Callable] = None
-    make_layer: Optional[Callable] = None
-    make_body: Optional[Callable] = None
-    make_head: Optional[Callable] = None
+    make_stem: Optional[Callable[["ModelCfg"], nn.Module]] = None
+    make_layer: Optional[Callable[["ModelCfg"], nn.Module]] = None
+    make_body: Optional[Callable[["ModelCfg"], nn.Module]] = None
+    make_head: Optional[Callable[["ModelCfg"], nn.Module]] = None
 
     class Config:
         arbitrary_types_allowed = True

src/model_constructor/twist.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from .layers import ConvLayer, noop, act_fn, SimpleSelfAttention
+from .layers import ConvLayer, noop, act, SimpleSelfAttention
 
 import torch
 import torch.nn as nn
@@ -112,7 +112,7 @@ class NewResBlockTwist(nn.Module):
     Now YaResBlock.'''
 
     def __init__(self, expansion, ni, nh, stride=1,
-                 conv_layer=ConvLayer, act_fn=act_fn, bn_1st=True,
+                 conv_layer=ConvLayer, act_fn=act, bn_1st=True,
                  pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False, zero_bn=True, **kwargs):
         super().__init__()
         nf, ni = nh * expansion, ni * expansion
@@ -139,7 +139,7 @@ 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,
+                 conv_layer=ConvLayer, act_fn=act, zero_bn=True, bn_1st=True,
                  pool=nn.AvgPool2d(2, ceil_mode=True), sa=False, sym=False, **kwargs):
         super().__init__()
         nf, ni = nh * expansion, ni * expansion

src/model_constructor/xresnet.py

@@ -2,7 +2,7 @@
 from collections import OrderedDict
 
 from .base_constructor import Net
-from .layers import ConvLayer, Noop, act_fn
+from .layers import ConvLayer, Noop, act
 
 
 __all__ = ['DownsampleLayer', 'XResBlock', 'xresnet18', 'xresnet34', 'xresnet50']
@@ -25,7 +25,7 @@ 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):
+                 conv_layer=ConvLayer, act_fn=act, **kwargs):
         super().__init__()
         nf, ni = nh * expansion, ni * expansion
         layers = [('conv_0', conv_layer(ni, nh, 3, stride=stride, act_fn=act_fn, **kwargs)),

src/model_constructor/yaresnet.py

@@ -7,7 +7,7 @@
 import torch.nn as nn
 from torch.nn import Mish
 
-from .layers import ConvBnAct
+from .layers import ConvBnAct, get_act
 from .model_constructor import ModelConstructor
 
 __all__ = [
@@ -34,7 +34,7 @@ def __init__(
         groups: int = 1,
         dw: bool = False,
         div_groups: Union[None, int] = None,
-        pool: Union[Callable[[Any], nn.Module], None] = None,
+        pool: Union[Callable[[], nn.Module], None] = None,
         se: Union[nn.Module, None] = None,
         sa: Union[nn.Module, None] = None,
     ):
@@ -115,7 +115,7 @@ def __init__(
             )
         else:
             self.id_conv = None
-        self.merge = act_fn()
+        self.merge = get_act(act_fn)
 
     def forward(self, x):
         if self.reduce: