Mecanismo de atenção ECANet - implementação pytorch

Paper Portal: ECA-Net: Canal de Atenção Eficiente para Redes Neurais Convolucionais Profundas

Objetivo da ECANet:

Adicione um módulo de mecanismo de atenção de canal ECAblock (Efficient Channel Attention) à rede para melhorar o desempenho da rede.

Fonte do ECAblock:

O autor acredita que existem deficiências no mecanismo de atenção SE:
①É ineficiente e desnecessário usar a camada totalmente conectada para calcular a relação entre todos os canais;
②A redução do canal (redução da dimensionalidade do canal) na camada totalmente conectada reduzirá o desempenho de atenção do mecanismo de força.

A estrutura do ECAblock:

Semelhante ao SEblock , use global average pooling (GAP, global average
pooling) para a camada de feição de entrada para obter o vetor de feição e, em seguida, use a convolução 1d (observe a distinção entre conv1d e 1x1conv2d) e a função de ativação sigmoide para calcular o peso de importância de cada canal, Finalmente, a camada de feição de entrada é multiplicada pelo peso de importância do canal correspondente.
Alt
Entre eles, o tamanho do kernel da convolução 1d é determinado pela seguinte fórmula (ímpar significa número ímpar), no texto original, $γ$ é definido como 2, $b$ é definido como 1.
Alt

Vantagens do ECAblock:

①Sem redução de canal (sem DR);
②Interação entre canais;
③Leve.
Alt

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相乘