move xresnet from model_constructor to xresnet module

Author: ayasyrev

src/model_constructor/model_constructor.py

@@ -34,7 +34,8 @@ def init_cnn(module: nn.Module) -> None:
 
 
 class BasicBlock(nn.Module):
-    """Basic Resnet block."""
+    """Basic Resnet block.
+    Configurable - can use pool to reduce at identity path, change act etc. """
 
     def __init__(
         self,
@@ -117,7 +118,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
 
 
 class BottleneckBlock(nn.Module):
-    """Bottleneck Resnet block."""
+    """Bottleneck Resnet block.
+    Configurable - can use pool to reduce at identity path, change act etc. """
 
     def __init__(
         self,
@@ -211,21 +213,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
 
 
 def make_stem(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
-    """Create xResnet stem -> 3 conv 3*3 instead 1 conv 7*7"""
-    len_stem = len(cfg.stem_sizes)
+    """Create Resnet stem."""
     stem: ListStrMod = [
         (
-            f"conv_{i}",
+            "conv_1",
             cfg.conv_layer(
-                cfg.stem_sizes[i - 1] if i else cfg.in_chans,  # type: ignore
-                cfg.stem_sizes[i],
-                stride=2 if i == cfg.stem_stride_on else 1,
-                bn_layer=(not cfg.stem_bn_end) if i == (len_stem - 1) else True,
+                cfg.in_chans,  # type: ignore
+                cfg.stem_sizes[-1],
+                kernel_size=7,
+                stride=2,
+                padding=3,
+                bn_layer=not cfg.stem_bn_end,
                 act_fn=cfg.act_fn,
                 bn_1st=cfg.bn_1st,
             ),
         )
-        for i in range(len_stem)
     ]
     if cfg.stem_pool:
         stem.append(("stem_pool", cfg.stem_pool()))
@@ -295,9 +297,7 @@ class ModelCfg(BaseModel):
     layers: list[int] = [2, 2, 2, 2]
     norm: type[nn.Module] = nn.BatchNorm2d
     act_fn: type[nn.Module] = nn.ReLU
-    pool: Callable[[Any], nn.Module] = partial(
-        nn.AvgPool2d, kernel_size=2, ceil_mode=True
-    )
+    pool: Optional[Callable[[Any], nn.Module]] = None
     expansion: int = 1
     groups: int = 1
     dw: bool = False
@@ -309,7 +309,7 @@ class ModelCfg(BaseModel):
     bn_1st: bool = True
     zero_bn: bool = True
     stem_stride_on: int = 0
-    stem_sizes: list[int] = [32, 32, 64]
+    stem_sizes: list[int] = [64]
     stem_pool: Union[Callable[[], nn.Module], None] = partial(
         nn.MaxPool2d, kernel_size=3, stride=2, padding=1
     )
@@ -368,7 +368,7 @@ def print_changed(self) -> None:
 
 
 class ModelConstructor(ModelCfg):
-    """Model constructor. As default - xresnet18"""
+    """Model constructor. As default - resnet18"""
 
     @validator("se")
     def set_se(  # pylint: disable=no-self-argument
@@ -446,10 +446,10 @@ def __repr__(self) -> str:
         )
 
 
-class XResNet34(ModelConstructor):
+class ResNet34(ModelConstructor):
     layers: list[int] = [3, 4, 6, 3]
 
 
-class XResNet50(XResNet34):
+class ResNet50(ResNet34):
     block: type[nn.Module] = BottleneckBlock
     block_sizes: list[int] = [256, 512, 1024, 2048]

src/model_constructor/xresnet.py

@@ -0,0 +1,55 @@
+from functools import partial
+from typing import Any, Callable, Optional, Union
+
+from torch import nn
+
+from .helpers import nn_seq
+from .model_constructor import (BottleneckBlock, ListStrMod, ModelCfg,
+                                ModelConstructor)
+
+__all__ = [
+    "XResNet",
+    "XResNet34",
+    "XResNet50",
+]
+
+
+def xresnet_stem(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
+    """Create xResnet stem -> 3 conv 3*3 instead 1 conv 7*7"""
+    len_stem = len(cfg.stem_sizes)
+    stem: ListStrMod = [
+        (
+            f"conv_{i}",
+            cfg.conv_layer(
+                cfg.stem_sizes[i - 1] if i else cfg.in_chans,  # type: ignore
+                cfg.stem_sizes[i],
+                stride=2 if i == cfg.stem_stride_on else 1,
+                bn_layer=(not cfg.stem_bn_end) if i == (len_stem - 1) else True,
+                act_fn=cfg.act_fn,
+                bn_1st=cfg.bn_1st,
+            ),
+        )
+        for i in range(len_stem)
+    ]
+    if cfg.stem_pool:
+        stem.append(("stem_pool", cfg.stem_pool()))
+    if cfg.stem_bn_end:
+        stem.append(("norm", cfg.norm(cfg.stem_sizes[-1])))  # type: ignore
+    return nn_seq(stem)
+
+
+class XResNet(ModelConstructor):
+    make_stem: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = xresnet_stem
+    stem_sizes: list[int] = [32, 32, 64]
+    pool: Optional[Callable[[Any], nn.Module]] = partial(
+        nn.AvgPool2d, kernel_size=2, ceil_mode=True
+    )
+
+
+class XResNet34(XResNet):
+    layers: list[int] = [3, 4, 6, 3]
+
+
+class XResNet50(XResNet34):
+    block: type[nn.Module] = BottleneckBlock
+    block_sizes: list[int] = [256, 512, 1024, 2048]

src/model_constructor/yaresnet.py

@@ -1,8 +1,8 @@
 # YaResBlock - former NewResBlock.
 # Yet another ResNet.
 
-from collections import OrderedDict
-from typing import Callable, Union
+from functools import partial
+from typing import Any, Callable, Optional, Union
 
 import torch
 from torch import nn
@@ -11,7 +11,9 @@
 from model_constructor.helpers import nn_seq
 
 from .layers import ConvBnAct, get_act
-from .model_constructor import ListStrMod, ModelConstructor
+from .model_constructor import ListStrMod, ModelConstructor, ModelCfg
+from .xresnet import xresnet_stem
+
 
 __all__ = [
     "YaBasicBlock",
@@ -202,9 +204,13 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
 
 
 class YaResNet(ModelConstructor):
-    block: type[nn.Module] = YaBasicBlock
+    make_stem: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = xresnet_stem
     stem_sizes: list[int] = [3, 32, 64, 64]
+    block: type[nn.Module] = YaBasicBlock
     act_fn: type[nn.Module] = Mish
+    pool: Optional[Callable[[Any], nn.Module]] = partial(
+        nn.AvgPool2d, kernel_size=2, ceil_mode=True
+    )
 
 
 class YaResNet34(YaResNet):