instantiate from module string

Author: ayasyrev

src/model_constructor/helpers.py

@@ -1,11 +1,11 @@
+import importlib
 from collections import OrderedDict
 from functools import partial
-from typing import Iterable, Optional, Union
-from pydantic import BaseModel
+from typing import Any, Iterable, Optional, Union
 
+from pydantic import BaseModel
 from torch import nn
 
-
 ListStrMod = list[tuple[str, nn.Module]]
 ModSeq = Union[nn.Module, nn.Sequential]
 
@@ -25,6 +25,49 @@ def init_cnn(module: nn.Module) -> None:
         init_cnn(layer)
 
 
+def is_module(val: Any) -> bool:
+    """Check if val is a nn.Module or partial of nn.Module."""
+
+    to_check = val
+    if isinstance(val, partial):
+        to_check = val.func
+    try:
+        return issubclass(to_check, nn.Module)
+    except TypeError:
+        return False
+
+
+def instantiate_module(
+    name: str,
+    default_path: Optional[str] = None,
+) -> nn.Module:
+    """Instantiate model from name."""
+    if default_path is None:
+        path_list = name.rsplit(".", 1)
+        if len(path_list) == 1:
+            default_path = "torch.nn"
+            name = path_list[0]
+        else:
+            if path_list[0] == "nn":
+                default_path = "torch.nn"
+                name = path_list[1]
+            else:
+                default_path = path_list[0]
+                name = path_list[1]
+    try:
+        mod = importlib.import_module(default_path)
+    except ImportError:
+        raise ImportError(f"Module {default_path} not found")
+    if hasattr(mod, name):
+        module = getattr(mod, name)
+        if is_module(module):
+            return module
+        else:
+            raise ImportError(f"Module {name} is not a nn.Module")
+    else:
+        raise ImportError(f"Module {name} not found at {default_path}")
+
+
 class Cfg(BaseModel):
     """Base class for config."""
 

src/model_constructor/model_constructor.py

@@ -7,7 +7,8 @@
 from torch import nn
 
 from .blocks import BasicBlock, BottleneckBlock
-from .helpers import Cfg, ListStrMod, ModSeq, init_cnn, nn_seq
+from .helpers import (Cfg, ListStrMod, ModSeq, init_cnn, instantiate_module,
+                      is_module, nn_seq)
 from .layers import ConvBnAct, SEModule, SimpleSelfAttention
 
 __all__ = [
@@ -24,16 +25,7 @@
 }
 
 
-def is_module(val: Any) -> bool:
-    """Check if val is a nn.Module or partial of nn.Module."""
-
-    to_check = val
-    if isinstance(val, partial):
-        to_check = val.func
-    try:
-        return issubclass(to_check, nn.Module)
-    except TypeError:
-        return False
+nnModule = Union[type[nn.Module], Callable[[], nn.Module], str]
 
 
 class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
@@ -42,13 +34,13 @@ class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
     name: Optional[str] = None
     in_chans: int = 3
     num_classes: int = 1000
-    block: type[nn.Module] = BasicBlock
-    conv_layer: type[nn.Module] = ConvBnAct
+    block: nnModule = BasicBlock
+    conv_layer: nnModule = 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: Optional[Callable[[Any], nn.Module]] = None
+    norm: nnModule = nn.BatchNorm2d
+    act_fn: nnModule = nn.ReLU
+    pool: Optional[nnModule] = None
     expansion: int = 1
     groups: int = 1
     dw: bool = False
@@ -61,11 +53,22 @@ class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
     zero_bn: bool = True
     stem_stride_on: int = 0
     stem_sizes: list[int] = [64]
-    stem_pool: Optional[Callable[[], nn.Module]] = partial(
+    stem_pool: Optional[nnModule] = partial(
         nn.MaxPool2d, kernel_size=3, stride=2, padding=1
     )
     stem_bn_end: bool = False
 
+    @field_validator("act_fn", "block", "conv_layer", "norm", "pool", "stem_pool")
+    def set_modules(  # pylint: disable=no-self-argument
+        cls, value: Union[type[nn.Module], str], info: FieldValidationInfo,
+    ) -> Union[type[nn.Module], Callable[[], nn.Module]]:
+        """Check values, if string, convert to nn.Module."""
+        if is_module(value):
+            return value
+        if isinstance(value, str):
+            return instantiate_module(value)
+        raise ValueError(f"{info.field_name} must be str or nn.Module")
+
     @field_validator("se", "sa")
     def set_se(  # pylint: disable=no-self-argument
         cls, value: Union[bool, type[nn.Module]], info: FieldValidationInfo,