语义分割之SegFormer分享

今年可以说是分割算法爆发的一年,首先Vit通过引入transform将ADE20K mIOU精度第一次刷到50%,超过了之前HRnet+OCR效果,然后再是Swin屠榜各大视觉任务,在分类,语义分割和实例分割都做到了SOTA,斩获ICCV2021的bset paper,然后Segformer有凭借对transform再次深层次优化,在拿到更高精度的基础之上还大大提升了模型的实时性。

代码:https://github.com/NVlabs/SegFormer
论文:https://arxiv.org/abs/2105.15203

一般的语义分割算法,endoer一般会输出高分辨率的粗粒度特征和低分辨率的细粒度特征,好处是像素分类更准以及边缘更精细。

Vit是最早将transform引入语义分割的一批算法。在ADE20K等公开数据集上相比之前基于CNN的一票分割算法在精度层面有着明显优势,但是也还有很多缺点:

1.ViT-large 参数和计算量非常大,有300M+参数

2.ViT的结构不太适合做语义分割,只能输出固定分辨率的feature map, 比如1/16, 过低分辨率对轮廓要求比较精细的场景太不友好

3.ViT的柱状结构n2复杂度,对显存的负担非常大

4.位置编码, 但是语义分割在测试的时候往往图片的分辨率不是固定的,这时要么对positional embedding做双线性插值,这会损害性能, 要么滑窗,这样效率很低而且很不灵活

语义分割之SegFormer分享_第1张图片

在这里插入图片描述

改进:提出了

Efficient Self-Attention:

语义分割之SegFormer分享_第2张图片
这里直接上代码:
这是segformer的self attention实现:
EfficientMultiheadAttention(MultiheadAttention)
注意这里的qkv计算实际上是继承的pytorch的
nn.MultiheadAttention计算

    def forward(self, x, hw_shape, identity=None):

        x_q = x
        if self.sr_ratio > 1:
            x_kv = nlc_to_nchw(x, hw_shape)
            x_kv = self.sr(x_kv)
            x_kv = nchw_to_nlc(x_kv)
            x_kv = self.norm(x_kv)
        else:
            x_kv = x

        if identity is None:
            identity = x_q
        out = self.attn(query=x_q, key=x_kv, value=x_kv, need_weights=False)[0]

        return identity + self.dropout_layer(self.proj_drop(out))

可以看到作者提出来的Efficient Self-Attention本质其实是增加了一个sr_ratio超参,通过sr_ratio来控制KV参数矩阵的尺寸,具体实现是这样的:

self.sr_ratio = sr_ratio
        if sr_ratio > 1:
            self.sr = Conv2d(
                in_channels=embed_dims,
                out_channels=embed_dims,
                kernel_size=sr_ratio,
                stride=sr_ratio)

Overlapped Patch Merging
这种直接分割方式不能保证这些patch的局部连续性

优化:

作者设置了超参S
将切割后的小图按照一定步长(S)
合并,通过控制小图大小/padding大小/步长产生没有Overlap时的大小
语义分割之SegFormer分享_第3张图片

Mix-FFN

ViT使用位置编码(PE)来介绍位置信息。然而,PE的分辨率是固定的。因此,当测试分辨率不同于训练率时,位置代码需要进行插值,这通常会导致精度下降。作者认为,位置编码实际上并不是语义分割的必要条件,Segform引入了Mix-FFN,它通过直接在前馈网络(FFN)中使用3×3Conv传递位置信息

class MixFFN(BaseModule):
    def __init__(self,
                 embed_dims,
                 feedforward_channels,
                 act_cfg=dict(type='GELU'),
                 ffn_drop=0.,
                 dropout_layer=None,
                 init_cfg=None):
        super(MixFFN, self).__init__(init_cfg)
        self.embed_dims = embed_dims
        self.feedforward_channels = feedforward_channels
        self.act_cfg = act_cfg
        self.activate = build_activation_layer(act_cfg)
        in_channels = embed_dims
        fc1 = Conv2d(
            in_channels=in_channels,
            out_channels=feedforward_channels,
            kernel_size=1,
            stride=1,
            bias=True)
        # 3x3 depth wise conv to provide positional encode information
        pe_conv = Conv2d(
            in_channels=feedforward_channels,
            out_channels=feedforward_channels,
            kernel_size=3,
            stride=1,
            padding=(3 - 1) // 2,
            bias=True,
            groups=feedforward_channels)
        fc2 = Conv2d(
            in_channels=feedforward_channels,
            out_channels=in_channels,
            kernel_size=1,
            stride=1,
            bias=True)
        drop = nn.Dropout(ffn_drop)
        layers = [fc1, pe_conv, self.activate, drop, fc2, drop]
        self.layers = Sequential(*layers)
        self.dropout_layer = build_dropout(
            dropout_layer) if dropout_layer else torch.nn.Identity()
    def forward(self, x, hw_shape, identity=None):
        out = nlc_to_nchw(x, hw_shape)
        out = self.layers(out)
        out = nchw_to_nlc(out)
        if identity is None:
            identity = x
        return identity + self.dropout_layer(out)

上面代码本质上就是利用3X3的conv学习位置信息

这里注意一下,在transform算法中,激活函数不再是我们所熟悉的Relu,而是Gelu,这个也不是说随便就换上去的。

在神经网络的建模过程中,模型很重要的性质就是非线性,同时为了模型泛化能力,需要加入随机正则,例如dropout(随机置一些输出为0,其实也是一种变相的随机非线性激活), 而随机正则与非线性激活是分开的两个事情, 而其实模型的输入是由非线性激活与随机正则两者共同决定的。

GELUs正是在激活中引入了随机正则的思想,是一种对神经元输入的概率描述,直观上更符合自然的认识,同时实验效果要比Relus与ELUs都要好。

GELU也会为inputs乘以0或者1,但不同于以上的或有明确值或随机,GELU所加的0-1mask的值是随机的,同时是依赖于inputs的分布的。可以理解为:GELU的权值取决于当前的输入input有多大的概率大于其余的inputs

因为神经元的输入趋向于正态分布,这么设定使得当输入x减小的时候,输入会有一个更高的概率被dropout掉,这样的激活变换就会随机依赖于输入了。

在这里插入图片描述
论文也给了一种近似表示:
在这里插入图片描述

bert源码:

def gelu(input_tensor):
	cdf = 0.5 * (1.0 + tf.erf(input_tensor / tf.sqrt(2.0)))
	return input_tesnsor*cdf

在transform中确实采用GELU效果会更理想一些
语义分割之SegFormer分享_第4张图片

语义分割之SegFormer分享_第5张图片
差不多4个点左右的提升

encoder总结:
语义分割之SegFormer分享_第6张图片
过去三年语义分割从Deeplab系列到PSPNet到DANet等等都是在研究如何设计
更好的decoder(encoder一般通过backbone提取)decoder越来越重越来越复杂

对于语义分割来说最重要的问题就是如何增大感受野,。首先对于CNN encoder来说,有效感受野是比较小且局部的,所以需要一些decoder 的设计来增大有效感受野,比如ASPP里利用了不同大小的空洞卷积来实现这一目的。
但是对于Transformer encoder来说,由于 self-attention存在,有效感受野变得非常大,因此decoder 不需要更多操作来提高感受野(作者试了一堆分割头,基本没有提升),下面是deeplab和segformer有效感受野可视化的对比(有效感受野:
Understanding the Effective Receptive Field in Deep Convolutional Neural Networks)

语义分割之SegFormer分享_第7张图片
由于自注意力机制的存在,segformer encoder阶段感受野就足够大了
,所以decoder不需要很重的head

首先对不同层的Feature分别过一个linear层确保他们的channel维度一样,其次都上采样到1/4分辨率并concat起来,再用一个linear层融合,最后一个linear层预测结果。整个decoder只有四部分一共6个linear层
没有dilate conv
没有3x3 conv.

所以参数非常少

语义分割之SegFormer分享_第8张图片

class SegformerHead(BaseDecodeHead):
    def __init__(self, interpolate_mode='bilinear', **kwargs):
        super().__init__(input_transform='multiple_select', **kwargs)
        self.interpolate_mode = interpolate_mode
        num_inputs = len(self.in_channels)
        assert num_inputs == len(self.in_index)
        self.convs = nn.ModuleList()
        for i in range(num_inputs):
            self.convs.append(
                ConvModule(
                    in_channels=self.in_channels[i],
                    out_channels=self.channels,
                    kernel_size=1,
                    stride=1,
                    norm_cfg=self.norm_cfg,
                    act_cfg=self.act_cfg))
        self.fusion_conv = ConvModule(
            in_channels=self.channels * num_inputs,
            out_channels=self.channels,
            kernel_size=1,
            norm_cfg=self.norm_cfg)
    def forward(self, inputs):
        # Receive 4 stage backbone feature map: 1/4, 1/8, 1/16, 1/32
        inputs = self._transform_inputs(inputs)
        outs = []
        for idx in range(len(inputs)):
            x = inputs[idx]
            conv = self.convs[idx]
            outs.append(
                resize(
                    input=conv(x),
                    size=inputs[0].shape[2:],
                    mode=self.interpolate_mode,
                    align_corners=self.align_corners))
        out = self.fusion_conv(torch.cat(outs, dim=1))
        out = self.cls_seg(out)
        return out

上图为Segformer的分割头部分,可以看出整个结构非常简洁

损失函数层面没有改进,依然还是最经典的交叉熵损失函数。

语义分割之SegFormer分享_第9张图片
实验结果:

语义分割之SegFormer分享_第10张图片
语义分割之SegFormer分享_第11张图片
鲁棒性测试:

语义分割之SegFormer分享_第12张图片

SegFormer一作的思考:

对于语义分割,特征提取非常重要,Transformer已经在分类上证明了比CNN更强大的特征提取能力。但是分类和分割还是有一定的GAP, 因此如何设计对分割友好的更好的Transformer结构,还可以继续研究。

有了很好的特征,decoder该如何设计才能进一步提高性能。这里用了一个很简单的MLP decoder取得了不错的效果,而传统的ASPP之类的decoder 在Transformer的基础上基本上没有帮助,未来如何针对性的设计更好的decoder也比较值得探索。

关于tf的思考:downsample、position encoding也都开始倾向于用conv了,退化回CNN架构的设计方式(SwinT使用了CNN的local和hierarchical思想)。到最后,Vision Transformer相比于CNN,可能只有local self-attention是有进步意义的。

你可能感兴趣的:(深度学习,人工智能,机器学习)