Add CL SE module

timm/models/layers/squeeze_excite.py

@@ -72,3 +72,31 @@ class EffectiveSEModule(nn.Module):
 
 
 EffectiveSqueezeExcite = EffectiveSEModule  # alias
+
+
+class SqueezeExciteCl(nn.Module):
+    """ SE Module as defined in original SE-Nets with a few additions
+    Additions include:
+        * divisor can be specified to keep channels % div == 0 (default: 8)
+        * reduction channels can be specified directly by arg (if rd_channels is set)
+        * reduction channels can be specified by float rd_ratio (default: 1/16)
+        * global max pooling can be added to the squeeze aggregation
+        * customizable activation, normalization, and gate layer
+    """
+    def __init__(
+            self, channels, rd_ratio=1. / 16, rd_channels=None, rd_divisor=8,
+            bias=True, act_layer=nn.ReLU, gate_layer='sigmoid'):
+        super().__init__()
+        if not rd_channels:
+            rd_channels = make_divisible(channels * rd_ratio, rd_divisor, round_limit=0.)
+        self.fc1 = nn.Linear(channels, rd_channels, bias=bias)
+        self.act = create_act_layer(act_layer, inplace=True)
+        self.fc2 = nn.Linear(rd_channels, channels, bias=bias)
+        self.gate = create_act_layer(gate_layer)
+
+    def forward(self, x):
+        x_se = x.mean((1, 2), keepdims=True)  # FIXME avg dim [1:n-1], don't assume 2D NHWC
+        x_se = self.fc1(x_se)
+        x_se = self.act(x_se)
+        x_se = self.fc2(x_se)
+        return x * self.gate(x_se)