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

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

引用

CSDN

1.

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

CBAM（Convolutional Block Attention Module）是一种常见的注意力机制，通过通道注意力模块和空间注意力模块来增强特征表示。本文将详细介绍CBAM的原理，并提供完整的PyTorch实现代码。

CBAM原理详解

CBAM的整体结构如下图所示：

首先，输入特征图会分别送入通道注意力模块和空间注意力模块，最终得到精制的特征图。

Channel Attention Module

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

具体步骤如下：

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

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的工作原理，并掌握其在PyTorch中的实现方法。