eca depth wise

components/ECA.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: ECA.py
+# Created Date: Tuesday February 23rd 2021
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Tuesday, 23rd February 2021 9:14:28 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2021 Shanghai Jiao Tong University
+#############################################################
+
+import math
+import torch
+from torch import nn
+from torch.nn.parameter import Parameter
+
+class eca_layer(nn.Module):
+    """Constructs a ECA module.
+    Args:
+        channel: Number of channels of the input feature map
+        k_size: Adaptive selection of kernel size
+    """
+    def __init__(self, channel): 
+        super(eca_layer, self).__init__()
+
+        b = 1
+        gamma = 2
+        k_size = int(abs(math.log(channel,2)+b)/gamma)
+        k_size = k_size if k_size % 2 else k_size+1
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.conv = nn.Conv1d(1, 1, kernel_size=k_size, padding=(k_size - 1) // 2, bias=False) 
+        self.sigmoid = nn.Sigmoid()
+
+    def forward(self, x):
+        # x: input features with shape [b, c, h, w]
+        # b, c, h, w = x.size()
+
+        # feature descriptor on the global spatial information
+        y = self.avg_pool(x)
+
+        # Two different branches of ECA module
+        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
+
+        # Multi-scale information fusion
+        y = self.sigmoid(y)
+
+        return x * y.expand_as(x)