SeaFormer：超越TopFormer的轻量级移动端语义分割模型

ICLR2023最新论文，移动端福音！SeaFormer: Squeeze-enhanced Axial Transformer

AI Studio

2892人浏览 · 2023-03-14 14:53:17

AI Studio · 2023-03-14 14:53:17 发布

★★★ 本文源自AlStudio社区精品项目，【点击此处】查看更多精品内容 >>>

SeaFormer: 超越TopFormer的轻量级移动端语义分割模型

１.1 简介

本项目基于PaddlePaddle 复现了ICLR 2023 上发表的论文：SeaFormer: Squeeze-enhanced Axial Transformer for Mobile Semantic Segmentation

背景

Transformer在服务器端的多类别语义分割数据集上均获得了碾压的效果，Topformer第一次在移动端上使用了Transformer的注意力机制增强了特征，但是使用的注意力还是不够高效，因此只能在下采样64倍的分辨率上进行特征增强。现有的高效注意力机制包含基于窗口的注意力机制和轴线注意力等，但是这些注意力在高分辨率上的延时对于移动端来说还是损失了全局信息/不够高效。为此这篇文章设计增加了位置信息的轴线压缩增强注意力，一方面将QKV特征进行轴线压缩后再注意力增强，另一方面将QKV特征使用卷积网络提升局部信息，最后将二者融合，输出增强特征。

网络结构

如上面的框架图所示，该网络先对图像进行1/2、1/4和1/8的下采样，再分别用两个分支进行处理，红色的是上下文分支，蓝色的是空间分支。上下文分支交替使用了MobileNetV2 Block（MV2）和SeaFormer Layers，中间使用Fusion block对两个分支进行融合，使用卷积和Sigmoid提取权重信息，再将权重信息与空间分支相乘，连续迭代三次后使用Light segmentation head进行处理。

SeaFormer Layer的结构如下：

SeaFormer采用了轴向注意力机制，先对Q，K，V分别进行池化，再转换维度，分别使用横向压缩和纵向压缩得到Hx1xCp和1xWxCp的两个向量，对这两个向量分别计算多头自注意力，再将得到的结果进行广播，将两个结果扩大为相同的维度再进行相加。为了增强细节信息，SeaFormer中使用了Detail enhancement kernel，通过该kernel得到的结果与压缩的轴向注意力的计算结果进行相乘，得到最终的注意力。

SeaFormer注意力的核心代码：

class Sea_Attention(paddle.nn.Layer):
    def __init__(self, dim, key_dim, num_heads,
                 attn_ratio=2.,
                 activation=None
                 ):
        super().__init__()
        self.num_heads = num_heads
        self.scale = key_dim ** -0.5
        self.key_dim = key_dim
        self.nh_kd = nh_kd = key_dim * num_heads  # num_head key_dim
        self.d = int(attn_ratio * key_dim)
        self.dh = int(attn_ratio * key_dim) * num_heads
        self.attn_ratio = attn_ratio

        self.to_q = layers.ConvBN(dim, nh_kd, 1, bias_attr=False)
        self.to_k = layers.ConvBN(dim, nh_kd, 1, bias_attr=False)
        self.to_v = layers.ConvBN(dim, self.dh, 1, bias_attr=False)

        self.proj = paddle.nn.Sequential(activation(), layers.ConvBN(
            self.dh, dim, kernel_size=1, bias_attr=False))
        self.proj_encode_row = paddle.nn.Sequential(activation(), layers.ConvBN(
            self.dh, self.dh, kernel_size=1, bias_attr=False))
        self.pos_emb_rowq = SqueezeAxialPositionalEmbedding(nh_kd, 16)
        self.pos_emb_rowk = SqueezeAxialPositionalEmbedding(nh_kd, 16)

        self.proj_encode_column = paddle.nn.Sequential(activation(), layers.ConvBN(
            self.dh, self.dh, kernel_size=1, bias_attr=False))
        self.pos_emb_columnq = SqueezeAxialPositionalEmbedding(nh_kd, 16)
        self.pos_emb_columnk = SqueezeAxialPositionalEmbedding(nh_kd, 16)

        self.dwconv = layers.ConvBN(self.dh + 2 * self.nh_kd, 2 * self.nh_kd + self.dh,
                             kernel_size=3, stride=1, padding=1, dilation=1,
                             groups=2 * self.nh_kd + self.dh, bias_attr=False)
        self.act = activation()
        self.pwconv = layers.ConvBN(2 * self.nh_kd + self.dh, dim, kernel_size=1, bias_attr=False)
        self.sigmoid = HSigmoid()

    def forward(self, x):
        B, C, H, W = x.shape

        q = self.to_q(x)
        k = self.to_k(x)
        v = self.to_v(x)

        # detail enhance
        qkv = paddle.concat([q, k, v], axis=1)
        qkv = self.act(self.dwconv(qkv))
        qkv = self.pwconv(qkv)

        # squeeze axial attention
        ## squeeze row
        qrow = self.pos_emb_rowq(q.mean(-1)).reshape((B, self.num_heads, -1, H)).transpose((0, 1, 3, 2))
        krow = self.pos_emb_rowk(k.mean(-1)).reshape((B, self.num_heads, -1, H))
        vrow = v.mean(-1).reshape((B, self.num_heads, -1, H)).transpose((0, 1, 3, 2))
        attn_row = paddle.matmul(qrow, krow) * self.scale
        attn_row = F.softmax(attn_row, axis=-1)
        xx_row = paddle.matmul(attn_row, vrow)  # B nH H C
        xx_row = self.proj_encode_row(xx_row.transpose((0, 1, 3, 2)).reshape((B, self.dh, H, 1)))

        ## squeeze column
        qcolumn = self.pos_emb_columnq(q.mean(-2)).reshape((B, self.num_heads, -1, W)).transpose((0, 1, 3, 2))
        kcolumn = self.pos_emb_columnk(k.mean(-2)).reshape((B, self.num_heads, -1, W))
        vcolumn = v.mean(-2).reshape((B, self.num_heads, -1, W)).transpose((0, 1, 3, 2))
        attn_column = paddle.matmul(qcolumn, kcolumn) * self.scale
        attn_column = F.softmax(attn_column, axis=-1)
        xx_column = paddle.matmul(attn_column, vcolumn)  # B nH W C
        xx_column = self.proj_encode_column(xx_column.transpose((0, 1, 3, 2)).reshape((B, self.dh, 1, W)))

        xx = xx_row.add(xx_column)
        xx = v.add(xx)
        xx = self.proj(xx)

        xx = self.sigmoid(xx) * qkv
        return xx

1.2 复现精度

原文的测试效果如图。左边为不同分辨率输入下，模型的效率。右边为不同型号的模型在ADE20K数据集上的512x512分辨率训练的性能对比，其中SeaFormer从轻量到重量模型依次为SeaFormer—Tiny, SeaFormer—Small，SeaFormer—Base，SeaFormer—Large。

本项目在数据集的测试效果如下表。

Frame	NetWork	iters	opt	lr	resolution	batch_size	dataset	card	mIoU
本项目Paddle	SeaFormer—Base	160k	AdamW	0.00025	512X512	16	ADE20K	1xV100	39.75

可见，本项目成功用PaddlePaddle复现了论文结果（SeaFormer—Base： 40.2）

快速开始

2.1 数据准备

数据集网站：ADE20K

AiStudio上的数据集：ADE20K

数据集介绍：

ADE20K拥有超过25,000张图像（20 K train，2k val，3k test），覆盖100个thing和50个stuff类别。

%cd ~/data/data26423/
!unzip -oq ade20k.zip
!mv ADEChallengeData2016 ..
%cd ~/

2.2 训练

本项目基于PaddleSeg进行了封装，可以便捷地训练评估。

代码位置位于paddleseg/models/seaformer.py，配置文件位于configs/seaformer文件夹下。

需要注意的是backbone预训练模型是从pytorch转换过来的。在配置yml文件中修改其路径。
我们以SeaFormer-Base为例进行训练：其骨干网络预训练模型为SeaFormer_B_cls_76.4.pdiparams。

转换好的骨干模型预训练模型数据集在：AIStudio上SeaFormer_ImageNet_Pretrained
可以获取。

%cd /home/aistudio/PaddleSeg/
!python setup.py install # 安装自己的PaddleSeg
!python tools/train.py --config configs/seaformer/seaformer_base_ade20k_512x512_160k.yml --do_eval --use_vdl --save_interval 4000 --num_workers 8

2.3 验证

%cd /home/aistudio/PaddleSeg/
!python tools/val.py --config configs/seaformer/seaformer_base_ade20k_512x512_160k.yml --model_path /home/aistudio/PaddleSeg/output/best_model/model.pdparams

2.4 预测

下面以一张图片进行预测：

%cd /home/aistudio/PaddleSeg/
!python tools/predict.py --config configs/seaformer/seaformer_base_ade20k_512x512_160k.yml --model_path /home/aistudio/PaddleSeg/output/best_model/model.pdparams --image_path  /home/aistudio/ADE_val_00000003.jpg

复现心得与相关信息

利用Paddle快速复现新论文的经验分享可以从两个方面来讲：
第一步是熟悉论文的核心思想和参考代码的基本结构和核心代码，对复现的难度等有一个大概的把握。
第二个是快速的代码对齐。这部分主要是需要熟悉不同框架与Paddle的api函数的功能，不熟悉也没关系，可以通过查阅官网的手册和利用X2Paddle提供的对齐文档进行快速上对齐。

相关信息