typing

Author: ayasyrev

src/model_constructor/blocks.py

@@ -1,4 +1,4 @@
-from typing import Callable, Union
+from typing import Callable, Optional
 
 import torch
 from torch import nn
@@ -22,10 +22,10 @@ def __init__(
         bn_1st: bool = True,
         groups: int = 1,
         dw: bool = False,
-        div_groups: Union[None, int] = None,
-        pool: Union[Callable[[], nn.Module], None] = None,
-        se: Union[nn.Module, None] = None,
-        sa: Union[nn.Module, None] = None,
+        div_groups: Optional[int] = None,
+        pool: Optional[Callable[[], nn.Module]] = None,
+        se: Optional[nn.Module] = None,
+        sa: Optional[nn.Module] = None,
     ):
         super().__init__()
         # pool defined at ModelConstructor.
@@ -107,10 +107,10 @@ def __init__(
         bn_1st: bool = True,
         groups: int = 1,
         dw: bool = False,
-        div_groups: Union[None, int] = None,
-        pool: Union[Callable[[], nn.Module], None] = None,
-        se: Union[nn.Module, None] = None,
-        sa: Union[nn.Module, None] = None,
+        div_groups: Optional[int] = None,
+        pool: Optional[Callable[[], nn.Module]] = None,
+        se: Optional[nn.Module] = None,
+        sa: Optional[nn.Module] = None,
     ):
         super().__init__()
         # pool defined at ModelConstructor.

src/model_constructor/model_constructor.py

@@ -33,11 +33,11 @@ class ModelCfg(Cfg, arbitrary_types_allowed=True, extra="forbid"):
     expansion: int = 1
     groups: int = 1
     dw: bool = False
-    div_groups: Union[int, None] = None
+    div_groups: Optional[int] = None
     sa: Union[bool, type[nn.Module]] = False
     se: Union[bool, type[nn.Module]] = False
-    se_module: Union[bool, None] = None
-    se_reduction: Union[int, None] = None
+    se_module: Optional[bool] = None
+    se_reduction: Optional[int] = None
     bn_1st: bool = True
     zero_bn: bool = True
     stem_stride_on: int = 0
@@ -186,7 +186,7 @@ def from_cfg(cls, cfg: ModelCfg):
 
     @classmethod
     def create_model(
-        cls, cfg: Union[ModelCfg, None] = None, **kwargs: dict[str, Any]
+        cls, cfg: Optional[ModelCfg] = None, **kwargs: dict[str, Any]
     ) -> nn.Sequential:
         if cfg:
             return cls(**cfg.model_dump())()

src/model_constructor/universal_blocks.py

@@ -1,4 +1,4 @@
-from typing import Callable, Union
+from typing import Callable, Optional, Union
 
 import torch
 from torch import nn
@@ -32,10 +32,10 @@ def __init__(
         bn_1st: bool = True,
         groups: int = 1,
         dw: bool = False,
-        div_groups: Union[None, int] = None,
-        pool: Union[Callable[[], nn.Module], None] = None,
-        se: Union[nn.Module, None] = None,
-        sa: Union[nn.Module, None] = None,
+        div_groups: Optional[int] = None,
+        pool: Optional[Callable[[], nn.Module]] = None,
+        se: Optional[nn.Module] = None,
+        sa: Optional[nn.Module] = None,
     ):
         super().__init__()
         # pool defined at ModelConstructor.
@@ -156,10 +156,10 @@ def __init__(
         bn_1st: bool = True,
         groups: int = 1,
         dw: bool = False,
-        div_groups: Union[None, int] = None,
-        pool: Union[Callable[[], nn.Module], None] = None,
-        se: Union[type[nn.Module], None] = None,
-        sa: Union[type[nn.Module], None] = None,
+        div_groups: Optional[int] = None,
+        pool: Optional[Callable[[], nn.Module]] = None,
+        se: Optional[type[nn.Module]] = None,
+        sa: Optional[type[nn.Module]] = None,
     ):
         super().__init__()
         # pool defined at ModelConstructor.

src/model_constructor/yaresnet.py

@@ -38,10 +38,10 @@ def __init__(
         bn_1st: bool = True,
         groups: int = 1,
         dw: bool = False,
-        div_groups: Union[None, int] = None,
-        pool: Union[Callable[[], nn.Module], None] = None,
-        se: Union[nn.Module, None] = None,
-        sa: Union[nn.Module, None] = None,
+        div_groups: Optional[int] = None,
+        pool: Optional[Callable[[], nn.Module]] = None,
+        se: Optional[nn.Module] = None,
+        sa: Optional[nn.Module] = None,
     ):
         super().__init__()
         # pool defined at ModelConstructor.
@@ -123,10 +123,10 @@ def __init__(
         bn_1st: bool = True,
         groups: int = 1,
         dw: bool = False,
-        div_groups: Union[None, int] = None,
-        pool: Union[Callable[[], nn.Module], None] = None,
-        se: Union[nn.Module, None] = None,
-        sa: Union[nn.Module, None] = None,
+        div_groups: Optional[int] = None,
+        pool: Optional[Callable[[], nn.Module]] = None,
+        se: Optional[nn.Module] = None,
+        sa: Optional[nn.Module] = None,
     ):
         super().__init__()
         # pool defined at ModelConstructor.