SegNeXt:重新思考语义分割中的卷积注意力设计.

语义分割任务的主流模型逐渐采用Vision Transformer结构。本文对已有成功分割模型进行了重审视并发现了几个有助于性能提升的关键成分,进而设计了一种新型的卷积注意力架构方案SegNeXt

本文对语义分割领域代表性方案(DeepLabV3+, HRNet, SETR, SegFormer)进行重审视,总结出成功的语义分割方案应具有的几点关键属性:

基于上述考量,本文对卷积注意力设计进行了重思考并提出了一种简单而有效的编码器-解码器架构SegNeXt。不同于已有Transformer方案,SegNeXt对编码器模块采用传统卷积模块设计但引入了多尺度卷积注意力,对解码器模块采用了Hamberger(自注意力的一种替代方案)进一步提取全局上下文信息。因此SegNeXt能够从局部到全局提取多尺度上下文信息,能在空域与通道维度达成自适应性,能从底层到高层进行信息聚合。

1. 编码器

Encoder部分采用了金字塔架构,每个构成模块采用了类似ViT的结构,但不同之处在于:自注意力模块通过一种多尺度卷积注意力模块MSCA实现。

MSCA由三部分构成:深度卷积用于聚合局部信息、多分支深度卷积用于捕获多尺度上下文信息、$1\times 1$卷积用于在通道维度进行相关性建模。卷积的输出将作为注意力权值对MSCA的输入进行重加权。此外主要注意的是:MSCAN的每个模块采用的是BN,而非LN

上表给出了通过堆叠MSCA而得到的不同MSCAN骨干信息以及SegNeXt架构信息。MSCA的实现参考如下。


class AttentionModule(BaseModule):
    def __init__(self, dim):
        super().__init__()
        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
        self.conv0_1 = nn.Conv2d(dim, dim, (1, 7), padding=(0, 3), groups=dim)
        self.conv0_2 = nn.Conv2d(dim, dim, (7, 1), padding=(3, 0), groups=dim)
 
        self.conv1_1 = nn.Conv2d(dim, dim, (1, 11), padding=(0, 5), groups=dim)
        self.conv1_2 = nn.Conv2d(dim, dim, (11, 1), padding=(5, 0), groups=dim)
 
        self.conv2_1 = nn.Conv2d(dim, dim, (1, 21), padding=(0, 10), groups=dim)
        self.conv2_2 = nn.Conv2d(dim, dim, (21, 1), padding=(10, 0), groups=dim)
        self.conv3 = nn.Conv2d(dim, dim, 1)
 
    def forward(self, x):
        u = x.clone()
        attn = self.conv0(x)
 
        attn_0 = self.conv0_1(attn)
        attn_0 = self.conv0_2(attn_0)
 
        attn_1 = self.conv1_1(attn)
        attn_1 = self.conv1_2(attn_1)
 
        attn_2 = self.conv2_1(attn)
        attn_2 = self.conv2_2(attn_2)
        attn = attn + attn_0 + attn_1 + attn_2
 
        attn = self.conv3(attn)
        return attn * u

2. 解码器

常规语义分割模型的骨干往往在ImageNet上预训练得到,为捕获高级语义信息,通常需要一个Decoder模块。本文则对以下三种简单Decoder架构进行了探索:

需要注意的是,SegFormer的解码器对Stage1Stage4的特征进行聚合,而本文方案则仅对Stage2-Stage4的特征进行聚合。这是因为: