注意力机制——CBAM原理详解及源码解析

注意力机制——CBAM原理详解及源码解析

引用

CSDN

1.

https://blog.csdn.net/qq_47233366/article/details/137103025

注意力机制——CBAM原理详解及源码解析

CBAM原理详解

CBAM全称为Convolutional Block Attention Module，可以翻译成卷积块的注意模型。其整体结构如下图所示：

从图中可以看出，CBAM主要由两个模块组成：通道注意力模块（Channel Attention Module）和空间注意力模块（Spatial Attention Module）。输入特征图首先经过这两个模块的处理，最终得到精制的特征图（Refined Feature）。

Channel Attention Module

通道注意力模块的结构如下图所示：

具体步骤如下：

1. 对输入特征图F分别进行全局最大池化和全局平均池化，得到两个1 × 1 × C的特征图。
2. 将这两个特征图送入两个全连接层（MLP）中，输出两个1 × 1 × C的特征图。
3. 将两个特征图相加并经过sigmoid激活函数，得到通道注意力M_c，其尺寸为1 × 1 × C。

这个过程可以用以下公式表示：

接下来，将输入特征图F与通道注意力M_c相乘，得到特征图F'，其尺寸与F一致，都是H × W × C。

Spatial Attention Module

空间注意力模块的结构如下图所示：

具体步骤如下：

1. 对特征图F'分别进行全局最大池化和全局平均池化，得到两个H × W × 1的特征图。
2. 将这两个特征图在channel维度拼接，得到H × W × 2的特征图。
3. 对拼接后的特征图进行卷积操作，得到H × W × 1的特征图。
4. 经过sigmoid激活函数，得到空间注意力M_s，其尺寸为H × W × 1。

这个过程可以用以下公式表示：

最后，将特征图F'与空间注意力M_s相乘，得到最终的输出结果，其尺寸为H × W × C。

CBAM代码详解

下面是CBAM的PyTorch实现代码：

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

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu(self.fc1(self.max_pool(x))))
        out = self.sigmoid(avg_out + max_out)
        return out

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1
        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x = torch.cat([avg_out, max_out], dim=1)
        out = self.sigmoid(self.conv1(x))
        return out

class CBAM(nn.Module):
    def __init__(self, in_planes, ratio=1, kernel_size=7):
        super(CBAM, self).__init__()
        self.channel_att = ChannelAttention(in_planes, ratio)
        self.spatial_att = SpatialAttention(kernel_size)

    def forward(self, x):
        out = self.channel_att(x) * x
        print(self.channel_att(x).shape)
        print(f"channel Attention Module:{out.shape}")
        out = self.spatial_att(out) * out
        print(self.spatial_att(out).shape)
        #print(f"Spatial Attention Module:{out.shape}")
        return out

if __name__ == '__main__':
    # Testing
    model = CBAM(3)
    input_tensor = torch.ones((1, 3, 224, 224))
    output_tensor = model(input_tensor)
    print(f'Input shape: {input_tensor.shape})')
    print(f'Output shape: {output_tensor.shape}')

实验结果如下：

从图中可以看出：

• 第一行表示理论部分的特征图M_c
• 第二行表示F'
• 第三行表示特征图M_s
• 第四行表示输入
• 第五行表示CBAM的最终输出

小结

本文详细介绍了CBAM注意力机制的原理和实现，包括通道注意力模块和空间注意力模块的结构和代码实现。CBAM作为一种常用的注意力机制，在实际项目中应用广泛，希望本文能帮助读者更好地理解和应用CBAM。