即插即用的轻量注意力机制

以下是一个即插即用的轻量级注意力机制的示例代码，它可以在CNN中使用：

import torch
import torch.nn as nn

class ECALayer(nn.Module):
    def __init__(self, channel, gamma=2, b=1):
        super(ECALayer, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv2d(channel, channel // gamma, kernel_size=1, bias=False)
        self.sigmoid = nn.Sigmoid()
        self.b = b

    def forward(self, x):
        y = self.avg_pool(x)
        y = self.conv(y)
        y = self.sigmoid(y)
        return x * (self.b + y)

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.ecal1 = ECALayer(64)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.ecal2 = ECALayer(128)

    def forward(self, x):
        x = self.conv1(x)
        x = self.ecal1(x)
        x = self.conv2(x)
        x = self.ecal2(x)
        return x

在这个示例中，ECALayer是一个即插即用的轻量级注意力机制，它可以在CNN中使用。ECALayer包括一个自适应平均池化层，一个卷积层和一个sigmoid激活函数。在前向传递过程中，输入特征图通过自适应平均池化层进行平均池化，然后通过卷积层进行卷积，最后通过sigmoid激活函数进行激活。ECALayer的输出是输入特征图和注意力图的加权和，其中注意力图是通过sigmoid激活函数计算得出的。在CNN中，ECALayer可以嵌入到卷积层之间，以增强CNN的表达能力。