ECA-Net代码理解——CVPR 2020

原文链接： https://arxiv.org/abs/1910.03151
代码链接： https://github.com/BangguWu/ECANet

核心代码：

class eca_layer(nn.Module):
    def __init__(self, channel, k_size):
        super(eca_layer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.k_size = k_size
        self.conv = nn.Conv1d(channel, channel, kernel_size=k_size, bias=False, groups=channel)
        self.sigmoid = nn.Sigmoid()


    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x)
        y = nn.functional.unfold(y.transpose(-1, -3), kernel_size=(1, self.k_size), padding=(0, (self.k_size - 1) // 2))
        y = self.conv(y.transpose(-1, -2)).unsqueeze(-1)
        y = self.sigmoid(y)
        x = x * y.expand_as(x)
        return x

1 ECA

前面说过，SE模块由三个部分组成：squeeze、excitation和scale。ECA主要更新SE模块中excitation部分，将SE模块中带通道维度缩放的两层全连接，替换为一层的局部连接。

1.1 局部连接的实现

局部连接的实现结合了nn.functional.unfold()和nn.Conv1d()的使用。

假设输入尺寸为(B, C, H, W)，全局平均池化后后为(B, C, 1, 1)；下面是局部连接的操作，分两步走：

• 在nn.functional.unfold()中首先交换维度为(B, 1, 1, C)，接着以stride=1、kernel_size=(1, self.k_size)滑动提取像素值，纵向摆放，操作后维度为(B, 3, C)；
• 在nn.Conv1d()中首先交换维度为(B, C, 3)，接着使用分组groups=channel卷积提取特征，最后增加一维，操作后维度回到(B,C, 1, 1)

unfold()用法参考：https://blog.csdn.net/guihaiyuan123/article/details/113455775