pytorch实现CBAM

CBAM (Convolutional Block Attention Module) 是一种用于计算卷积层中特征图的注意力权重的方法。在 PyTorch 中实现 CBAM 需要使用自定义层。

首先，需要计算特征图的通道注意力权重和空间注意力权重。通道注意力权重可以使用全连接层和 softmax 函数计算，空间注意力权重可以使用卷积层和 sigmoid 函数计算。

然后，将计算出的两种注意力权重乘在特征图上，得到最终的输出特征图。

示例代码如下：

import torch
import torch.nn as nn

class CBAM(nn.Module):
    def __init__(self, in_channels, reduction_ratio=16):
        super(CBAM, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc1 = nn.Linear(in_channels, in_channels // reduction_ratio)
        self.fc2 = nn.Linear(in_channels // reduction_ratio, in_channels)
        self.conv_after_concat = nn.Conv2d(in_channels * 2, in_channels, kernel_size=1, stride=1)
        
    def forward(self, x):
        # calculate channel attention weight
        avg_out = self.avg_pool(x)
        avg_out = self.fc1(avg_out.view(avg_out.size(0), -1))
        avg_out = nn.ReLU()(avg_out)
        avg_out = self.fc2(avg_out)
        avg_out = nn.Sigmoid()(avg_out)
        
        max_out = self.max_pool(x)
        max_out = self.fc1(max_out.view(max_out.size(0), -1))
        max_out = nn.ReLU()(max_out)
        max_out = self.fc2(max_out)
        max_out = nn.Sigmoid()(max_out)
        
        # calculate spatial attention weight
        x = self.conv_after_con