Merge pull request #110 from ayasyrev/test_typing

Test typing

Author: ayasyrev

noxfile_conda.py

@@ -1,7 +1,7 @@
 import nox
 
 
-@nox.session(python=["3.9", "3.10", "3.11"], venv_backend="mamba")
+@nox.session(python=["3.8", "3.9", "3.10", "3.11"], venv_backend="mamba")
 def conda_tests(session: nox.Session) -> None:
     args = session.posargs or ["--cov"]
     session.conda_install("pytest", "pytest-cov")

src/model_constructor/blocks.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from typing import Callable, Optional, Type
 
 import torch
 from torch import nn
@@ -16,8 +16,8 @@ def __init__(
         in_channels: int,
         out_channels: int,
         stride: int = 1,
-        conv_layer: type[ConvBnAct] = ConvBnAct,
-        act_fn: type[nn.Module] = nn.ReLU,
+        conv_layer: Type[ConvBnAct] = ConvBnAct,
+        act_fn: Type[nn.Module] = nn.ReLU,
         zero_bn: bool = True,
         bn_1st: bool = True,
         groups: int = 1,
@@ -101,8 +101,8 @@ def __init__(
         out_channels: int,
         stride: int = 1,
         expansion: int = 4,
-        conv_layer: type[ConvBnAct] = ConvBnAct,
-        act_fn: type[nn.Module] = nn.ReLU,
+        conv_layer: Type[ConvBnAct] = ConvBnAct,
+        act_fn: Type[nn.Module] = nn.ReLU,
         zero_bn: bool = True,
         bn_1st: bool = True,
         groups: int = 1,

src/model_constructor/convmixer.py

@@ -3,11 +3,13 @@
 # Adopted from https://github.com/tmp-iclr/convmixer
 # Home for convmixer: https://github.com/locuslab/convmixer
 from collections import OrderedDict
-from typing import Callable, List, Optional, Union
+from typing import Callable, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 
+from .helpers import ListStrMod
+
 
 class Residual(nn.Module):
     def __init__(self, fn: Callable[[torch.Tensor], torch.Tensor]):
@@ -59,15 +61,15 @@ def __init__(
         self,
         in_channels: int,
         out_channels: int,
-        kernel_size: Union[int, tuple[int, int]],
+        kernel_size: Union[int, Tuple[int, int]],
         stride: int = 1,
         act_fn: nn.Module = nn.GELU(),
         padding: Union[int, str] = 0,
         groups: int = 1,
         bn_1st: bool = False,
         pre_act: bool = False,
     ):
-        conv_layer: List[tuple[str, nn.Module]] = [
+        conv_layer: ListStrMod = [
             (
                 "conv",
                 nn.Conv2d(
@@ -80,7 +82,7 @@ def __init__(
                 ),
             )
         ]
-        act_bn: List[tuple[str, nn.Module]] = [
+        act_bn: ListStrMod = [
             ("act_fn", act_fn),
             ("bn", nn.BatchNorm2d(out_channels)),
         ]

src/model_constructor/helpers.py

@@ -1,16 +1,16 @@
 import importlib
 from collections import OrderedDict
 from functools import partial
-from typing import Any, Iterable, Optional, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from pydantic import BaseModel
 from torch import nn
 
-ListStrMod = list[tuple[str, nn.Module]]
+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:
+def nn_seq(list_of_tuples: Iterable[Tuple[str, nn.Module]]) -> nn.Sequential:
     """return nn.Sequential from OrderedDict from list of tuples"""
     return nn.Sequential(OrderedDict(list_of_tuples))
 
@@ -86,22 +86,22 @@ def _get_str_value(self, field: str) -> str:
     def __repr__(self) -> str:
         return f"{self.__repr_name__()}(\n  {self.__repr_str__(chr(10) + '  ')})"
 
-    def __repr_args__(self) -> list[tuple[str, str]]:
+    def __repr_args__(self) -> List[Tuple[str, str]]:
         return [
             (field, str_value)
             for field in self.model_fields
             if (str_value := self._get_str_value(field))
         ]
 
-    def __repr_set_fields__(self) -> list[str]:
+    def __repr_set_fields__(self) -> List[str]:
         """Return list repr for fields set at init"""
         return [
             f"{field}: {self._get_str_value(field)}"
             for field in self.model_fields_set  # pylint: disable=E1133
             if field != "name"
         ]
 
-    def __repr_changed_fields__(self) -> list[str]:
+    def __repr_changed_fields__(self) -> List[str]:
         """Return list repr for changed fields"""
         return [
             f"{field}: {self._get_str_value(field)}"
@@ -110,7 +110,7 @@ def __repr_changed_fields__(self) -> list[str]:
         ]
 
     @property
-    def changed_fields(self) -> dict[str, Any]:
+    def changed_fields(self) -> Dict[str, Any]:
         # return "\n".join(self.__repr_changed_fields__())
         return {
             field: self._get_str_value(field)

src/model_constructor/layers.py

@@ -1,10 +1,12 @@
 from collections import OrderedDict
-from typing import List, Optional, Type, Union
+from typing import Optional, Type, Union
 
 import torch
 from torch import nn
 from torch.nn.utils.spectral_norm import spectral_norm
 
+from .helpers import ListStrMod
+
 __all__ = [
     "Flatten",
     "noop",
@@ -72,7 +74,7 @@ def __init__(
     ):
         if padding is None:
             padding = kernel_size // 2
-        layers: List[tuple[str, nn.Module]] = [
+        layers: ListStrMod = [
             (
                 "conv",
                 self.convolution_module(
@@ -262,10 +264,10 @@ def __init__(
         reduction: int = 16,
         rd_channels: Optional[int] = None,
         rd_max: bool = False,
-        se_layer: type[nn.Module] = nn.Linear,
+        se_layer: Type[nn.Module] = nn.Linear,
         act_fn: nn.Module = nn.ReLU(inplace=True),
         use_bias: bool = True,
-        gate: type[nn.Module] = nn.Sigmoid,
+        gate: Type[nn.Module] = nn.Sigmoid,
     ):
         super().__init__()
         reducted = max(channels // reduction, 1)  # preserve zero-element tensors
@@ -302,10 +304,10 @@ def __init__(
         reduction: int = 16,
         rd_channels: Optional[int] = None,
         rd_max: bool = False,
-        se_layer: type[nn.Module] = nn.Conv2d,
+        se_layer: Type[nn.Module] = nn.Conv2d,
         act_fn: nn.Module = nn.ReLU(inplace=True),
         use_bias: bool = True,
-        gate: type[nn.Module] = nn.Sigmoid,
+        gate: Type[nn.Module] = nn.Sigmoid,
     ):
         super().__init__()
         #       rd_channels = math.ceil(channels//reduction/8)*8

src/model_constructor/model_constructor.py

@@ -1,14 +1,21 @@
 from collections import OrderedDict
 from functools import partial
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Union, Type
 
 from pydantic import field_validator
 from pydantic_core.core_schema import FieldValidationInfo
 from torch import nn
 
 from .blocks import BasicBlock, BottleneckBlock
-from .helpers import (Cfg, ListStrMod, ModSeq, init_cnn, instantiate_module,
-                      is_module, nn_seq)
+from .helpers import (
+    Cfg,
+    ListStrMod,
+    ModSeq,
+    init_cnn,
+    instantiate_module,
+    is_module,
+    nn_seq,
+)
 from .layers import ConvBnAct, SEModule, SimpleSelfAttention
 
 __all__ = [
@@ -25,7 +32,7 @@
 }
 
 
-nnModule = Union[type[nn.Module], Callable[[], nn.Module]]
+nnModule = Union[Type[nn.Module], Callable[[], nn.Module]]
 
 
 class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
@@ -36,8 +43,8 @@ class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
     num_classes: int = 1000
     block: Union[nnModule, str] = BasicBlock
     conv_layer: Union[nnModule, str] = ConvBnAct
-    block_sizes: list[int] = [64, 128, 256, 512]
-    layers: list[int] = [2, 2, 2, 2]
+    block_sizes: List[int] = [64, 128, 256, 512]
+    layers: List[int] = [2, 2, 2, 2]
     norm: Union[nnModule, str] = nn.BatchNorm2d
     act_fn: Union[nnModule, str] = nn.ReLU
     pool: Union[nnModule, str, None] = None
@@ -52,15 +59,16 @@ class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
     bn_1st: bool = True
     zero_bn: bool = True
     stem_stride_on: int = 0
-    stem_sizes: list[int] = [64]
+    stem_sizes: List[int] = [64]
     stem_pool: Union[nnModule, str, None] = 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[nnModule, str],
+        cls,
+        value: Union[nnModule, str],
     ) -> nnModule:
         """Check values, if string, convert to nn.Module."""
         if is_module(value):
@@ -69,7 +77,9 @@ def set_modules(  # pylint: disable=no-self-argument
 
     @field_validator("se", "sa")
     def set_se(  # pylint: disable=no-self-argument
-        cls, value: Union[bool, nnModule, str], info: FieldValidationInfo,
+        cls,
+        value: Union[bool, nnModule, str],
+        info: FieldValidationInfo,
     ) -> nnModule:
         if isinstance(value, (int, bool)):
             return DEFAULT_SE_SA[info.field_name]
@@ -154,8 +164,8 @@ def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:  # type: ignore
 
 
 def make_body(
-        cfg: ModelCfg,
-        layer_constructor: Callable[[ModelCfg, int], nn.Sequential] = make_layer,
+    cfg: ModelCfg,
+    layer_constructor: Callable[[ModelCfg, int], nn.Sequential] = make_layer,
 ) -> nn.Sequential:
     """Create model body."""
     if hasattr(cfg, "make_layer"):
@@ -203,7 +213,7 @@ def from_cfg(cls, cfg: ModelCfg):
 
     @classmethod
     def create_model(
-        cls, cfg: Optional[ModelCfg] = None, **kwargs: dict[str, Any]
+        cls, cfg: Optional[ModelCfg] = None, **kwargs: Dict[str, Any]
     ) -> nn.Sequential:
         if cfg:
             return cls(**cfg.model_dump(exclude_none=True))()
@@ -226,9 +236,9 @@ def __call__(self) -> nn.Sequential:
 
 
 class ResNet34(ModelConstructor):
-    layers: list[int] = [3, 4, 6, 3]
+    layers: List[int] = [3, 4, 6, 3]
 
 
 class ResNet50(ResNet34):
-    block: type[nn.Module] = BottleneckBlock
-    block_sizes: list[int] = [256, 512, 1024, 2048]
+    block: Type[nn.Module] = BottleneckBlock
+    block_sizes: List[int] = [256, 512, 1024, 2048]

src/model_constructor/mxresnet.py

@@ -1,17 +1,18 @@
+from typing import List, Type
 from torch import nn
 
 from .xresnet import XResNet
 
 
 class MxResNet(XResNet):
-    stem_sizes: list[int] = [3, 32, 64, 64]
-    act_fn: type[nn.Module] = nn.Mish
+    stem_sizes: List[int] = [3, 32, 64, 64]
+    act_fn: Type[nn.Module] = nn.Mish
 
 
 class MxResNet34(MxResNet):
-    layers: list[int] = [3, 4, 6, 3]
+    layers: List[int] = [3, 4, 6, 3]
 
 
 class MxResNet50(MxResNet34):
     expansion: int = 4
-    block_sizes: list[int] = [256, 512, 1024, 2048]
+    block_sizes: List[int] = [256, 512, 1024, 2048]

src/model_constructor/universal_blocks.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from typing import Callable, List, Optional, Type
 
 import torch
 from torch import nn
@@ -26,8 +26,8 @@ def __init__(
         in_channels: int,
         mid_channels: int,
         stride: int = 1,
-        conv_layer: type[ConvBnAct] = ConvBnAct,
-        act_fn: type[nn.Module] = nn.ReLU,
+        conv_layer: Type[ConvBnAct] = ConvBnAct,
+        act_fn: Type[nn.Module] = nn.ReLU,
         zero_bn: bool = True,
         bn_1st: bool = True,
         groups: int = 1,
@@ -150,16 +150,16 @@ def __init__(
         in_channels: int,
         mid_channels: int,
         stride: int = 1,
-        conv_layer: type[ConvBnAct] = ConvBnAct,
-        act_fn: type[nn.Module] = nn.ReLU,
+        conv_layer: Type[ConvBnAct] = ConvBnAct,
+        act_fn: Type[nn.Module] = nn.ReLU,
         zero_bn: bool = True,
         bn_1st: bool = True,
         groups: int = 1,
         dw: bool = False,
         div_groups: Optional[int] = None,
         pool: Optional[Callable[[], nn.Module]] = None,
-        se: Optional[type[nn.Module]] = None,
-        sa: Optional[type[nn.Module]] = None,
+        se: Optional[Type[nn.Module]] = None,
+        sa: Optional[Type[nn.Module]] = None,
     ):
         super().__init__()
         # pool defined at ModelConstructor.
@@ -265,7 +265,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 def make_stem(cfg: ModelCfg) -> nn.Sequential:
     """Create xResnet stem -> 3 conv 3*3 instead of 1 conv 7*7"""
     len_stem = len(cfg.stem_sizes)
-    stem: list[tuple[str, nn.Module]] = [
+    stem: ListStrMod = [
         (
             f"conv_{i}",
             cfg.conv_layer(
@@ -341,11 +341,11 @@ class XResNet(ModelConstructor):
     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
+    block: Type[nn.Module] = XResBlock
 
 
 class XResNet34(XResNet):
-    layers: list[int] = [3, 4, 6, 3]
+    layers: List[int] = [3, 4, 6, 3]
 
 
 class XResNet50(XResNet34):
@@ -357,13 +357,13 @@ class YaResNet(XResNet):
     YaResBlock, Mish activation, custom stem.
     """
 
-    block: type[nn.Module] = YaResBlock
-    stem_sizes: list[int] = [3, 32, 64, 64]
-    act_fn: type[nn.Module] = nn.Mish
+    block: Type[nn.Module] = YaResBlock
+    stem_sizes: List[int] = [3, 32, 64, 64]
+    act_fn: Type[nn.Module] = nn.Mish
 
 
 class YaResNet34(YaResNet):
-    layers: list[int] = [3, 4, 6, 3]
+    layers: List[int] = [3, 4, 6, 3]
 
 
 class YaResNet50(YaResNet34):