ECANet注意力机制——pytorch实现

论文传送门：ECA-Net: Efficient Channel Attention for Deep Convolutional Neural Networks

ECANet的目的：

为网络中添加ECAblock(Efficient Channel Attention)通道注意力机制模块，提高网络性能。

ECAblock的来源：

作者认为SE注意力机制存在缺点：
①使用全连接层来计算所有通道之间的相互关系是低效且不必要的；
②在全连接层进行了通道的缩减(channel dimensionality reduction)，会降低注意力机制的性能。

ECAblock的结构：

与SEblock相似，对输入特征层使用全局平均池化(GAP, global average
pooling)得到特征向量，然后使用1d卷积(注意区分conv1d和1x1conv2d)和sigmoid激活函数计算其每个通道的重要性权重，最后将输入的特征层乘上对应的通道重要性权重。

其中，1d卷积的卷积核尺寸由下式决定(odd表示奇数)，在原文中，$\gamma$设为2，$b$设为1。

ECAblock的优点：

①没有通道缩减(No DR)；
②跨通道相互作用(Cross-channel Interaction)；
③轻量(Lightweight)。

import torch
import torch.nn as nn
import math


class ECAblock(nn.Module):  # Efficient Channel Attention block
    def __init__(self, channels: int, gamma=2, b=1):
        super(ECAblock, self).__init__()
        self.kernel_size = int(abs((math.log(channels, 2) + b) / gamma))  # 计算1d卷积核尺寸
        self.kernel_size = self.kernel_size if self.kernel_size % 2 else self.kernel_size + 1  # 计算1d卷积核尺寸

        self.avgpool = nn.AdaptiveAvgPool2d(1)  # avgpool
        self.conv = nn.Conv1d(1, 1, self.kernel_size, 1, int(self.kernel_size / 2), bias=False)  # 1dconv
        self.sigmoid = nn.Sigmoid()  # sigmoid，将输出压缩到(0,1)

    def forward(self, x):
        weights = self.avgpool(x)  # (N C W H) -> (N C 1 1)
        weights = self.conv(weights.squeeze(-1).transpose(-1, -2))  # (N C 1 1) -> (N C 1) -> (N 1 C) -> (N 1 C)
        weights = self.sigmoid(weights.transpose(-1, -2).unsqueeze(-1))  # (N 1 C) -> (N C 1) -> (N C 1 1) -> (N C 1 1)
        return weights * x  # 将计算得到的weights与输入的feature map相乘