CBAM(基于卷积块的注意机制): Convolutional Block Attention Module

 1.介绍

  CBAM的中文名字是基于卷积块的注意机制，从结构上来看，它结合了空间注意力机制和通道注意力机制，从效果上来看，它能提高分类和检测的正确率。更加详细的内容可以参加论文：CBAM: Convolutional Block Attention Module

  

2.模型结构

     对于以下结构，我想用比较具体的数字来表示整个流程，比如输入的Feature维度为1*16*10*10（各维度代表BCHW），经过通道注意力机制之后，得到的权重维度为1*16*1*1（权重值越大，代表着那个通道特征图越重要），经过空间注意力机制之后，得到的权重为1*1*10*10（权重越大，代表特征图上的内容越重要）

    通道注意力机制：在下图中，MaxPool的操作就是提取一副特征图上的最大值，有多少通道就提取多少个；AvgPool的操作就是提取一副特征图上的平均值，也是有多少通道就提取多少个；MLP是一个共享（共用，只有一个）的全连接操作，后面的操作比较简单，就不介绍了。

    空间注意力机制：在下图中，MaxPool的操作就是在通道上提取最大值，提取的次数是高乘以宽；AvgPool的操作就是在通道上提取平均值，提取的次数也是是高乘以宽；接着将前面所提取到的特征图（通道数都为1）合并得到一个2通道的特征图，后面的操作比较简单，就不介绍了。

3.模型特点

      1.在通道注意力机制中引入全连接，并通过全连接降维，有利于提取更重要的信息（相当于PCA操作）

      2.一般注意力机制是用AvgPool操作，本文引入MaxPool操作，可以提取到不同信息，增强特征的多样性。

 4.代码实现Pytorch

# -*-coding:utf-8-*-
import math

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


class Flatten(nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)


class ChannelGate(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16):
        super(ChannelGate, self).__init__()
        self.gate_channels = gate_channels
        self.mlp = nn.Sequential(
            Flatten(),
            nn.Linear(gate_channels, gate_channels // reduction_ratio),
            nn.ReLU(),
            nn.Linear(gate_channels // reduction_ratio, gate_channels)
        )

    def forward(self, x):
        avg_pool = F.avg_pool2d(
            x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
        channel_avg = self.mlp(avg_pool)
        max_pool = F.max_pool2d(
            x, (x.size(2), x.size(3)), stride=(x.size(2), x.size(3)))
        channel_max = self.mlp(max_pool)

        channel_att_sum = channel_avg + channel_max

        scale = torch.sigmoid(channel_att_sum).unsqueeze(
            2).unsqueeze(3).expand_as(x)
        return x * scale


class ChannelPool(nn.Module):
    def forward(self, x):
        return torch.cat((torch.max(x, 1)[0].unsqueeze(1), torch.mean(x, 1).unsqueeze(1)), dim=1)


class SpatialGate(nn.Module):
    def __init__(self):
        super(SpatialGate, self).__init__()
        self.compress = ChannelPool()
        self.spatial = torch.nn.Conv2d(2, 1, kernel_size=(7, 7), stride=(1, 1), padding=(3, 3))

    def forward(self, x):
        x_compress = self.compress(x)
        x_out = self.spatial(x_compress)
        scale = torch.sigmoid(x_out)  # broadcasting
        return x * scale


class CBAM(nn.Module):
    def __init__(self, gate_channels, reduction_ratio=16, pool_types=['avg', 'max']):
        super(CBAM, self).__init__()
        self.ChannelGate = ChannelGate(
            gate_channels, reduction_ratio, pool_types)
        self.SpatialGate = SpatialGate()

    def forward(self, x):
        x_out = self.ChannelGate(x)
        x_out = self.SpatialGate(x_out)
        return x_out