type ModSeq

Author: ayasyrev

noxfile_conda.py

@@ -4,7 +4,6 @@
 @nox.session(python=["3.9", "3.10", "3.11"], venv_backend="mamba")
 def conda_tests(session: nox.Session) -> None:
     args = session.posargs or ["--cov"]
-    # session.install("pytest", "pytest-cov")
     session.conda_install("pytest", "pytest-cov")
     session.conda_install("pytorch")
     session.conda_install("pydantic")

src/model_constructor/helpers.py

@@ -1,12 +1,13 @@
 from collections import OrderedDict
 from functools import partial
-from typing import Iterable, Optional
+from typing import Iterable, Optional, Union
 from pydantic import BaseModel
 
 from torch import nn
 
 
 ListStrMod = list[tuple[str, nn.Module]]
+ModSeq = Union[nn.Module, nn.Sequential]
 
 
 def nn_seq(list_of_tuples: Iterable[tuple[str, nn.Module]]) -> nn.Sequential:

src/model_constructor/model_constructor.py

@@ -6,7 +6,7 @@
 from torch import nn
 
 from .blocks import BasicBlock, BottleneckBlock
-from .helpers import Cfg, ListStrMod, init_cnn, nn_seq
+from .helpers import Cfg, ListStrMod, ModSeq, init_cnn, nn_seq
 from .layers import ConvBnAct, SEModule, SimpleSelfAttention
 
 __all__ = [
@@ -138,10 +138,10 @@ class ModelConstructor(ModelCfg):
     """Model constructor. As default - resnet18"""
 
     init_cnn: Callable[[nn.Module], None] = init_cnn
-    make_stem: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = make_stem  # type: ignore
-    make_layer: Callable[[ModelCfg, int], Union[nn.Module, nn.Sequential]] = make_layer  # type: ignore
-    make_body: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = make_body  # type: ignore
-    make_head: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = make_head  # type: ignore
+    make_stem: Callable[[ModelCfg], ModSeq] = make_stem  # type: ignore
+    make_layer: Callable[[ModelCfg, int], ModSeq] = make_layer  # type: ignore
+    make_body: Callable[[ModelCfg], ModSeq] = make_body  # type: ignore
+    make_head: Callable[[ModelCfg], ModSeq] = make_head  # type: ignore
 
     @field_validator("se")
     def set_se(  # pylint: disable=no-self-argument

src/model_constructor/universal_blocks.py

@@ -1,9 +1,9 @@
-from typing import Callable, Optional, Union
+from typing import Callable, Optional
 
 import torch
 from torch import nn
 
-from .helpers import nn_seq
+from .helpers import ModSeq, nn_seq
 from .layers import ConvBnAct, get_act
 from .model_constructor import ListStrMod, ModelCfg, ModelConstructor
 
@@ -337,10 +337,10 @@ def make_head(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
 class XResNet(ModelConstructor):
     """Base Xresnet constructor."""
 
-    make_stem: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = make_stem
-    make_layer: Callable[[ModelCfg, int], Union[nn.Module, nn.Sequential]] = make_layer
-    make_body: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = make_body
-    make_head: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = make_head
+    make_stem: Callable[[ModelCfg], ModSeq] = make_stem
+    make_layer: Callable[[ModelCfg, int], ModSeq] = make_layer
+    make_body: Callable[[ModelCfg], ModSeq] = make_body
+    make_head: Callable[[ModelCfg], ModSeq] = make_head
     block: type[nn.Module] = XResBlock
 
 

src/model_constructor/xresnet.py

@@ -1,10 +1,10 @@
 from functools import partial
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional
 
 from torch import nn
 
 from .blocks import BottleneckBlock
-from .helpers import ListStrMod, nn_seq
+from .helpers import ListStrMod, nn_seq, ModSeq
 from .model_constructor import ModelCfg, ModelConstructor
 
 __all__ = [
@@ -39,7 +39,7 @@ def xresnet_stem(cfg: ModelCfg) -> nn.Sequential:  # type: ignore
 
 
 class XResNet(ModelConstructor):
-    make_stem: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = xresnet_stem
+    make_stem: Callable[[ModelCfg], ModSeq] = xresnet_stem
     stem_sizes: list[int] = [32, 32, 64]
     pool: Optional[Callable[[Any], nn.Module]] = partial(
         nn.AvgPool2d, kernel_size=2, ceil_mode=True

src/model_constructor/yaresnet.py

@@ -2,12 +2,12 @@
 # Yet another ResNet.
 
 from functools import partial
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Optional
 
 import torch
 from torch import nn
 
-from model_constructor.helpers import nn_seq
+from model_constructor.helpers import ModSeq, nn_seq
 
 from .layers import ConvBnAct, get_act
 from .model_constructor import ListStrMod, ModelConstructor, ModelCfg
@@ -203,7 +203,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:  # type: ignore
 
 
 class YaResNet(ModelConstructor):
-    make_stem: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = xresnet_stem
+    make_stem: Callable[[ModelCfg], ModSeq] = xresnet_stem
     stem_sizes: list[int] = [32, 64, 64]
     block: type[nn.Module] = YaBasicBlock
     act_fn: type[nn.Module] = nn.Mish