ECA attention mechanism

Overall network structure diagram

insert image description here

  1. Input the feature map, first go through the global average pooling, change the h and w dimensions to 1, and only keep the channel dimension, then the feature map will be stretched into strips,
  2. After 1D convolution, the channel of each layer interacts with the channel of the adjacent layer to share weights.
  3. Use Sigmoid for processing.
  4. Multiply the input feature map with the processed feature map weight, then the weight will be added to the feature map.

the code

Use pytorch to write:

from numpy import pad
import torch
import torch.nn as nn
import math

"""
    ECA注意力机制,和SE注意力主要的区别是避免降维,减少了参数数量的增加。
    提出了一种自适应卷积核的1D卷积。在通道之间进行信息交互,获取信息。

    步骤:全局平均池化-》1D卷积-》Sigmoid激活-》输出
"""
class ECA(nn.Module):
    def __init__(self, channel, gamma=2, b=1):
        super(ECA, self).__init__()
        # 计算卷积核大小
        kernel_size = int(abs((math.log(channel, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
        # 计算padding
        padding = kernel_size // 2
        self.avg = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(
            1, 1, kernel_size=kernel_size, padding=padding, bias=False
        )
        self.sig = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.size()
        y = self.avg(x).view([b, 1, c])
        y = self.conv(y)
        y = self.sig(y).view([b, c, 1, 1])
        out = x * y
        return out


model = ECA(8)
model = model.cuda()
input = torch.randn(1, 8, 12, 12).cuda()
output = model(input)
print(output.shape)  # (1, 8, 12, 12)

Guess you like

Origin blog.csdn.net/m0_59967951/article/details/126999735