Merge pull request #82 from ayasyrev/model_name

Model name refactored

Author: ayasyrev

src/model_constructor/model_constructor.py

@@ -1,6 +1,6 @@
 from collections import OrderedDict
 from functools import partial
-from typing import Any, Callable, List, Type, TypeVar, Union
+from typing import Any, Callable, List, Optional, Type, TypeVar, Union
 
 import torch.nn as nn
 from pydantic import BaseModel, root_validator
@@ -212,7 +212,7 @@ def make_head(cfg: TModelCfg) -> nn.Sequential:  # type: ignore
 class ModelCfg(BaseModel):
     """Model constructor Config. As default - xresnet18"""
 
-    name: str = "MC"
+    name: Optional[str] = None
     in_chans: int = 3
     num_classes: int = 1000
     block: Type[nn.Module] = ResBlock
@@ -291,17 +291,36 @@ def from_cfg(cls, cfg: ModelCfg):
         return cls(**cfg.dict())
 
     def __call__(self):
-        model = nn.Sequential(
+        model_name = self.name or self.__class__.__name__
+        named_sequential = type(model_name, (nn.Sequential, ), {})
+        model = named_sequential(
             OrderedDict([("stem", self.stem), ("body", self.body), ("head", self.head)])
         )
         self.init_cnn(model)  # pylint: disable=too-many-function-args
-        model.extra_repr = lambda: f"{self.name}"
+        extra_repr = self.get_extra_repr()
+        if extra_repr:
+            model.extra_repr = lambda: extra_repr
         return model
 
+    def get_extra_repr(self) -> str:
+        return " ".join(
+            f"{field}: {self.get_str_value(field)},"
+            for field in self.__fields_set__ if field != "name"
+        )[:-1]
+
+    def get_str_value(self, field: str) -> str:
+        value = getattr(self, field)
+        if isinstance(value, type):
+            value = value.__name__
+        if isinstance(value, partial):
+            value = f"{value.func.__name__} {value.keywords}"
+        return value
+
     def __repr__(self):
         se_repr = self.se.__name__ if self.se else "False"  # type: ignore
+        model_name = self.name or self.__class__.__name__
         return (
-            f"{self.name} constructor\n"
+            f"{model_name}\n"
             f"  in_chans: {self.in_chans}, num_classes: {self.num_classes}\n"
             f"  expansion: {self.expansion}, groups: {self.groups}, dw: {self.dw}, div_groups: {self.div_groups}\n"
             f"  act_fn: {self.act_fn.__name__}, sa: {self.sa}, se: {se_repr}\n"
@@ -312,10 +331,8 @@ def __repr__(self):
 
 
 class XResNet34(ModelConstructor):
-    name: str = "xresnet34"
     layers: list[int] = [3, 4, 6, 3]
 
 
 class XResNet50(XResNet34):
-    name: str = "xresnet50"
     expansion: int = 4

src/model_constructor/yaresnet.py

@@ -124,13 +124,11 @@ def forward(self, x):
 
 
 class YaResNet34(ModelConstructor):
-    name: str = 'YaResnet34'
     block: Type[nn.Module] = YaResBlock
     expansion: int = 1
     layers: List[int] = [3, 4, 6, 3]
     act_fn: Type[nn.Module] = Mish
 
 
 class YaResNet50(YaResNet34):
-    name: str = 'YaResnet50'
     expansion: int = 4

tests/test_mc.py

@@ -11,6 +11,8 @@ def test_MC():
     """test ModelConstructor"""
     img_size = 16
     mc = ModelConstructor()
+    assert "name=None" in str(mc)
+    mc.name = "MC"
     assert "name='MC'" in str(mc)
     model = mc()
     xb = torch.randn(bs_test, 3, img_size, img_size)