STDC 实时语义分割网络

1 简介

1.1 背景

为了满足速度更快的需求，很多研究者提出设计低延迟、高效并且满足分割准确率的 CNN 模型。这些实时语义分割方法在各种各样的基准中实现了有希望的性能。对于实时推理，一些工作，比如 DFANet 和 BiSeNetV1 选择轻量级主干和研究特征融合或聚合模块的方法来补偿准确率的丢失。然而，由于特定任务设计的不足，这些从图像分类任务中借用的轻量级主干可能对图像分割问题不是完美的。除了选择轻量级主干，限制输入图像的尺寸是提高推理速度的另一种常用方法。更小的输入分辨率看似高效，但是它容易忽略围绕边界的详细外观和小目标。为了解决这个问题，如图 2(a)，BiSeNet 采用多路径框架来组合低级细节和高级语义。然而，添加额外的路径来获取低级特征是耗时的，并且辅助路径通常缺少低级信息指导。

BiSeNet 已经被证明是流行的实时分割双流网络。然而，其增加额外路径对空间信息进行编码的原理非常耗时，并且由于特定任务设计的不足，从预训练任务（例如图像分类）中借用的主干可能对图像分割效率低下。

1.2 提出的方法

因为上述原因，作者提出了一个新颖的手工网络，以实现更快的推理速度、可解释的结构、有竞争力的性能等目的。首先，作者设计了一个新颖的结构，称为短期密集级联模块（Short-Term Dense Concatenate module, STDC module），用几个参数获得可变的感受野。然后，STDC 模块无缝地集成到 U-net 架构中，以形成 STDC 网络，这极大地提高了语义分割任务中网络性能。

细节上，如图 3 所示，作者从多个连续的层中连接相应的特征图，这些特征图在不同的尺度和感受野上编码了输入图像（或特征），实现了多尺度特征表示。为了加速，在可以忽略分割性能损失的情况下，层中的滤波器尺寸逐渐减少。STDC 网络的详细结构可见表 2。

在解码阶段，如图 2(b) 所示，与其使用额外的、耗时的路径，作者采用了细节引导（Detail Guidance）来引导低级层学习空间细节。作者首先使用细节引导模块来生成详细的 ground-truth。然后，二值交叉熵损失和 dice 损失被用于优化细节信息的学习任务，这被当成是一种辅助信息学习。应该注意，这种辅助信息在推理时是不需要的。最后，来自低级层的空间细节和来自深层的语义信息被融合，以预测语义分割结果。整体架构如图 4 所示。

在 Cityscapes 和 CamVid 数据集上的大量实验证明了作者的方法的有效性，在分割精度和推理速度之间实现了有希望的权衡。在 Cityscapes 上，作者在 NVIDIA GTX 1080Ti 上以 250.4 FPS 的速度在测试装置上实现了 71.9% 的 mIoU，比最新方法快了 45.2%，并在推理更高分辨率的图像时，以 97.0 FPS 实现了 76.8% 的 mIoU。

2 方法

2.1 编码网络的设计

2.1.1 短期密集级联模块（Short-Term Dense Concatenate module, STDC module）

作者提出的网络的关键组件是短期密集级联模块（STDC module）。图 3(b) 和（c）说明了 STDC 模块的布局。具体来讲，每个模块是分散到若干块的，作者使用 $Conv{X}_i$ 来表示第 $i$ 个块的操作。因此，第 $i$ 个块的输出计算如下：
$$x_i=Conv{X}_i(x_{i-1},k_i)\text{\hspace{1em}(1)}$$
其中，$x_{i-1}$ 和 $x_i$ 是第 $i$ 个块的输入和输出。ConvX 包括一个卷积层、一个批量归一化层和 ReLU 激活层，$k_i$ 是卷积层的核大小。

在 STDC 模块，第一个块的核大小为 1，其余的简单设置为 3。给定 STDC 模块的输出通道数 N，第 $i$ 个块的卷积层滤波器数量是 $N/2^i$，除了最后一个卷积层的滤波器，它的数量和先前的卷积层一样。在图像分类任务中，通常在更高层中使用更多的通道。但在语义分割任务中，作者关注于可变的感受野和多尺度信息。低级层需要足够的通道在小的感受野下编码更多的细粒度信息，而高级层有大的感受野，关注于更多高级信息归纳，设置和低级层相同的通道可能导致信息冗余。下采样仅发生在 Block2。为了丰富特征信息，作者通过跳跃路径连接 $x_1$ 到 $x_n$ 的特征图作为 STDC 模块的输出。在连接之前，STDC 模块不同块的特征图通过池化尺寸为 $3\times 3$ 的平均池化操作下采样到相同的空间尺寸，如图 3（c）所示。在作者的设置中，STDC 模块的最后输出是：
$$x_{output}=F(x_1,x_2,\cdots ,x_n)\text{\hspace{1em}(2)}$$
其中 $x_{output}$ 表示 STDC 模块的输出，$F$ 是融合操作，其中 $x_1,x_2,\cdots ,x_n$ 是来自所有 $n$ 个块的特征图。在效率的考虑下，作者采用连接作为融合操作。在作者的方法中，作者使用 STDC 模块在 4 个块中。

表 1 表示了 STDC 模块中的感受野，$x_{output}$ 聚合了所有块的多尺度信息。作者声称 STDC 模块有两个优点：（1）作者通过以几何级数的方式逐渐减少，精心调整了块的滤波器尺寸，导致了极大地减少计算复杂性。（2）STDC 模块的最后输出是连接了所有块，保留了可变感受野和多尺度信息。

给定输入通道维数 M 和输出通道维数 N，STDC 模块的参数量为：
$$\begin{gathered} S_{param}=M\times 1\times 1\times \frac{N}{2^1}+\sum _{i=2}^{n-1}\frac{N}{2^{i-1}}\times 3\times 3\times \frac{N}{2^i}+\frac{N}{2^{n-1}}\times 3\times 3\times \frac{N}{2^{n-1}} \\ =\frac{NM}{2}+\frac{9N^2}{2^3}\times \sum _{i=0}^{n-3}\frac{1}{2^{2i}}+\frac{9N^2}{2^{2n-2}}=\frac{NM}{2}+\frac{3N^2}{2}\times (1+\frac{1}{2^{2n-3}})\text{\hspace{1em}(3)} \end{gathered}$$
如公式 (3) 所示，STDC 模块的参数量由预先定义的输入和输出通道维度决定，因此块的数量对参数大小影响较小。尤其是，如果 n 达到最大极限，STDC 模块的参数量通常是常数，仅由 M 和 N 定义。

2.1.2 网络架构

图 3(a) 是作者的网络架构。它由除了输入层和预测层外的 6 个阶段组成。通常，阶段 1~5 以步长 2 下采样输入的空间分辨率，并且阶段 6 通过一个 ConvX、一个全卷平均池化层和两个全连接层输出预测逻辑。

阶段 1 和 2 通常被视为低级层，提取外观特征。在效率的追求上，作者只在阶段 1 和 2 上使用一个卷积块，已通过实验证明这是足够的。STDC 模块在阶段 3， 4， 5 的数量是仔细调整过的。在这些阶段中，每个阶段的第一个 STDC 模块以步长 2 下采样空间分辨率。其余的 STDC 模块保持空间分辨率不变。

作者表示每个阶段的输出通道数为 $N_l$，其中 $l$ 是阶段的索引。实际上，作者经验性地设置 $N_6$ 为 1024，并小心地调整余下地阶段通道数，直到在准确率和效率之间达到好的平衡。表 2 显示了 STDC 网络的详细结构。

2.2 解码器的设计

2.2.1 分割架构

作者使用预训练的 STDC 网络作为编码器的主干，并采用 BiSeNet 的上下文路径来编码上下文信息。如图 4(a)，作者使用阶段 3， 4， 5 来生成特征图，它们的下采样率分别是 1/8, 1/16, 1/32。然后，作者使用全局平均池化来提供具有大感受野的全局上下文信息。U 形结构被用于从全局特征上采样特征，并将它们中的每一个与编码阶段最后两个阶段的对应项（阶段 4 和 5）组合。根据 BiSeNet，作者使用注意力细化模块来细化每两个阶段的组合特征。对于最后语义分割预测，作者采用 BiSeNet 的特征融合模块来融合编码器中的阶段 3 的 1/8 下采样特征和解码器中的对应项。作者声称，这两个阶段的特征是不同级别的特征表示。编码主干的特征保留了丰富细节信息，解码器的特征包含上下文信息，由于输入来自全局池化层。具体来讲，分割头包含一个 $3\times 3$ Conv-BN-ReLU 操作，后跟一个 $1\times 1$ 卷积来获得输出维度 N，设置为类别的数量。作者采用具有 Online Hard Example Mining 的交叉熵损失来优化语义分割学习任务。

2.2.2 低级特征的细节引导

作者在图 5(b) 可视化 BiSeNet 的空间路径特征。比较主干相同下采样率的低级层（阶段 3），空间路径可以编码更多空间细节，例如，边界、角。基于这个观察，作者提出了一个细节引导模块来引导低级层以单流方式学习空间信息。作者建模细节预测为一个二值分割任务。作者首先通过 Laplacian 操作从分割 ground-truth 生成细节图 ground-truth，如图 4（c）。如图 4(a) 的说明，作者在阶段 3 插入 Detail Head 来生成细节特征图。然后作者使用细节 ground-truth 作为细节特征图的引导，来引导低级层学习空间细节信息。如图 5(d) 所示，具有细节引导的特征图可以比之前提到的图 5（c）的结果编码更多空间细节。最后，学习到的细节特征和解码器的深层块的上下文特征融合，用于分割预测。

细节 Ground-truth 生成：作者通过 Detail Aggregation 模块从语义分割 ground-truth 中生成二值细节 ground-truth，如图 4（c）中的虚线蓝色框。这个操作可以通过称为 Laplacian 核的 2-D 卷积核和一个可训练的 $1\times 1$ 卷积执行。如图 4(e)，作者使用 Laplcian 操作来生成软细细节特征图，并使用不同步长来获得多尺度细节信息。然后，作者上采样细节特征图到原始尺寸，并使用可训练的 $1\times 1$ 卷积来融合它们，以动态重加权。最后，作者采用一个阈值 0.1 来转换预测的细节为最后的二值细节 ground-truth，使其具有边界和角信息。

细节损失：由于细节像素的数量远小于非细节像素的数量，细节预测是一个类不平衡问题。由于加权的交叉熵损失通常导致粗糙的结果，根据 [7]，作者采用二值交叉熵和 dice 损失来联合优化细节学习。Dice 损失度量预测图和 ground-truth 之间的重叠。同时，它也对前景和背景像素的数量不那么敏感，意味着它能缓解类不平衡问题。因此，对预测的高为 H 宽为 W 的细节图，细节损失 $L_{detail}$ 表示如下：
$$L_{detail}(p_d,g_d)=L_{dice}(p_d,g_d)+L_{bce}(p_d,g_d)\text{\hspace{1em}(4)}$$
其中 $p_d\in {\mathbb{R}}^{H\times W}$ 表示预测的细节，$g_d\in {\mathbb{R}}^{H\times W}$ 表示对应的细节 ground-truth。$L_{dice}$ 表示为：
$$L_{dice}(p_d,g_d)=1-\frac{2{\sum }_i^{H\times W}{p}_d^i{g}_d^i+ϵ}{{\sum }_i^{H\times W}(p_d^i{)}^2+{\sum }_i^{H\times W}(g_d^i{)}^2+ϵ}\text{\hspace{1em}(5)}$$

如图 4(b) 所示，作者使用一个 Detail Head 来生成细节图，引导浅层编码空间信息。Detail Head 包括一个 $3\times 3$ Conv-BN-ReLU 操作，后跟一个 $1\times 1$ 卷积来获得输出的细节图。注意，这个分支在推理阶段不使用。