注意力机制 python

CBAM是一种用于计算卷积层中特征图注意力权重的方法。在PyTorch中实现CBAM需要使用自定义层。可以使用以下代码作为参考:

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):
        # 计算通道注意力权重
        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)