注意力模块用在CNN上注意力模块怎么玩

转载

mob6454cc7c268c 2024-06-15 13:09:41

文章标签 注意力模块用在CNN上深度学习池化全连接 文章分类 机器学习人工智能

1.啥是CBAM？

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

2.CBAM的结构图

注意力模块用在CNN上注意力模块怎么玩_注意力模块用在CNN上

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

注意力模块用在CNN上注意力模块怎么玩_深度学习_02

如图，先对输入特征图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的空间注意力模块

注意力模块用在CNN上注意力模块怎么玩_池化_03

如图，将上一步输出的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])