ECA-Net:卷积神经网络的高效通道注意力机制.

ECA-Net设计了一种高效通道注意力(Efficient Channel Attention)模块,通过把通道注意力模块中的全连接层替换为一维卷积层,实现了轻量级的注意力模块。

作者经验性地给出了特征通道数$C$和一维卷积核大小$k$之间的关系:

\[C = 2^{\gamma k - b}\]

用以指导卷积核大小$k$的选择:

\[k = | \frac{\log_2C}{\gamma}+\frac{b}{\gamma} |_{odd}\]

实验中设置$\gamma=2,b=1$。实验结果表明这种自适应的核大小选择策略比固定大小的卷积核表现更好:

模型实现如下:

import math
from torch import nn

class EfficientCA(nn.Module):
    def __init__(self, channel, gamma=2, b=1):
        super(EfficientCA, self).__init__()
        t = int(abs((math.log(channel,2) + b) / gamma))
        k = t if t % 2 else t + 1
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size=k, padding=int(k/2)) 
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.size()
        y = self.avg_pool(x) # [b, c, 1, 1]
        y = self.conv(y.squeeze(-1).transpose(-1, -2)) # [b, 1, c]
        y = y.transpose(-1, -2).unsqueeze(-1) # [b, c, 1, 1]
        y = self.sigmoid(y)
        return x * y.expand_as(x)
		
if __name__ == "__main__":
    t = torch.ones((32, 256, 24, 24))
    eca = EfficientCA(256)
    out = eca(t)
    print(out.shape)

实验结果表明,ECA-Net比最先进的图像分类方法具有更低的模型复杂性,同时获得了非常有竞争力的性能。