Merge pull request #76 from ayasyrev/typing

v0.3.1 pydantic, typing, tests

Author: ayasyrev

.pylintrc

@@ -21,9 +21,6 @@
 class Flatten(nn.Module):
     """flat x to vector"""
 
-    def __init__(self):
-        super().__init__()
-
     def forward(self, x):
         return x.view(x.size(0), -1)
 
@@ -36,9 +33,6 @@ def noop(x):
 class Noop(nn.Module):
     """Dummy module"""
 
-    def __init__(self):
-        super().__init__()
-
     def forward(self, x):
         return x
 

src/model_constructor/layers.py

@@ -1,6 +1,6 @@
 from collections import OrderedDict
 from functools import partial
-from typing import Any, Callable, List, Optional, Type, Union
+from typing import Any, Callable, List, Type, TypeVar, Union
 
 import torch.nn as nn
 from pydantic import BaseModel, root_validator
@@ -16,6 +16,19 @@
 ]
 
 
+TModelCfg = TypeVar("TModelCfg", bound="ModelCfg")
+
+
+def init_cnn(module: nn.Module):
+    "Init module - kaiming_normal for Conv2d and 0 for biases."
+    if getattr(module, "bias", None) is not None:
+        nn.init.constant_(module.bias, 0)  # type: ignore
+    if isinstance(module, (nn.Conv2d, nn.Linear)):
+        nn.init.kaiming_normal_(module.weight)
+    for layer in module.children():
+        init_cnn(layer)
+
+
 class ResBlock(nn.Module):
     """Universal Resnet block. Basic block if expansion is 1, otherwise is Bottleneck."""
 
@@ -116,86 +129,28 @@ def forward(self, x):
         return self.act_fn(self.convs(x) + identity)
 
 
-class ModelCfg(BaseModel):
-    """Model constructor Config. As default - xresnet18"""
-
-    name: str = "MC"
-    in_chans: int = 3
-    num_classes: int = 1000
-    block: Type[nn.Module] = ResBlock
-    conv_layer: Type[nn.Module] = ConvBnAct
-    block_sizes: List[int] = [64, 128, 256, 512]
-    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)
-    expansion: int = 1
-    groups: int = 1
-    dw: bool = False
-    div_groups: Union[int, None] = None
-    sa: Union[bool, int, Type[nn.Module]] = False
-    se: Union[bool, int, Type[nn.Module]] = False
-    se_module: Union[bool, None] = None
-    se_reduction: Union[int, None] = None
-    bn_1st: bool = True
-    zero_bn: bool = True
-    stem_stride_on: int = 0
-    stem_sizes: List[int] = [32, 32, 64]
-    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[["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
-        extra = "forbid"
-
-    def extra_repr(self) -> str:
-        res = ""
-        for k, v in self.dict().items():
-            if v is not None:
-                res += f"{k}: {v}\n"
-        return res
-
-    def pprint(self) -> None:
-        print(self.extra_repr())
-
-
-def init_cnn(module: nn.Module):
-    "Init module - kaiming_normal for Conv2d and 0 for biases."
-    if getattr(module, "bias", None) is not None:
-        nn.init.constant_(module.bias, 0)  # type: ignore
-    if isinstance(module, (nn.Conv2d, nn.Linear)):
-        nn.init.kaiming_normal_(module.weight)
-    for layer in module.children():
-        init_cnn(layer)
-
-
-def make_stem(self: ModelCfg) -> nn.Sequential:
+def make_stem(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
     stem: List[tuple[str, nn.Module]] = [
-        (f"conv_{i}", self.conv_layer(
-            self.stem_sizes[i],  # type: ignore
-            self.stem_sizes[i + 1],
-            stride=2 if i == self.stem_stride_on else 1,
-            bn_layer=(not self.stem_bn_end)
-            if i == (len(self.stem_sizes) - 2)
+        (f"conv_{i}", cfg.conv_layer(
+            cfg.stem_sizes[i],  # type: ignore
+            cfg.stem_sizes[i + 1],
+            stride=2 if i == cfg.stem_stride_on else 1,
+            bn_layer=(not cfg.stem_bn_end)
+            if i == (len(cfg.stem_sizes) - 2)
             else True,
-            act_fn=self.act_fn,
-            bn_1st=self.bn_1st,
+            act_fn=cfg.act_fn,
+            bn_1st=cfg.bn_1st,
         ),)
-        for i in range(len(self.stem_sizes) - 1)
+        for i in range(len(cfg.stem_sizes) - 1)
     ]
-    if self.stem_pool:
-        stem.append(("stem_pool", self.stem_pool()))
-    if self.stem_bn_end:
-        stem.append(("norm", self.norm(self.stem_sizes[-1])))  # type: ignore
+    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.Sequential(OrderedDict(stem))
 
 
-def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:
+def make_layer(cfg: TModelCfg, layer_num: int) -> nn.Sequential:  # type: ignore
     #  expansion, in_channels, out_channels, blocks, stride, sa):
     # if no pool on stem - stride = 2 for first layer block in body
     stride = 1 if cfg.stem_pool and layer_num == 0 else 2
@@ -231,7 +186,7 @@ def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:
     )
 
 
-def make_body(cfg: ModelCfg) -> nn.Sequential:
+def make_body(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
     return nn.Sequential(
         OrderedDict(
             [
@@ -245,7 +200,7 @@ def make_body(cfg: ModelCfg) -> nn.Sequential:
     )
 
 
-def make_head(cfg: ModelCfg) -> nn.Sequential:
+def make_head(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
     head = [
         ("pool", nn.AdaptiveAvgPool2d(1)),
         ("flat", nn.Flatten()),
@@ -254,22 +209,59 @@ def make_head(cfg: ModelCfg) -> nn.Sequential:
     return nn.Sequential(OrderedDict(head))
 
 
+class ModelCfg(BaseModel):
+    """Model constructor Config. As default - xresnet18"""
+
+    name: str = "MC"
+    in_chans: int = 3
+    num_classes: int = 1000
+    block: Type[nn.Module] = ResBlock
+    conv_layer: Type[nn.Module] = ConvBnAct
+    block_sizes: List[int] = [64, 128, 256, 512]
+    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)
+    expansion: int = 1
+    groups: int = 1
+    dw: bool = False
+    div_groups: Union[int, None] = None
+    sa: Union[bool, int, Type[nn.Module]] = False
+    se: Union[bool, int, Type[nn.Module]] = False
+    se_module: Union[bool, None] = None
+    se_reduction: Union[int, None] = None
+    bn_1st: bool = True
+    zero_bn: bool = True
+    stem_stride_on: int = 0
+    stem_sizes: List[int] = [32, 32, 64]
+    stem_pool: Union[Callable[[], nn.Module], None] = partial(nn.MaxPool2d, kernel_size=3, stride=2, padding=1)
+    stem_bn_end: bool = False
+    init_cnn: Callable[[nn.Module], None] = init_cnn
+    make_stem: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_stem  # type: ignore
+    make_layer: Callable[[TModelCfg, int], Union[nn.Module, nn.Sequential]] = make_layer  # type: ignore
+    make_body: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_body  # type: ignore
+    make_head: Callable[[TModelCfg], Union[nn.Module, nn.Sequential]] = make_head  # type: ignore
+
+    class Config:
+        arbitrary_types_allowed = True
+        extra = "forbid"
+
+    def extra_repr(self) -> str:
+        res = ""
+        for k, v in self.dict().items():
+            if v is not None:
+                res += f"{k}: {v}\n"
+        return res
+
+    def pprint(self) -> None:
+        print(self.extra_repr())
+
+
 class ModelConstructor(ModelCfg):
     """Model constructor. As default - xresnet18"""
 
     @root_validator
-    def post_init(cls, values):
-        if values["init_cnn"] is None:
-            values["init_cnn"] = init_cnn
-        if values["make_stem"] is None:
-            values["make_stem"] = make_stem
-        if values["make_layer"] is None:
-            values["make_layer"] = make_layer
-        if values["make_body"] is None:
-            values["make_body"] = make_body
-        if values["make_head"] is None:
-            values["make_head"] = make_head
-
+    def post_init(cls, values):  # pylint: disable=E0213
         if values["stem_sizes"][0] != values["in_chans"]:
             values["stem_sizes"] = [values["in_chans"]] + values["stem_sizes"]
         if values["se"] and isinstance(values["se"], (bool, int)):  # if se=1 or se=True
@@ -284,15 +276,15 @@ def post_init(cls, values):
 
     @property
     def stem(self):
-        return self.make_stem(self)  # type: ignore
+        return self.make_stem(self)  # pylint: disable=too-many-function-args
 
     @property
     def head(self):
-        return self.make_head(self)  # type: ignore
+        return self.make_head(self)  # pylint: disable=too-many-function-args
 
     @property
     def body(self):
-        return self.make_body(self)  # type: ignore
+        return self.make_body(self)  # pylint: disable=too-many-function-args
 
     @classmethod
     def from_cfg(cls, cfg: ModelCfg):
@@ -302,12 +294,12 @@ def __call__(self):
         model = nn.Sequential(
             OrderedDict([("stem", self.stem), ("body", self.body), ("head", self.head)])
         )
-        self.init_cnn(model)  # type: ignore
+        self.init_cnn(model)  # pylint: disable=too-many-function-args
         model.extra_repr = lambda: f"{self.name}"
         return model
 
     def __repr__(self):
-        se_repr = self.se.__name__ if self.se else "False"
+        se_repr = self.se.__name__ if self.se else "False"  # type: ignore
         return (
             f"{self.name} constructor\n"
             f"  in_chans: {self.in_chans}, num_classes: {self.num_classes}\n"

src/model_constructor/model_constructor.py

@@ -1 +1 @@
-__version__ = "0.3"
+__version__ = "0.3.1"

src/model_constructor/version.py

@@ -1,10 +1,10 @@
-import torch.nn as nn
 from collections import OrderedDict
 
+import torch.nn as nn
+
 from .base_constructor import Net
 from .layers import ConvLayer, Noop, act
 
-
 __all__ = ['DownsampleLayer', 'XResBlock', 'xresnet18', 'xresnet34', 'xresnet50']
 
 

src/model_constructor/xresnet.py

@@ -2,7 +2,7 @@
 # Yet another ResNet.
 
 from collections import OrderedDict
-from typing import Any, Callable, List, Type, Union
+from typing import Callable, List, Type, Union
 
 import torch.nn as nn
 from torch.nn import Mish
@@ -11,13 +11,12 @@
 from .model_constructor import ModelConstructor
 
 __all__ = [
-    'YaResBlock',
+    "YaResBlock",
+    "YaResNet34",
+    "YaResNet50",
 ]
 
 
-# act_fn = nn.ReLU(inplace=True)
-
-
 class YaResBlock(nn.Module):
     '''YaResBlock. Reduce by pool instead of stride 2'''
 
@@ -129,7 +128,7 @@ class YaResNet34(ModelConstructor):
     block: Type[nn.Module] = YaResBlock
     expansion: int = 1
     layers: List[int] = [3, 4, 6, 3]
-    act_fn: nn.Module = Mish()
+    act_fn: Type[nn.Module] = Mish
 
 
 class YaResNet50(YaResNet34):

src/model_constructor/yaresnet.py

@@ -0,0 +1,14 @@
+from torch import nn
+
+
+def value_name(value) -> str:  # pragma: no cover
+    name = getattr(value, "__name__", None)
+    if name is not None:
+        return name
+    if isinstance(value, nn.Module):
+        return value._get_name()  # pylint: disable=W0212
+    return value
+
+
+def ids_fn(key, value):
+    return [f"{key[:2]}_{value_name(v)}" for v in value]