2021刷新COCO和Cityscapes | Polarized Self-Attention:极化自注意力机制(keras实现)

 

2021刷新COCO和Cityscapes | Polarized Self-Attention:极化自注意力机制(keras实现)_第1张图片

Polarized Self-Attention: Towards High-quality Pixel-wise Regression 

paper:https://arxiv.org/pdf/2107.00782.pdf

code:https://github.com/DeLightCMU/PSA

摘要

像素级回归是细粒度计算机视觉任务中最常见的问题,如估计关键点热图和分割掩模。这些回归问题非常具有挑战性,特别是因为它们需要在低计算开销下,建模对高分辨率输入/输出的长期依赖关系来估计高度非线性的像素语义。虽然深度卷积神经网络(DCNNs)中的注意机制已经因增强长期依赖而流行,但元素特定的注意,如非局部块,学习高度复杂和噪声敏感,大多数简化的注意力混合试图在多种类型的任务中达到最佳的折衷。在本文中,提出了极化自注意(PSA)块,它包含了两个高质量像素回归:

(1)极化滤波:在通道和空间注意计算中保持高内部分辨率,同时沿其对应维度完全折叠输入张量。

(2)增强:组成非线性,直接适合典型的细粒度回归的输出分布,如二维高斯分布(关键点热图),或二维二项分布(二进制分割掩模)。

PSA似乎已经耗尽了其仅通道和仅空间分支内的表示能力,因此其串行和并行布局之间只有边际度量差异。

2021刷新COCO和Cityscapes | Polarized Self-Attention:极化自注意力机制(keras实现)_第2张图片

论文背景

注意机制广泛应用于图像分类、目标检测、实例分割。具体来说,注意方法有两种类型,分别是通道注意和空间注意。通道注意力机制如SE、ECANett等,仅通过通道不同权重,但对不同的空间位置施加相同的权重。因此分类任务由于其空间信息最终通过池化而崩溃。随着空间和通道注意力机制的提出,很自然的,结合空间和通道两个维度的双重注意力机制也被提出如DANet、CBAM。该论文提出了一种更加精细的双重注意力机制——极化自注意力(Polarized Self-Attention)。在通道和空间维度保持比较高的resolution(在通道上保持C/2的维度,在空间上保持[H,W]的维度 ),这一步能够减少降维度造成的信息损失,并且采用细粒度回归输出分布的非线性函数。作为一个即插即用的模块,在人体姿态估计和语义分割任务上,作者将它用在了以前的SOTA模型上,并达到了新的SOTA性能。

论文主要思想

2021刷新COCO和Cityscapes | Polarized Self-Attention:极化自注意力机制(keras实现)_第3张图片

提出极化自注意力(PSA),分为两个之路,一个分支做通道维度的自注意力机制,另一个分支做空间维度的自注意力机制,最后将这两个的分支的结果进行融合,得到极化自注意力结构的输出。极化注意力机制有两种模式一种是串联型(mode='s'),另一种是并联型(mode='p')。

通道维度注意力机制

先用了1x1的卷积将输入的特征X转换成了q和v,其中q的通道被完全压缩,而v的通道维度依旧保持在一个比较高的水平(C/2)。因为q的通道维度被压缩,就需要通过HDR进行信息的增强,因此作者用Softmax对q的信息进行了增强。然后将q和k进行矩阵乘法,并在后面接上1x1卷积、LN将通道上C/2的维度升为C(原始代码还加了一个1x1卷积)。最后用Sigmoid函数使得所有的参数都保持在0-1之间。

空间维度注意力机制

作者先用了1x1的卷积将输入的特征转换为了q和v,其中,对于q特征,作者还用了GlobalPooling对空间维度压缩,转换成了1x1的大小;而v特征的空间维度则保持在一个比较大的水平(HxW)。由于q的空间维度被压缩了,所以作者就用了Softmax对q的信息进行增强。然后将q和K进行矩阵乘法,然后接上reshape和Sigmoid使得所有的参数都保持在0-1之间。

Keras实现

以下是根据论文和pytorch源码实现的keras版本(支持Tensorflow1.x)。

注:安装LayerNormlization模块 pip install keras-layer-normalization

PSA模块

def _PSA(x, mode='p'):
    context_channel = spatial_pool(x, mode)
    if mode == 'p':
        context_spatial = channel_pool(x)
        out = Add()([context_spatial, context_channel])
    elif mode == 's':
        out = channel_pool(context_channel)
    else:
        out = x
    return out

通道注意力机制 


def spatial_pool(x, mode, ratio=4):
    channel_axis = 1 if K.image_data_format() == 'channels_first' else -1

    if channel_axis == -1:
        batch, height, width, channels = K.int_shape(x)
        assert channels % 2 == 0
        channel = channels//2
        input_x = Conv2D(channel, kernel_size=1, strides=1, padding='same', use_bias=False, kernel_initializer='he_normal')(x)
        input_x = Reshape((width*height, channel))(input_x)

        context_mask = Conv2D(1, kernel_size=1, strides=1, padding='same', use_bias=False, kernel_initializer='he_normal')(x)
        context_mask = Reshape((width*height, 1))(context_mask)
        context_mask = Softmax(axis=1)(context_mask)
        context = Lambda(lambda x: tf.matmul(x[0], x[1], transpose_a=True))([input_x, context_mask])
        context = Permute((2, 1))(context)
        context = Reshape((1, 1, channel))(context)

    else:
        batch, channels, height, width = K.int_shape(x)
        assert channels % 2 == 0
        channel = channels // 2
        input_x = Conv2D(channel, kernel_size=1, strides=1, padding='same', use_bias=False,
                         kernel_initializer='he_normal')(x)
        input_x = Reshape((channel, width * height))(input_x)

        context_mask = Conv2D(1, kernel_size=1, strides=1, padding='same', use_bias=False,
                              kernel_initializer='he_normal')(x)
        context_mask = Reshape((width * height, 1))(context_mask)
        context_mask = Softmax(axis=1)(context_mask)
        context = Lambda(lambda x: tf.matmul(x[0], x[1]))([input_x, context_mask])
        context = Reshape((channel, 1, 1))(context)

    if mode == 'p':
        context = Conv2D(channels, kernel_size=1, strides=1, padding='same')(context)
    else:
        context = Conv2D(channel // ratio, kernel_size=1, strides=1, padding='same')(context)
        context = LayerNormalization()(context) # pip install keras-layer-normalization
        context = Conv2D(channels, kernel_size=1, strides=1, padding='same')(context)

    mask_ch = Activation('sigmoid')(context)
    return Multiply()([x, mask_ch])

空间注意力机制 


def channel_pool(x):
    channel_axis = 1 if K.image_data_format() == 'channels_first' else -1
    if channel_axis == -1:
        batch, height, width, channels = K.int_shape(x)
        assert channels % 2 == 0
        channel = channels//2
        g_x = Conv2D(channel, kernel_size=1, strides=1, padding='same', use_bias=False, kernel_initializer='he_normal')(x)
        avg_x = GlobalAvgPool2D()(g_x)
        avg_x = Softmax()(avg_x)
        avg_x = Reshape((channel, 1))(avg_x)

        theta_x = Conv2D(channel, kernel_size=1, strides=1, padding='same', use_bias=False, kernel_initializer='he_normal')(x)
        theta_x = Reshape((height*width, channel))(theta_x)
        context = Lambda(lambda x: tf.matmul(x[0], x[1]))([theta_x, avg_x])
        context = Reshape((height*width,))(context)
        mask_sp = Activation('sigmoid')(context)
        mask_sp = Reshape((height, width, 1))(mask_sp)
    else:
        batch, channels, height, width = K.int_shape(x)
        assert channels % 2 == 0
        channel = channels//2
        g_x = Conv2D(channel, kernel_size=1, strides=1, padding='same', use_bias=False, kernel_initializer='he_normal')(x)
        avg_x = GlobalAvgPool2D()(g_x)
        avg_x = Softmax()(avg_x)
        avg_x = Reshape((1, channel))(avg_x)

        theta_x = Conv2D(channel, kernel_size=1, strides=1, padding='same', use_bias=False, kernel_initializer='he_normal')(x)
        theta_x = Reshape((channel, height*width))(theta_x)
        context = Lambda(lambda x: tf.matmul(x[0], x[1]))([avg_x, theta_x])
        context = Reshape((height*width,))(context)
        mask_sp = Activation('sigmoid')(context)
        mask_sp = Reshape((1, height, width))(mask_sp)

    return Multiply()([x, mask_sp])

声明:本内容来源网络,版权属于原作者,图片来源原论文。如有侵权,联系删除。

创作不易,欢迎大家点赞评论收藏关注!(想看更多最新的注意力机制文献欢迎关注浏览我的博客)

你可能感兴趣的:(深度学习之注意力机制)