ECANet アテンションメカニズム - pytorch の実装

論文ポータル: ECA-Net: ディープ畳み込みニューラルネットワークの効率的なチャネルアテンション

ECANet の目的:

ECAblock (Efficient Channel Attendance) チャネルアテンションメカニズムモジュールをネットワークに追加して、ネットワークパフォーマンスを向上させます。

ECAblock のソース:

著者は、SE アテンションのメカニズムには、
①すべてのチャネル間の関係を計算するために全結合層を使用するのは非効率的かつ不必要である、
②全結合層でのチャネルの削減 (チャネルの次元削減) により、注目のフォース機構の性能。

ECAblock の構造:

SEblockと同様に、入力フィーチャレイヤーに対してグローバル平均プーリング(GAP、グローバル平均
プーリング)を使用して特徴ベクトルを取得し、次に1d 畳み込み(conv1d と 1x1conv2d の違いに注意してください) とシグモイドアクティベーション関数を使用して重要度の重みを計算します。最後に、入力フィーチャレイヤーに、対応するチャネル重要度の重みが乗算されます。
オルタナティブ
このうち、1d コンボリューションのコンボリューションカーネルサイズは次の式で決まります (odd は奇数を意味します)。原文では、 $γ$ は 2、 $b$ は1に設定されます。

ECAblock の利点:

①チャンネル削減なし（DRなし）、
②クロスチャンネルインタラクション、
③軽量。
オルタナティブ

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