空间注意力机制

第一步是沿着通道维度进行最大池化和平均池化，比如下面3*3的特征图，有3个通道。

第二步新特征图进行拼接并经过卷积调整通道数

第三步经过Sigmoid函数后乘到输入上

代码：

class SpatialAttention(layers.Layer):
    def __init__(self):
        super(SpatialAttention, self).__init__()
        self.conv2d = layers.Conv2D(filters=1, strides=1, kernel_size=3, padding="same", activation="sigmoid")

    def call(self, inputs, *args, **kwargs):
        MAX = tf.reduce_max(inputs, axis=-1)#[B,w,h,channel]，沿着最后一个维度：通道，取最大值，得到的维度为【B,w,h】
        MAX = tf.expand_dims(MAX, axis=-1)#增加通道维度，[B,h,w,1]
        print(MAX.shape)

        average = tf.reduce_mean(inputs, axis=-1)#同上述一样
        average = tf.expand_dims(average, axis=-1)

        combine = tf.concat([MAX, average], axis=-1)#[B,h,w,2]
        print(combine.shape)

        x = self.conv2d(combine)#[B,h,w,1]
        return inputs*x

代码实战用法

手写数字识别-卷积神经网络-通道注意力机制-空间注意力机制

SE注意力机制-通道注意力机制