Implement SK Module in Glasses

glasses/nn/att/SelectiveKernel.py

@@ -0,0 +1,98 @@
+import torch
+import torch.nn as nn
+from typing import Union, List
+
+from glasses.nn.att.utils import make_divisible
+from ..blocks import ConvBnAct
+from einops.layers.torch import Rearrange, Reduce
+
+def _kernel_valid(k):
+    if isinstance(k, (list, tuple)):
+        for ki in k:
+            return _kernel_valid(ki)
+    assert k >=3 and k % 2
+
+class SelectiveKernelAtt(nn.Module):
+    def __init__(
+        self,
+        features: int,
+        num_paths: int = 2,
+        mid_features: int = 32,
+        act_layer: nn.Module = nn.ReLU,
+        norm_layer: nn.Module = nn.BatchNorm2d,
+    ):
+        super().__init__()
+        self.num_paths = num_paths
+        self.att = nn.Sequential(
+            Reduce("b n c h w -> b c h w", reduction="sum"),
+            Reduce("b c h w -> b c 1 1", reduction="mean"),
+            nn.Conv2d(features, mid_features, kernel_size=1, bias=False),
+            norm_layer(mid_features),
+            act_layer(inplace=True),
+            nn.Conv2d(mid_features, features * num_paths, kernel_size=1, bias=False),
+            Rearrange('b (n c) h w -> b n c h w', n=num_paths, c=features),
+            nn.Softmax(dim=1),
+        )
+
+    def forward(self, x):
+        assert x.shape[1] == self.num_paths
+        x = self.att(x)
+        return x
+
+
+class SelectiveKernel(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int = None,
+        kernel_size: Union[List, int] = None,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        reduction: int = 16,
+        reduction_divisor: int = 8,
+        reduced_features: int = None,
+        keep_3x3: bool = True,
+        activation: nn.Module = nn.ReLU,
+        normalization: nn.Module = nn.BatchNorm2d,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        kernel_size = kernel_size or [3, 5]
+        _kernel_valid(kernel_size)
+        if not isinstance(kernel_size, list):
+            kernel_size = [kernel_size] * 2
+        if keep_3x3:
+            dilation = [dilation * (k - 1) // 2 for k in kernel_size]
+            kernel_size = [3] * len(kernel_size)
+        else:
+            dilation = [dilation] * len(kernel_size)
+        self.num_paths = len(kernel_size)
+        self.in_features = in_features
+        self.out_features = out_features,
+        groups = min(out_features, groups)
+
+        self.paths = nn.ModuleList([
+            ConvBnAct(in_features = in_features, 
+                      out_features = out_features, 
+                      activation = activation, 
+                      normalization=normalization,
+                      mode = "same",
+                      stride=stride,
+                      kernel_size=k, 
+                      dilation=d)
+            for k, d in zip(kernel_size, dilation)
+        ])
+
+        attn_features = reduced_features or make_divisible(out_features // reduction, divisor=reduction_divisor)
+        self.attn = SelectiveKernelAtt(out_features, self.num_paths, attn_features)
+
+    def forward(self, x):
+        x_paths = [op(x) for op in self.paths]
+        x = torch.stack(x_paths, dim=1)
+        x_attn = self.attn(x)
+        x = x * x_attn
+        return torch.sum(x, dim=1)
+
+
+

glasses/nn/att/__init__.py

@@ -1,4 +1,5 @@
 from .se import SpatialSE, ChannelSE, SpatialChannelSE
+from .SelectiveKernel import SelectiveKernel
 from .ECA import ECA
 from .CBAM import CBAM
 from .utils import WithAtt
@@ -9,5 +10,6 @@
     "SpatialChannelSE",
     "SpatialSE",
     "CBAM",
+    "SelectiveKernel",
     "WithAtt"
 ]

glasses/nn/att/utils.py

@@ -23,3 +23,12 @@ def __call__(
         b = self.block(in_features, out_features, *args, **kwargs)
         b.block.add_module("se", self.att(out_features))
         return b
+
+
+def make_divisible(v, divisor=8, min_value=None, round_limit=.9):
+    min_value = min_value or divisor
+    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
+    # Make sure that round down does not go down by more than 10%.
+    if new_v < round_limit * v:
+        new_v += divisor
+    return new_v

test/test_att.py

@@ -1,9 +1,9 @@
 import torch
-from glasses.nn.att import ChannelSE, ECA, SpatialChannelSE, SpatialSE, CBAM
+from glasses.nn.att import ChannelSE, ECA, SpatialChannelSE, SpatialSE, CBAM, SelectiveKernel
 
 
 def test_att():
-    x = torch.rand(1, 48, 8, 8)
+    x = torch.rand(2, 48, 8, 8)
 
     se = SpatialSE(x.shape[1])
     res = se(x)
@@ -33,4 +33,9 @@ def test_att():
 
     res = cbam(x)
     assert res.shape == x.shape
+
+    sk = SelectiveKernel(x.shape[1], kernel_size=[3, 5])
+
+    res = sk(x)
+    assert res.shape == x.shape