cbam注意力机制pytorch

CBAM（Convolutional Block Attention Module）是一种针对卷积神经网络（CNN）的注意力机制，在PyTorch中可以使用以下代码实现：

import torch
import torch.nn as nn
import torch.nn.functional as F

class CBAM_Module(nn.Module):

    def __init__(self, channels, reduction=16):
        super(CBAM_Module, self).__init__()
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)

        self.fc1 = nn.Linear(channels, channels // reduction, bias=False)
        self.relu = nn.ReLU(inplace=True)
        self.fc2 = nn.Linear(channels // reduction, channels, bias=False)

        self.sigmoid_channel = nn.Sigmoid()
        self.sigmoid_spatial = nn.Sigmoid()

        self.conv = nn.Conv2d(2, 1, kernel_size=7, stride=1, padding=3)

    def forward(self, x):
        max_out = self.max_pool(x)
        avg_out = self.avg_pool(x)
        max_out = max_out.view(max_out.size(0), -1)
        avg_out = avg_out.view(avg_out.size(0), -1)

        out = torch.cat((max_out, avg_out), dim=1)
        out = self.fc1(out)
        out = self.relu(out)
        out = self.fc2(out)

        channel_attention = self.sigmoid_channel(out)
        channel_attention = channel_attention.view(channel_attention.size(0), channel_attention.size(1), 1, 1)
        x = x * channel_attention

        spatial_attention = torch.mean(x, dim=1, keepdim=True)
        spatial_attention = self.conv(spatial_attention)
        spatial_attention = self.sigmoid_spatial(spatial_attention)
        x = x * spatial_attention
        return x

这个模块包含两个部分：通道注意力和空间注意力。通道注意力用来对每个通道的特征进行加权，从而强化重要的特征，减弱不重要的特征。空间注意力则用来对特征图中的每个位置进行加权，以便在不同的位置上聚焦于重要的特征。

通过在CNN中使用CBAM模块，可以提高模型的性能和准确性。