神经网络学习小记录78——Keras CA（Coordinate attention）注意力机制的解析与代码详解

学习前言

CA注意力机制是最近提出的一种注意力机制，全面关注特征层的空间信息和通道信息。

代码下载

Github源码下载地址为：
https://github.com/bubbliiiing/yolov4-tiny-keras

复制该路径到地址栏跳转。

CA注意力机制的概念与实现

该文章的作者认为现有的注意力机制（如CBAM、SE）在求取通道注意力的时候，通道的处理一般是采用全局最大池化/平均池化，这样会损失掉物体的空间信息。作者期望在引入通道注意力机制的同时，引入空间注意力机制，作者提出的注意力机制将位置信息嵌入到了通道注意力中。

CA注意力的实现如图所示，可以认为分为两个并行阶段：

将输入特征图分别在为宽度和高度两个方向分别进行全局平均池化，分别获得在宽度和高度两个方向的特征图。假设输入进来的特征层的形状为[H, W, C]，在经过宽方向的平均池化后，获得的特征层shape为[H, 1, C]，此时我们将特征映射到了高维度上；在经过高方向的平均池化后，获得的特征层shape为[1, W, C]，此时我们将特征映射到了宽维度上。

然后将两个并行阶段合并，将宽和高转置到同一个维度，然后进行堆叠，将宽高特征合并在一起，此时我们获得的特征层为：[1, H+W, C]，利用卷积+标准化+激活函数获得特征。需要注意的是，这里的卷积通道数一般会小一点，做一个缩放，可以减少参数量。卷积后的特征层的shape为[1, H+W, C/r]，其中r为缩放系数。

之后再次分开为两个并行阶段，再将宽高分开成为：[1, H, C/r]和[1, W, C/r]，之后进行转置。获得两个特征层[H, 1, C/r]和[1, W, C/r]。

然后利用1x1卷积调整通道数后取sigmoid获得宽高维度上的注意力情况，前者在宽上拓展一下，后者在高上拓展一下，然后一起乘上原有的特征就是CA注意力机制。

实现的python代码为：

def ca_block(input_feature, ratio=16, name=""):
	channel = input_feature._keras_shape[-1]
	h		= input_feature._keras_shape[1]
	w		= input_feature._keras_shape[2]
 
	x_h = Lambda(lambda x: K.mean(x, axis=2, keepdims=True))(input_feature)
	x_h = Lambda(lambda x: K.permute_dimensions(x, [0, 2, 1, 3]))(x_h)
	x_w = Lambda(lambda x: K.max(x, axis=1, keepdims=True))(input_feature)
	
	x_cat_conv_relu = Concatenate(axis=2)([x_w, x_h])
	x_cat_conv_relu = Conv2D(channel // ratio, kernel_size=1, strides=1, use_bias=False, name = "ca_block_conv1_"+str(name))(x_cat_conv_relu)
	x_cat_conv_relu = BatchNormalization(name = "ca_block_bn_"+str(name))(x_cat_conv_relu)
	x_cat_conv_relu = Activation('relu')(x_cat_conv_relu)
 
	x_cat_conv_split_h, x_cat_conv_split_w = Lambda(lambda x: tf.split(x, num_or_size_splits=[h, w], axis=2))(x_cat_conv_relu)
	x_cat_conv_split_h = Lambda(lambda x: K.permute_dimensions(x, [0, 2, 1, 3]))(x_cat_conv_split_h)
	x_cat_conv_split_h = Conv2D(channel, kernel_size=1, strides=1, use_bias=False, name = "ca_block_conv2_"+str(name))(x_cat_conv_split_h)
	x_cat_conv_split_h = Activation('sigmoid')(x_cat_conv_split_h)
 
	x_cat_conv_split_w = Conv2D(channel, kernel_size=1, strides=1, use_bias=False, name = "ca_block_conv3_"+str(name))(x_cat_conv_split_w)
	x_cat_conv_split_w = Activation('sigmoid')(x_cat_conv_split_w)
 
	output = multiply([input_feature, x_cat_conv_split_h])
	output = multiply([output, x_cat_conv_split_w])
	return output

注意力机制的应用

注意力机制是一个即插即用的模块，理论上可以放在任何一个特征层后面，可以放在主干网络，也可以放在加强特征提取网络。

由于放置在主干会导致网络的预训练权重无法使用，本文以YoloV4-tiny为例，将注意力机制应用加强特征提取网络上。

如下图所示，我们在主干网络提取出来的两个有效特征层上增加了注意力机制，同时对上采样后的结果增加了注意力机制。

实现代码如下：

attention = [se_block, cbam_block, eca_block, ca_block]

#---------------------------------------------------#
#   特征层->最后的输出
#---------------------------------------------------#
def yolo_body(input_shape, anchors_mask, num_classes, phi = 0):
    inputs = Input(input_shape)
    #---------------------------------------------------#
    #   生成CSPdarknet53_tiny的主干模型
    #   feat1的shape为26,26,256
    #   feat2的shape为13,13,512
    #---------------------------------------------------#
    feat1, feat2 = darknet_body(inputs)
    if phi >= 1 and phi <= 4:
        feat1 = attention[phi - 1](feat1, name='feat1')
        feat2 = attention[phi - 1](feat2, name='feat2')

    # 13,13,512 -> 13,13,256
    P5 = DarknetConv2D_BN_Leaky(256, (1,1))(feat2)
    # 13,13,256 -> 13,13,512 -> 13,13,255
    P5_output = DarknetConv2D_BN_Leaky(512, (3,3))(P5)
    P5_output = DarknetConv2D(len(anchors_mask[0]) * (num_classes+5), (1,1))(P5_output)
    
    # 13,13,256 -> 13,13,128 -> 26,26,128
    P5_upsample = compose(DarknetConv2D_BN_Leaky(128, (1,1)), UpSampling2D(2))(P5)
    if phi >= 1 and phi <= 4:
        P5_upsample = attention[phi - 1](P5_upsample, name='P5_upsample')

    # 26,26,256 + 26,26,128 -> 26,26,384
    P4 = Concatenate()([P5_upsample, feat1])
    
    # 26,26,384 -> 26,26,256 -> 26,26,255
    P4_output = DarknetConv2D_BN_Leaky(256, (3,3))(P4)
    P4_output = DarknetConv2D(len(anchors_mask[1]) * (num_classes+5), (1,1))(P4_output)
    
    return Model(inputs, [P5_output, P4_output])