首页 > 其他分享 >CBAM: 卷积注意力模块的学习、实现及其应用

CBAM: 卷积注意力模块的学习、实现及其应用

时间:2022-10-29 16:35:18浏览次数:85  
标签:nn 卷积 CBAM self attention channels stride 模块 注意力

简介

Convolutional Block Attention Module(CBAM), 卷积注意力模块。该论文发表在ECCV2018上(论文地址),这是一种用于前馈卷积神经网络的简单而有效的注意力模块。

CBAM融合了通道注意力(channel Attention)和空间注意力(Spatial Attention),同时该注意力模块非常轻量化,而且能够即插即用,可以用在现存的任何一个卷积神经网络中。

CBAM的流程如上图所示

首先,输入是一个中间特征图,将特征图输入至Channel Attention Module 获取通道注意力,然后将注意力权重作用于中间特征图。

然后,将施加通道注意力的特征图输入至Spatial Attention Module 获取空间注意力,然后将注意力权重作用到特征图上。

最终,经过这两个注意力模块的串行操作,最初的特征图就经过了通道和空间两个注意力机制的处理,自适应细化特征。

那么,CBAM的注意力到底如何计算如上图所示,我们将在下面进行讲解。

 

Channel Attention

通道注意力,将特征图在空间维度池化,保留通道的特征信息。在CBAM中,我们可以看到,全局平均池化和全局最大值池化均有使用(SENet只使用了平均池化,而cbam认为最大值池化被可以捕捉突出特征,实验证明也确实是有效的)。

它的计算流程如下:

  1. 将输入的特征图分别进行全局最大池化和全局平均池化,将空间维度压缩为1,保留通道信息。

  2. 将两个池化后的特征送入共享的多层感知机(MLP)提取特征。

  3. 将经过MLP的池化特征相加,经过sigmoid激活得到最终的通道注意力权重。

 

Spatial Attention

空间注意力,将特征图在通道维度池化,保留空间的特征信息。

它的计算流程如下:

  1. 将特征图(经过通道注意力计算后的)在通道维度分别进行最大值池化和平均池化,将通道维度压缩为1,保留空间信息。

  2. 将池化特征concatenate起来,经过一个卷积层提取特征,同时将通道维度降至1。

  3. 最终经过sigmoid激活,得到最终的空间注意力权重(包含进了通道注意力)。

 

在串行地进行完这两个步骤后,将空间注意力特征与原特征图相乘即可。从上述运算过程可以看出,CBAM作为独立的两个模块,可以直接添加在任何一个卷积神经网络中,带来的附加运算量开销也很小。

实现

接下来,我们用Pytorch实现CBAM。

Channel Attention

class ChannelAttention(nn.Module):
    def __init__(self, channels, reduction_radio=16):
        super().__init__()
        self.channels = channels
        self.inter_channels = self.channels  // reduction_radio
        self.maxpool = nn.AdaptiveMaxPool2d((1, 1))
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
​
        self.mlp = nn.Sequential(  # 使用1x1卷积代替线性层,可以不用调整tensor的形状
            nn.Conv2d(self.channels, self.inter_channels,
                    kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(self.inter_channels),
            nn.ReLU(),
            nn.Conv2d(self.inter_channels, self.channels,
                    kernel_size=1, stride=1, padding=0),
            nn.BatchNorm2d(self.channels)
        )
        self.sigmoid = nn.Sigmoid()
​
    def forward(self, x):  # (b, c, h, w)
        maxout = self.maxpool(x) # (b, c, 1, 1)
        avgout = self.avgpool(x) # (b, c, 1, 1)
​
        maxout = self.mlp(maxout) # (b, c, 1, 1)
        avgout = self.mlp(avgout) # (b, c, 1, 1)
​
        attention = self.sigmoid(maxout + avgout) #(b, c, 1, 1)
​
        return attention

 

Spatial Attention

class SpatialAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(in_channels=2, out_channels=1,
                kernel_size=7, stride=1, padding=3)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, x): # (b, c, h, w)
        maxpool = x.argmax(dim=1, keepdim=True) # (b, 1, h, w)
        avgpool = x.mean(dim=1, keepdim=True)   # (b, 1, h, w)
​
        out = torch.cat([maxpool, avgpool], dim=1) # (b, 2, h, w)
        out = self.conv(out)  # (b, 1, h, w)
​
        attention = self.sigmoid(out) #(b, 1, h, w)
        return attention
​

对于一个特征图X,CBAM的运算结果如下:

ca = ChannelAttention(64)
sa = SpatialAttention()
​
x = torch.randn(3, 64, 56, 56)
​
channel = ca(x)  # (3, 64, 1, 1)
x = channel * x  # (3, 64, 56, 56)
​
spatial = sa(x)  # (3, 1, 56, 56)
x = spatial * x  # (3, 64, 56, 56)

 

应用

在之前的一篇博客里,尝试使用Vision Transformer训练102种鲜花分类。但是因为从头训练Vision Transformer的效果不好,预训练权重又比较难获得,因此总体而言准确度较低。然后这次想用ResNet + CBAM重新试一下。最后也没太使劲调参,发现效果还不错。如何将CBAM插入到ResNet中如下图所示。

 

 

同时想尝试一下,将CBAM的注意力权重提取出来可视化观察一下效果,那么接下来就写一下ResNet + CBAM的实现(以ResNet34为例)。

代码主要参考了这个仓库,不过因为要可视化等等,还是进行了一些修改。github

BasicBlock

class BasicBlock(nn.Module):
    expansion = 1  # 通道升降维倍数
​
    def __init__(self, in_channels, channels, stride=1, downsample=None, attention=None):
        super().__init__()
​
        self.conv1 = nn.Conv2d(in_channels, channels,
                               kernel_size=3, stride=stride, padding=1)  # 第一个卷积层,通过stride进行下采样
        self.bn1 = nn.BatchNorm2d(channels)
        self.conv2 = nn.Conv2d(channels, channels,
                               kernel_size=3, stride=1, padding=1)  # 第二个卷积层,不进行下采样
        self.bn2 = nn.BatchNorm2d(channels)
​
        self.downsample = downsample
        self.attention = attention  # CBAM模块
        self.stride = stride
​
        self.relu = nn.ReLU(inplace=True)
​
    def forward(self, x):
        residual = x
​
        out = self.bn1(self.conv1(x))
        out = self.relu(out)
        out = self.bn2(self.conv2(out))
​
        if self.attention is not None:
            out = self.attention[0](out) * out  # 先进行通道注意力
            self.attention_weights = self.attention[1](out)  # CBAM的注意力图
            out = self.attention_weights * out  # 然后进行空间注意力
        else:
            self.attention_weights = None
​
        if self.downsample is not None:
            residual = self.downsample(x)  # 通道数不变,1x1卷积层仅用于降采样
​
        out += residual
        return self.relu(out)

 

ResNet

class ResNet(nn.Module):
    def __init__(self, block, layers, num_classes=1000):
        self.in_channels = 64
        self.layers = layers
        super().__init__()
​
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
​
        self.layer1 = self._make_layer(block, 64, layers[0])  # 第一个残差层不进行下采样
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
​
        self.attention_layer = [self.layer3, self.layer4]     # 仅在最后两个layer上添加注意力
​
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.flatten = nn.Flatten()
        self.fc = nn.Linear(512 * block.expansion, num_classes)
​
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.xavier_normal_(m.weight, gain=1)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
​
    def _make_layer(self, block, channels, blocks, stride=1):  # block:basicblock or bottleneck
        downsample = None
​
        if stride != 1 or self.in_channels != channels * block.expansion:  # 需要下采样or要融合通道
            downsample = nn.Sequential(
                nn.Conv2d(self.in_channels, channels * block.expansion, kernel_size=1,
                          stride=stride, bias=False),
                nn.BatchNorm2d(channels * block.expansion)
            )
        layers = []
        layers.append(block(self.in_channels, channels, stride, downsample))  # 第一个残差块
​
        self.in_channels = channels * block.expansion
        for i in range(1, blocks):
            attention = None
            if i > 1: # 在第2层往后才添加cbam
                attention = nn.Sequential(
                    ChannelAttention(self.in_channels),
                    SpatialAttention())
              
            layers.append(block(self.in_channels, channels, attention=attention))
​
        return nn.Sequential(*layers)
​
    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
​
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        self._attention_weights = [None] * len(self.attention_layer)  # 将带有注意力层的注意力权重拿出来
​
        for i, layer in enumerate(self.attention_layer):
            for j, (name, blk) in enumerate(layer.named_children()):
                self._attention_weights[i] = blk.attention_weights  # 覆盖,仅获取最后一个block的注意力
​
        x = self.avgpool(x)
        x = self.flatten(x)
        x = self.fc(x)
​
        return x
​
    @property
    def attention_weights(self):
        return self._attention_weights

 

ResNet34 with CBAM

model_urls = {
    'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
    'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
    'resnet50': 'https://download.pytorch.org/models/resnet50-19c8e357.pth',
    'resnet101': 'https://download.pytorch.org/models/resnet101-5d3b4d8f.pth',
    'resnet152': 'https://download.pytorch.org/models/resnet152-b121ed2d.pth',
}
​
​
def resnet34_cbam(pretrained=False, num_class=1000):
​
    model = ResNet(BasicBlock, [3, 4, 6, 3], num_class)
    if pretrained:
        pretrained_state_dict = model_zoo.load_url(model_urls['resnet34'])       # 预训练resnet的权重字典
        pretrained_state_dict = {k: v for k, v in pretrained_state_dict.items()  # 除去全连接层的预训练权重
                                 if (k in pretrained_state_dict and 'fc' not in k)}
​
        new_state_dict = model.state_dict()
        new_state_dict.update(pretrained_state_dict)  # 将预训练权重通过dict的update方式更新
​
        model.load_state_dict(new_state_dict)         # 将更新的网络权重载入到注意力resnet中
​
    return model

 

其他一些杂项

对于102种鲜花分类,只给了训练集和需要打标签的测试集。为了评估训练效果,最好从训练集中分出一部分来进行验证。本人就随机划分了10%的训练集用于验证,剩下的90%进行训练。

在训练的时候,从头开始训练相对来说容易过拟合。因此我们最好是迁移学习,使用自带的预训练权重,然后进行微调。然后把不带CBAM注意力模块的层的参数冻结(前面两层),让网络相对保持一个泛性(毕竟是在ImageNet上预训练的),不然相对而言还是比较容易过拟合。

在含有卷积注意力层的参数,使用较小的学习率进行微调。这样基本不用考虑太多其他的超参数,就可以得到一个很不错的效果(在这种方式下,第一次就得到了94.6%的准确率,比之前高了不少)。

 

整个训练代码以后应该会放在自己的github上,届时会在博客中贴出来。

 

可视化

因为添加了注意力,想尝试将注意力可视化出来观察一下效果。可视化的方法有很多,主要参考了这篇文章(知乎),将CBAM空间注意力的attention map(非特征图)可视化了出来,代码就不再放出来了。

大体效果是这样的,感觉还凑合。红色感兴趣的部分基本上是把这几朵花的大部分都包含进去了。

                             

可能以后再优化优化,用类别激活可视化(Class Activation Mapping,CAM)等方式将整个网络的特征图可视化一下。

 

若本文有错误的地方,欢迎大佬批评指正。

标签:nn,卷积,CBAM,self,attention,channels,stride,模块,注意力
From: https://www.cnblogs.com/Brisling/p/16838987.html

相关文章