From d63ae121d52463d86116021e9ee9160a4628e46a Mon Sep 17 00:00:00 2001
From: Chris Ha <15088501+VRandme@users.noreply.github.com>
Date: Thu, 6 Feb 2020 22:44:33 +0900
Subject: [PATCH] Clean up eca_module code

functionally similar
adjusted rwightman's version of reshaping and viewing.
Use F.pad for circular eca version for cleaner code
---
 timm/models/eca_module.py | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/timm/models/eca_module.py b/timm/models/eca_module.py
index 7664354e..5cb52d96 100644
--- a/timm/models/eca_module.py
+++ b/timm/models/eca_module.py
@@ -49,7 +49,19 @@ class eca_layer(nn.Module):
         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):
+        # feature descriptor on the global spatial information
+        y = self.avg_pool(x)
+        # reshape for convolution
+        y = y.view(x.shape[0], 1, -1)
+        # Two different branches of ECA module
+        y = self.conv(y)
+        # Multi-scale information fusion
+        y = self.sigmoid(y.view(x.shape[0], -1, 1, 1))
+        return x * y.expand_as(x)
 
+
+    '''original implementation
     def forward(self, x):
         # x: input features with shape [b, c, h, w]
         b, c, h, w = x.size()
@@ -62,8 +74,10 @@ class eca_layer(nn.Module):
 
         # Multi-scale information fusion
         y = self.sigmoid(y)
-
         return x * y.expand_as(x)
+    '''
+
+
 
 class ceca_layer(nn.Module):
     """Constructs a circular ECA module.
@@ -75,7 +89,6 @@ class ceca_layer(nn.Module):
     (accuracy, robustness), without signficantly impacting resource metrics 
     (parameter size, throughput,latency, etc)
 
-
     Args:
         channel: Number of channels of the input feature map
         k_size: Adaptive selection of kernel size
@@ -92,19 +105,16 @@ class ceca_layer(nn.Module):
         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)
         
-        #manually implement circular padding
-        y = torch.cat((y[:,:self.padding,:,:], y, y[:,-self.padding:,:,:]),dim=1)
-        
+        #manually implement circular padding, F.pad does not seemed to be bugged
+        y = F.pad(y.view(x.shape[0],1,-1),(self.padding,self.padding),mode='circular')
 
         # Two different branches of ECA module
-        y = self.conv(y.squeeze(-1).transpose(-1, -2)).transpose(-1, -2).unsqueeze(-1)
+        y = self.conv(y)
 
         # Multi-scale information fusion
-        y = self.sigmoid(y)
+        y = self.sigmoid(y.view(x.shape[0], -1, 1, 1))
 
         return x * y.expand_as(x)