SCTNet: Single-Branch CNN with Transformer Semantic Information for Real-Time Segmentation

论文地址: https://ojs.aaai.org/index.php/AAAI/article/view/28457
代码地址: https://github.com/xzz777/SCTNet

摘要

• 最近的实时语义分割方法通常采用一个额外的语义分支来追求丰富的远程上下文。
• 然而，额外的分支会带来额外的计算开销，并减慢推理速度
• 为了解决这种困境，我们提出了SCTNet——一个用于实时语义分割的具有Transformer语义信息的单分支CNN
• SCTNet具有推理解放语义分支的丰富语义表示，同时保持轻量级单分支CNN的高效率
• 考虑到它的强大的语义能力，SCTNet利用一个transformer作为只用来训练的语义分支
• 借助所提出的transformer类CNN模块CFBlock和语义信息对齐模块，SCTNet可以在训练中从transformer分支中捕获丰富的语义信息。
• 推理阶段，只有单个CNN分支需要部署。
• 对Cityscapes、ADE20K和COCO-Stuff-10K上进行实验，结果表示我们的方法达到了新的SOTA性能。

背景

• 最近的语义分割方法追求丰富的长程上下文信息。Transformer的自注意力机制被证明具有对长程上下文建模的显著能力。尽管这些工作得到了显著改进，但也会导致高昂的计算成本。基于自注意力的工作在图像分辨率方面甚至具有平方计算复杂度，这会显著增加处理高分辨率图像的延迟。
• 最近许多实时工作采用双边架构来快速提取高质量的语义信息。BiSeNet提出了一种双边网络，用于在早期阶段分离详细的空间特征和丰富的上下文信息，并并行处理它们，如下图a所示。BiSeNetV2和STDC做了进一步努力，以增强提取丰富的长程上下文的能力或降低空间分支的计算成本。为了平衡推理速度和精度，DDRNet、RTFormer和SeaFormer采用一种特征共享架构，该架构在深层阶段划分空间和上下文特征，如下图b所示。但是这些方法在两个分支之间引入了密集融合模块，以增强提取特征的语义信息。总之，所以的双边方法都由于额外的分支和多个融合模块而导致推理速度有限和计算成本高。
  

创新点

• 为了解决上述困境，提出了一个用于实时分割的单分支CNN，其中包含Transformer语义信息。它可以有效地提取语义信息，而不会因为双边网络而繁重的计算。
• SCTNet从仅用于训练的Transformer语义分支学习长程上下文信息到CNN分支。
• 为了减轻Transformer和CNN之间的语义差距，设计了一个类似Transformer的CNN模块，称为CFBklock。在对齐之前使用共享解码头。
• 通过在训练中对齐的语义信息，单分支CNN可以联合编码语义信息和空间细节。
• SCTNet可以保持轻量级单分支CNN架构在推理中的高效的同时，对齐来自Transformer架构的大感受野的语义表示。整体架构如上图c表示。主图如下图所示
  

方法

Conv-Former Block

CNN和Transformer提取的特征表示存在显著差异。直接对齐CNN和Transformer之间的特征会使学习过程变得困难，导致性能提升有限。为了让CNN分支能够轻松学习到从Transformer分支中提取到的 高质量的语义信息，设计了Conv-Former模块。

Cov-Former模块尽可能的模拟Transformer模块的结构，以便更好地学习Transformer分支的语义信息。同时，Conv-Former仅使用高效地卷积运算来实现注意力函数。

Conv-Former模块的结构类似于典型的Transformer编码器结构，如下图左侧所示。其过程为:
$$\begin{array}{rl}f& =Norm(x+ConvAttention(x))\\ y& =Norm(f+FFN(f))\end{array}$$
其中$Norm$表示batch Normalization。

卷积注意力机制

用于实时分割的注意力机制应具有低延迟和强大的语义提取能力。GFA是一个潜在的候选方案，提出的卷积注意力机制源于GFA。

GFA和提出的卷积注意力机制主要有两个区别:

1. 将GFA的矩阵乘法替换为像素级卷积运算。点卷积等于像素到像素的乘法，但不涉及特征扁平化和重塑操作。这些操作不利于保持固有的空间结构，并会增加额外的推理延迟。此外，卷积提高了一种更灵活的方式来扩展外部参数。
2. 由于Transformer和CNN之间的语义差距，仅仅计算几个可学习向量与每个像素之间的相似性，然后根据相似图和可学习向量增强像素，不足以捕捉丰富的上下文信息。为了更好地对齐Transformer的语义信息，将可学习向量扩展为可学习核。一方面，这将像素与可学习向量之间的相似度计算转换为像素块与可学习核之间的相似性计算；另一方面，使用可学习核的卷积操作在某种程度上保留了更多局部空间信息。

卷积注意力的操作可以总结如下:
$$X=\theta (X\otimes K)\otimes K^T,$$
其中$X\in {\mathbb{R}}^{C\times H\times W}$, $K\in {\mathbb{R}}^{C\times N\times k\times k}$, $K^T\in {\mathbb{R}}^{N\times C\times k\times k}$分别表示输入图像和可学习的query和key。$N$表示可学习参数的数量, $k$表示可学习参数的核大小。$\theta$表示分组双重归一化，它在$H\times W$维度上应用softmax, 并在$N$维度上应用分组$L_2$归一化，$\otimes$表示卷积操作。

考虑到效率, 使用条带卷积而不是标准卷积来实现卷积注意力。具体来说，利用一个$1\times k$卷积核一个$k\times 1$卷积来近似一个$k\times k$卷积层。上图右侧为卷积注意力的实现细节。

前馈网络FFN

经典的FFN在提供位置编码和嵌入通道方面起着至关重要的作用。最近的Transformer模型中的FFN由扩展点卷积、深度方向$3\times 3$卷积核压缩点卷积组成。与经典的FFN不同，本文的FFN由两个标准$3\times 3$卷积层组成。这样可以更高效，并提供更大的感受野。

语义信息对齐模块

提出了一个简单而有效地对齐模块来进行训练过程中的特征学习。它可以分为主干特征对齐和共享解码器头对齐。

主干特征对齐

由于Conv-Former模块类似Transformer的架构，对齐损失可以轻松地将Conv-Former模块的特征与Transformer的特征对齐。简而言之，主干特征对齐首先对来自Transformer和CNN分支的特征进行下采样或上采样以进行对齐。然后将CNN的特征投影到Transformer的维度。需要统一通道数量以及避免直接对齐。直接对齐会在训练过程中破坏CNN的真实标签监督。最后，将语义对齐损失应用于投影后的特征以对齐语义表示。

共享解码头对齐

Transformer解码器通常使用多个阶段的特征进行复杂解码，而SCTNet解码器仅选择stage2和stage4的特征。考虑到两者之间解码空间的显著差异，直接对齐解码头特征和输出logits只能获得有限的改进。因此，提出共享解码头对齐。具体来说，将单分支CNN的concatenation stage2和stage4特征输入到点卷积中以扩展维度。然后将高维特征通过Transformer解码器。Transformer解码器的新输出特征和logits用于计算与Transformer解码器原始输出的对齐损失。

总体架构

为了获得丰富的语义信息的同时降低计算成本，将流行的两分支架构简化为一个用于推理的快速CNN分支和一个只用于训练的语义对齐的Transformer分支。

backbone

为了提高推理速度，SCTNet采用典型的CNN主干网络。SCTNet起始于一个由两个连续的$3\times 3$卷积层组成的stem块。前两个阶段由堆叠的残差块组成，后两个阶段包含提出的类似Transformer的块，称为Conv-Former块(CFBlock)。CFBlock采用多个精心设计的卷积操作来执行与Transformer块类似的远程上下文捕获功能。在2-4阶段的开始应用一个由bacth normal和ReLU激活的步幅卷积组成的convdown层进行下采样。

解码器头

解码头由一个DAPPM和一个分割头组成。为了进一步丰富上下文信息，在第4阶段的输出之后添加了一个DAPPM。然后将输出与阶段2的特征图连接起来。最后，这个输出特征被传递到一个分割头。精确地说，分割头由一个$3\times 3$的Conv-BN-ReLU算子后跟一个$1\times 1$卷积分类器组成。

对齐损失

为了更好地对齐语义信息， 需要一种关注语义而不是空间信息的对齐损失。在实现中，使用CWD损失作为对齐损失, 它显示出比其他损失函数更好地结果。CWD损失如下:
$$\begin{array}{rl}\varphi (x_c)& =\frac{exp(\frac{x_{c,i}}{\tau })}{{\sum }_{i=1}^{W\cdot H}exp(\frac{x_{c,i}}{\tau })},\\ L_{cwd}& =\frac{{\tau }^2}{C}\sum _{c=1}^C\sum _{i=1}^{H\cdot W}\varphi (x_T^{c,i})\cdot \log \left[\frac{\varphi (x_T^{c,i})}{\varphi (x_T^{c,i})}\right],\end{array}$$

其中$c=1,2,\dots ,C$表示通道索引，$i=1,2,\dots ,H\cdot W$表示空间位置，$x^T$和$x^S$分别表示Transformer分支和CNN分支的特征图。$\varphi$将特征激活转换为逐通道概率分布，消除了Transformer和CNN紧凑CNN之间的尺度影响。为了最小化$L_{cwd}$, 当$\varphi (x_T^{c,i})$很大时，$\varphi (x_S^{c,i})$应该很大。但当$\varphi (x_T^{c,i})$很小时，$\varphi (x_S^{c,i})$的值无关紧要。这迫使CNN学习前景显著性的分布，其中包含语义信息。$\tau$是一个超参数，值越大，概率分布就越平滑。

实验

SOTA效果对比

Cityscapes:

ADE20K:

COCO-Stuff-10k:

Cityscapes可视化对比

结论

在本文中，我们提出了SCTNet，这是一种新颖的单分支架构，它可以在不增加额外推理计算成本的情况下提取高质量的远程上下文。 广泛的实验表明，SCTNet取得了新的最先进的结果。 此外，通过证明SCTNet的效率，我们为双边分支网络中的语义分支提供了一种新的见解，并提供了一种新的方法来提升实时分割社区，这不仅采用了Transformer的结构，而且还利用了它的知识。

补充

• 在ImageNet-1K上进行预训练。遵循swin-transformer在ImageNet-1K上的预训练设置。