注意力模型CBAM模块的Keras代码实现

1.啥是CBAM？

CBAM就是结合了通道注意力和空间注意力的一种注意力结构，与SE模块相比，多了空间注意力！

2.CBAM的结构图

如图，整体结构就是先对特征图进行通道注意力加权，然后再进行空间注意力加权操作，很简单。

2.1 CBAM的通道注意力模块

如图，先对输入特征图Input_feature（H×W×C）分别进行全局平均池化和全局平均池化得到两个向量M（1×1×C）和A（1×1×C），在将这俩分别进行两次全连接操作第一次全连接压缩通道为C/r（r自己调整），第二次全连接扩张通道为C。经过两次全连接后得到两个二维向量M*（1×1×C）和A*（1×1×C），将M与A进行Add操作，加到一起得到Ad（1×1×C），而这个Ad就是我们的通道注意力的权重啦！把这个Ad与输入特征图Input_feature逐像素相乘，最后就得到我们经过通道注意力加权后的特征图Channel_feature（H×W×C）。
2.2 CBAM的空间注意力模块

如图，将上一步输出的Channel_feature（H×W×C）分别进行通道维度的平均池化和最大池化。（有同学想问什么是通道维度的平均池化和最大池化呢？我只能说，你看图，横着池化，我很难描述清楚。）得到的是两个压缩后的特征图M（H×W×1）和A（H×W×1）将这俩特征图叠加，也就是Concatenate操作，得到C（H×W×2），将C进行一次卷积核为1的卷积操作得到C*（H×W×2）。将C*与输入的Channel_feature相乘后就得到我们的CBAM特征图啦。

3.代码实现

import keras
from keras.models import *
from keras.layers import *
from keras import layers
import keras.backend as K

IMAGE_ORDERING = 'channels_last'
def CBAM_block(cbam_feature,ratio=8):
    cbam_feature = channel_attention(cbam_feature, ratio)
    cbam_feature = spatial_attention(cbam_feature)
    return cbam_feature

def channel_attention(input_feature,ratio):

    channel = input_feature._keras_shape[-1]

    shared_layer_one = Dense(channel // ratio,
                             activation='relu',
                             kernel_initializer='he_normal',
                             use_bias=True,
                             bias_initializer='zeros')
    shared_layer_two = Dense(channel,
                             kernel_initializer='he_normal',
                             use_bias=True,
                             bias_initializer='zeros')
    avg_pool = GlobalAveragePooling2D()(input_feature)
    avg_pool = shared_layer_one(avg_pool)
    avg_pool = shared_layer_two(avg_pool)

    max_pool = GlobalMaxPool2D()(input_feature)
    max_pool = shared_layer_one(max_pool)
    max_pool = shared_layer_two(max_pool)

    cbam = Add()([avg_pool,max_pool])
    cbam = Activation('sigmoid')(cbam)

    return multiply([input_feature,cbam_feature])

def spatial_attention(input_feature):

    avg_pool = Lambda(lambda x:K.mean(x,axis=3,keepdims=True))(input_feature)
    max_pool = Lambda(lambda x:K.max(x,axis=3,keepdims=True))(input_feature)

    concat = Concatenate(axis=3)([avg_pool,max_pool])
    cbam_feature = Conv2D(1,(7,7),strides=1,padding='same',activation='sigmoid')(concat)

    return multiply([input_feature,cbam_feature])

4.怎么用呢

和SE结构相似，想怎么加怎么加，反正不改变输入特征图的尺寸！深度学习还是试了才知道效果，自由度很高。你也可以只进行平均池化（我想最大池化可能有噪声点的影响。）我只是有这个想法但还没进行试验。各位试过的可以在评论区讨论鸭！

有什么不懂的就在评论区提出！