Merge pull request #7 from dpxudong/master

add se attention

Author: likyoo

change_detection_pytorch/base/modules.py

@@ -140,6 +140,24 @@ def forward(self, x):
         return out
 
 
+class SEModule(nn.Module):
+    def __init__(self, in_channels, reduction=16):
+        super(SEModule, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(in_channels, in_channels // reduction, bias=False),
+            nn.ReLU(inplace=True),
+            nn.Linear(in_channels // reduction, in_channels, bias=False),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y.expand_as(x)
+
+
 class ArgMax(nn.Module):
 
     def __init__(self, dim=None):
@@ -196,6 +214,8 @@ def __init__(self, name, **params):
             self.attention = CBAMSpatial(**params)
         elif name == 'cbam':
             self.attention = CBAM(**params)
+        elif name == 'se':
+            self.attention = SEModule(**params)
         else:
             raise ValueError("Attention {} is not implemented".format(name))