black

Author: ayasyrev

src/model_constructor/activations.py

@@ -5,9 +5,24 @@
 from torch.nn import Mish
 
 
-__all__ = ['mish', 'Mish', 'mish_jit', 'MishJit', 'mish_jit_fwd', 'mish_jit_bwd', 'MishJitAutoFn', 'mish_me', 'MishMe',
-           'hard_mish_jit', 'HardMishJit', 'hard_mish_jit_fwd', 'hard_mish_jit_bwd', 'HardMishJitAutoFn',
-           'hard_mish_me', 'HardMishMe']
+__all__ = [
+    "mish",
+    "Mish",
+    "mish_jit",
+    "MishJit",
+    "mish_jit_fwd",
+    "mish_jit_bwd",
+    "MishJitAutoFn",
+    "mish_me",
+    "MishMe",
+    "hard_mish_jit",
+    "HardMishJit",
+    "hard_mish_jit_fwd",
+    "hard_mish_jit_bwd",
+    "HardMishJitAutoFn",
+    "hard_mish_me",
+    "HardMishMe",
+]
 
 
 def mish(x, inplace: bool = False):
@@ -40,7 +55,8 @@ def mish_jit(x, _inplace: bool = False):
 class MishJit(nn.Module):
     def __init__(self, inplace: bool = False):
         """Jit version of Mish.
-        Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681"""
+        Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+        """
         super(MishJit, self).__init__()
 
     def forward(self, x):
@@ -61,8 +77,9 @@ def mish_jit_bwd(x, grad_output):
 
 
 class MishJitAutoFn(torch.autograd.Function):
-    """ Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
     A memory efficient, jit scripted variant of Mish"""
+
     @staticmethod
     def forward(ctx, x):
         ctx.save_for_backward(x)
@@ -79,8 +96,9 @@ def mish_me(x, inplace=False):
 
 
 class MishMe(nn.Module):
-    """ Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
+    """Mish: A Self Regularized Non-Monotonic Neural Activation Function - https://arxiv.org/abs/1908.08681
     A memory efficient, jit scripted variant of Mish"""
+
     def __init__(self, inplace: bool = False):
         super(MishMe, self).__init__()
 
@@ -90,18 +108,19 @@ def forward(self, x):
 
 @torch.jit.script
 def hard_mish_jit(x, inplace: bool = False):
-    """ Hard Mish
+    """Hard Mish
     Experimental, based on notes by Mish author Diganta Misra at
       https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
     """
     return 0.5 * x * (x + 2).clamp(min=0, max=2)
 
 
 class HardMishJit(nn.Module):
-    """ Hard Mish
+    """Hard Mish
     Experimental, based on notes by Mish author Diganta Misra at
       https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
     """
+
     def __init__(self, inplace: bool = False):
         super(HardMishJit, self).__init__()
 
@@ -116,16 +135,17 @@ def hard_mish_jit_fwd(x):
 
 @torch.jit.script
 def hard_mish_jit_bwd(x, grad_output):
-    m = torch.ones_like(x) * (x >= -2.)
-    m = torch.where((x >= -2.) & (x <= 0.), x + 1., m)
+    m = torch.ones_like(x) * (x >= -2.0)
+    m = torch.where((x >= -2.0) & (x <= 0.0), x + 1.0, m)
     return grad_output * m
 
 
 class HardMishJitAutoFn(torch.autograd.Function):
-    """ A memory efficient, jit scripted variant of Hard Mish
+    """A memory efficient, jit scripted variant of Hard Mish
     Experimental, based on notes by Mish author Diganta Misra at
       https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
     """
+
     @staticmethod
     def forward(ctx, x):
         ctx.save_for_backward(x)
@@ -142,10 +162,11 @@ def hard_mish_me(x, inplace: bool = False):
 
 
 class HardMishMe(nn.Module):
-    """ A memory efficient, jit scripted variant of Hard Mish
+    """A memory efficient, jit scripted variant of Hard Mish
     Experimental, based on notes by Mish author Diganta Misra at
       https://github.com/digantamisra98/H-Mish/blob/0da20d4bc58e696b6803f2523c58d3c8a82782d0/README.md
     """
+
     def __init__(self, inplace: bool = False):
         super(HardMishMe, self).__init__()
 

src/model_constructor/convmixer.py

@@ -67,7 +67,6 @@ def __init__(
         bn_1st: bool = False,
         pre_act: bool = False,
     ):
-
         conv_layer: List[tuple[str, nn.Module]] = [
             (
                 "conv",

src/model_constructor/layers.py

@@ -70,7 +70,6 @@ def __init__(
         bn_1st: bool = True,
         zero_bn: bool = False,
     ):
-
         if padding is None:
             padding = kernel_size // 2
         layers: List[tuple[str, nn.Module]] = [
@@ -124,7 +123,6 @@ def __init__(
         groups=1,
         **kwargs  # pylint: disable=unused-argument
     ):
-
         if padding is None:
             padding = ks // 2
         layers = [
@@ -197,6 +195,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class SEBlock(nn.Module):
     """se block"""
+
     # first version
     se_layer = nn.Linear
     act_fn = nn.ReLU(inplace=True)
@@ -226,6 +225,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class SEBlockConv(nn.Module):
     """se block with conv on excitation"""
+
     # first version
     se_layer = nn.Conv2d
     act_fn = nn.ReLU(inplace=True)

src/model_constructor/mxresnet.py

@@ -7,6 +7,7 @@ class MxResNet(XResNet):
     stem_sizes: list[int] = [3, 32, 64, 64]
     act_fn: type[nn.Module] = nn.Mish
 
+
 class MxResNet34(XResNet34):
     stem_sizes: list[int] = [3, 32, 64, 64]
     act_fn: type[nn.Module] = nn.Mish

src/model_constructor/universal_blocks.py

@@ -297,14 +297,10 @@ def make_layer(cfg: ModelCfg, layer_num: int) -> nn.Sequential:  # type: ignore
             f"bl_{block_num}",
             cfg.block(
                 cfg.expansion,  # type: ignore
-                block_chs[layer_num]
-                if block_num == 0
-                else block_chs[layer_num + 1],
+                block_chs[layer_num] if block_num == 0 else block_chs[layer_num + 1],
                 block_chs[layer_num + 1],
                 stride if block_num == 0 else 1,
-                sa=cfg.sa
-                if (block_num == num_blocks - 1) and layer_num == 0
-                else None,
+                sa=cfg.sa if (block_num == num_blocks - 1) and layer_num == 0 else None,
                 conv_layer=cfg.conv_layer,
                 act_fn=cfg.act_fn,
                 pool=cfg.pool,
@@ -340,6 +336,7 @@ 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
@@ -359,6 +356,7 @@ class YaResNet(XResNet):
     """Base Yaresnet constructor.
     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