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架构进行了探索:
- a:源自SegFormer的解码器,是一种纯MLP架构;
- b:常被CNN方案使用,如ASPP、PSP、DANet等;
- c:本文采用的解码器,它采用轻量型Hamberger模块对后三个阶段的特性进行聚合以进行全局上下文建模。
需要注意的是,SegFormer的解码器对Stage1到Stage4的特征进行聚合,而本文方案则仅对Stage2-Stage4的特征进行聚合。这是因为:
- SegNeXt的Encoder采用了卷积架构,使得Stage1部分特征包含过多底层信息,进而导致其会影响语义分割性能。
- 对Stage1部分特征进行处理会带来过多的计算负载。