您的位置:

CBAM:一个深度学习模型中的重要改进

一、CBAM代码解析

CBAM是一个基于注意力机制的深度学习模型改进方法,其全称为“Convolutional Block Attention Module”(卷积块注意力模块)。CBAM在卷积神经网络(CNN)中引入了两个注意力机制,分别是通道注意力和空间注意力。关于CBAM代码的解析,我们可以看一下其代码实现:

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.relu1 = 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.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(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)
        x = self.conv1(x)
        return self.sigmoid(x)

上述代码中,ChannelAttention是用于通道注意力机制的类,SpatialAttention是用于空间注意力机制的类。在ChannelAttention类中,我们首先使用两个自适应池化层avg_pool和max_pool来对输入x进行平均值和最大值的计算。我们使用1x1的卷积层fc1和fc2对结果进行特征提取和映射,并使用ReLU激活函数对映射后的结果进行非线性处理。最后,我们使用Sigmoid函数将特征图映射到0到1的区间中,从而生成通道注意力图。

在对空间注意力机制进行处理时,我们使用一个卷积层对空间特征进行处理,然后使用Sigmoid激活函数将特征图映射到0到1的范围内。最终,我们将两个注意力机制的输出进行相乘,然后将结果与原始输入进行加和操作。这个加和操作中的权重由两个注意力机制的相乘结果确定。这样就完成了CBAM在深度学习模型中的引入。

二、CBAM代码实现

下面是一个使用CBAM的PyTorch代码实现,该模型可以用于进行图像分类:

import torch.nn as nn
from cbam import ChannelAttention, SpatialAttention

class CBAMBlock(nn.Module):
    def __init__(self, in_channels, ratio=16, kernel_size=7):
        super(CBAMBlock, self).__init__()
        self.ca = ChannelAttention(in_channels, ratio)
        self.sa = SpatialAttention(kernel_size)

    def forward(self, x):
        out = x * self.ca(x)
        out = out * self.sa(out)
        return out

class CBAMModel(nn.Module):
    def __init__(self, num_classes=10):
        super(CBAMModel, self).__init__()

        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.cbam1 = CBAMBlock(64)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(128)
        self.cbam2 = CBAMBlock(128)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn3 = nn.BatchNorm2d(256)
        self.cbam3 = CBAMBlock(256)
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(256, num_classes)

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.cbam1(out)
        out = F.relu(self.bn2(self.conv2(out)))
        out = self.cbam2(out)
        out = F.relu(self.bn3(self.conv3(out)))
        out = self.cbam3(out)
        out = self.avgpool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out

上述代码中,我们可以看到CBAM的实现方式与传统的卷积神经网络十分相似。我们首先使用三个卷积层和三个批量归一化层来进行特征提取和映射。然后,我们将每个卷积层的输出连接到一个对应的CBAMBlock中进行注意力机制处理。最后,我们使用一个自适应池化层和全连接层对最终特征提取结果进行分类。

三、CBAM代码选取

我们选取几个与CBAM代码相关的示例,以帮助更好地理解CBAM的实现。首先,我们来看一下CBAM模块中通道注意力机制的相关代码:

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.relu1 = 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.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)

上述代码中,我们在通道注意力机制中使用一个自适应池化层avg_pool和max_pool来对输入x进行平均值和最大值的计算。我们使用1x1的卷积层fc1和fc2对结果进行特征提取和映射,并使用ReLU激活函数对映射后的结果进行非线性处理。最后,我们使用Sigmoid函数将特征图映射到0到1的区间中,从而生成通道注意力图。

接着,我们来看一下CBAM模块中空间注意力机制的相关代码:

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)
        x = self.conv1(x)
        return self.sigmoid(x)

上述代码中,我们使用一个卷积层对空间特征进行处理,然后使用Sigmoid激活函数将特征图映射到0到1的范围内。最终,我们将两个注意力机制的输出进行相乘,然后将结果与原始输入进行加和操作。这个加和操作中的权重由两个注意力机制的相乘结果确定。这样就完成了CBAM在深度学习模型中的引入。

最后,我们来看一下使用CBAM模块的PyTorch代码示例:

class CBAMBlock(nn.Module):
    def __init__(self, in_channels, ratio=16, kernel_size=7):
        super(CBAMBlock, self).__init__()
        self.ca = ChannelAttention(in_channels, ratio)
        self.sa = SpatialAttention(kernel_size)

    def forward(self, x):
        out = x * self.ca(x)
        out = out * self.sa(out)
        return out

上述代码中,我们可以看到CBAM的具体使用方式,即将CBAMBlock插入到卷积神经网络的特征提取部分中,用于引入通道注意力和空间注意力机制。在CBAMBlock中,我们使用ChannelAttention和SpatialAttention分别进行通道注意力和空间注意力的计算,并将其相乘后再与原始输入相加,生成输出。

结语

本文详细地阐述了CBAM在深度学习模型中的重要改进,分别对CBAM的代码解析、代码实现和代码选取进行了详细的介绍。相信通过对CBAM深入了解,可以对深度学习模型的构建和优化有更深刻的认识和理解。