Tensorflow2.* keras CBAM 代码实现

理论推导

整体模型框架

试验结果显示串联效果优于并联结果，channel attention在前优于spatial attention在前。
因此，最后采取的是一个channel attention module与一个spatial attention module串联的形式。

表达式如下：
$$F^{\prime }=M_c(F)⨂F$$
$$F^{\prime \prime }=M_s(F^{\prime })⨂F^{\prime }$$
式中，
$F\in {\mathbb{R}}^{C\times H\times W}$表示module输入特征层
$M_c\in {\mathbb{R}}^{C\times 1\times 1}$表示1D channel attention
$M_s\in {\mathbb{R}}^{1\times H\times W}$表示2D spatial attention
$⨂$表示同位元素相乘

Channel attention module

1.同时进行了global maxpool和avgpool；
2. 共享权重且具有瓶颈机制两层全连接；
3. 同位相加后采用sigmoid激活，得到输出；

因此对应$M_c(F)$的表达式如下：
$$\begin{array}{rl}{M}_c(F)& =\sigma (MLP(AvgPool(F))+MLP(MaxPool(F))\\ & =\sigma (W_1(W_0(F_{avg}^c))+W_1(W_0(F_{max}^c)))\end{array}$$
式中，
$\sigma$表示sigmoid激活
$F_{avg}^c$和$F_{max}^c$分别表示global avgpooling和maxpooling操作；
$W_0\in {\mathbb{R}}^{\frac{C}{r}\times C}$和$W_1\in {\mathbb{R}}^{C\times \frac{C}{r}}$分别表示全连接权重

Spatial attention module

1. 对输入特征层沿着channels维度分别进行global maxpool和avgpool，并进行堆叠；
2. 采用7x7卷积核进行卷积操作；
3. 采用sigmoid激活，得到输出；
   
   $$\begin{array}{rl}{M}_s(F)& =\sigma (f^{7\times 7}([AvgPool(F);MaxPool(F)]))\\ & =\sigma (f^{7\times 7}([F_{avg}^s;F_{max}^s]))\end{array}$$
   式中：
   $\sigma$表示sigmoid激活
   $F_{avg}^s\in {\mathbb{R}}^{1\times H\times W}$和$F_{max}^s\in {\mathbb{R}}^{1\times H\times W}$分别表示沿着通道维度进行maxpool和avgpool
   $f^{7\times 7}$表示卷积核尺寸为7x7的卷积操作

对比SENet

相比SENet，CBAM的创新之处是在global pooling的时候同时进行了maxpool和avgpool。
原文介绍：
we show that those are suboptimal features in order to infer fine channel attention, and we suggest to use max-pooled features as well;
max-pooling gathers another important clue about distinctive object features to infer finer channel-wise attention;
max-pooled features are as meaningful as average-pooled features, comparing the accuracy improvement from the baseline;

channel pooling produces better accuracy, indicating that explicitly modeled pooling leads to finer attention inference rather than learnable weighted channel pooling

代码复现

基于之前搭建的 Tensorflow2.0 keras ResNet 50 101 152系列 代码实现。模型搭建请参看
Tensorflow 2.0 keras.models.Sequential() Model() 创建网络的若干方式 及共享权重问题

import  tensorflow as tf
from    tensorflow import keras
from    tensorflow.keras import layers, models, Sequential
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, Dense, Flatten,Reshape, Dropout, BatchNormalization, Activation, GlobalAveragePooling2D
from tensorflow.keras.layers import GlobalMaxPool2D, Concatenate

# 继承Layer,建立resnet50 101 152卷积层模块
def conv_block(inputs, filter_num, reduction_ratio, stride=1, name=None):
    
    x = inputs
    x = Conv2D(filter_num[0], (1,1), strides=stride, padding='same', name=name+'_conv1')(x)
    x = BatchNormalization(axis=3, name=name+'_bn1')(x)
    x = Activation('relu', name=name+'_relu1')(x)

    x = Conv2D(filter_num[1], (3,3), strides=1, padding='same', name=name+'_conv2')(x)
    x = BatchNormalization(axis=3, name=name+'_bn2')(x)
    x = Activation('relu', name=name+'_relu2')(x)

    x = Conv2D(filter_num[2], (1,1), strides=1, padding='same', name=name+'_conv3')(x)
    x = BatchNormalization(axis=3, name=name+'_bn3')(x)

    # Channel Attention
    avgpool = GlobalAveragePooling2D(name=name+'_channel_avgpool')(x)
    maxpool = GlobalMaxPool2D(name=name+'_channel_maxpool')(x)
    # Shared MLP
    Dense_layer1 = Dense(filter_num[2]//reduction_ratio, activation='relu', name=name+'_channel_fc1')
    Dense_layer2 = Dense(filter_num[2], activation='relu', name=name+'_channel_fc2')
    avg_out = Dense_layer2(Dense_layer1(avgpool))
    max_out = Dense_layer2(Dense_layer1(maxpool))

    channel = layers.add([avg_out, max_out])
    channel = Activation('sigmoid', name=name+'_channel_sigmoid')(channel)
    channel = Reshape((1,1,filter_num[2]), name=name+'_channel_reshape')(channel)
    channel_out = tf.multiply(x, channel)
    
    # Spatial Attention
    avgpool = tf.reduce_mean(channel_out, axis=3, keepdims=True, name=name+'_spatial_avgpool')
    maxpool = tf.reduce_max(channel_out, axis=3, keepdims=True, name=name+'_spatial_maxpool')
    spatial = Concatenate(axis=3)([avgpool, maxpool])

    spatial = Conv2D(1, (7,7), strides=1, padding='same',name=name+'_spatial_conv2d')(spatial)
    spatial_out = Activation('sigmoid', name=name+'_spatial_sigmoid')(spatial)

    CBAM_out = tf.multiply(channel_out, spatial_out)

    # residual connection
    r = Conv2D(filter_num[2], (1,1), strides=stride, padding='same', name=name+'_residual')(inputs)
    x = layers.add([CBAM_out, r])
    x = Activation('relu', name=name+'_relu3')(x)

    return x

def build_block (x, filter_num, blocks, reduction_ratio=16, stride=1, name=None):

    x = conv_block(x, filter_num, reduction_ratio, stride, name=name)

    for i in range(1, blocks):
        x = conv_block(x, filter_num, reduction_ratio, stride=1, name=name+'_block'+str(i))

    return x


# 创建resnet50 101 152
def SE_ResNet(Netname, nb_classes):

    ResNet_Config = {'ResNet50':[3,4,6,3],
                    'ResNet101':[3,4,23,3],
                    'ResNet152':[3,8,36,3]}
    layers_dims=ResNet_Config[Netname]

    filter_block1=[64, 64, 256]
    filter_block2=[128,128,512]
    filter_block3=[256,256,1024]
    filter_block4=[512,512,2048]

    # Reduction ratio in four blocks
    SE_reduction=[16,16,16,16]

    img_input = Input(shape=(224,224,3))
    # stem block 
    x = Conv2D(64, (7,7), strides=(2,2),padding='same', name='stem_conv')(img_input)
    x = BatchNormalization(axis=3, name='stem_bn')(x)
    x = Activation('relu', name='stem_relu')(x)
    x = MaxPooling2D((3,3), strides=(2,2), padding='same', name='stem_pool')(x)
    # convolution block
    x = build_block(x, filter_block1, layers_dims[0], SE_reduction[0], name='conv1')
    x = build_block(x, filter_block2, layers_dims[1], SE_reduction[1], stride=2, name='conv2')
    x = build_block(x, filter_block3, layers_dims[2], SE_reduction[2], stride=2, name='conv3')
    x = build_block(x, filter_block4, layers_dims[3], SE_reduction[3], stride=2, name='conv4')
    # top layer
    x = GlobalAveragePooling2D(name='top_layer_pool')(x)
    x = Dense(nb_classes, activation='softmax', name='fc')(x)

    model = models.Model(img_input, x, name=Netname)

    return model
    

if __name__=='__main__':
    model = SE_ResNet('ResNet50', 1000)
    model.summary()