STDC语义分割网络:BiSeNet的轻量化加强版（CVPR2021语义分割）

今天和大家分享CVPR2021中语义分割领域的一篇论文，论文标题及作者如下图所示：

论文下载地址：https://arxiv.org/abs/2104.13188

代码即将开源：https://github.com/MichaelFan01/STDC-Seg

0.动机

深度学习促进了计算机视觉领域中语义分割算法的发展；随着自动驾驶、视频监控等应用场景的迫切需求，面向终端的轻量化语义分割网络被广泛研究。

为了减少算法的耗时，一些诸如DFANet、BiSeNetV1的语义分割算法使用了轻量化backbone，但是它们都是直接使用为分类任务设计的backbone，作者认为，专为分类任务设计的backbone不能在语义分割领域中充分发挥它们的性能。

BiSeNet使用了multi-path结构的网络融合低层次特征和高层次特征，然而额外增加分支会增加网络的运行时间。

作者提出了STDC模块，能够使用较少的参数量提取多尺度特征，且能够很方便地集成到U-Net类型的语义分割网络中；对BiSeNet中的multi-path结构做出改进，在提取底层细节特征的同时减少网络计算量。

1.STDC（Short-Term Dense Concatenate）模块

STDC模块的结构如下图所示：

图中的$ConvX$表示“卷积+BN+ReLU”操作，$M$表示输入特征通道数，$N$表示输出特征通道数。

在STDC模块中，第1个block的卷积核尺寸为$1\times 1$，其余block的卷积核尺寸为$3\times 3$。

若STDC模块的最终输出通道数为$N$，除最后一个block外，该模块内第$i$个block的输出通道数为$N/2^i$；最后一个block的输出特征通道数与倒数第二个block保持一致。

与传统的backbone不同的是，STDC模块中深层的特征通道数少，浅层的特征通道数多。作者认为，浅层需要更多通道的特征编码细节信息；深层更关注高层次语义信息，过多的特征通道数量会导致信息冗余。

STDC模块最终的输出为各block输出特征的融合，即：

$$x_{\text{output }}=F\left(x_1,x_2,\dots ,x_n\right)$$

上式中的$F$表示融合函数，$x_1,x_2,\dots ,x_n$表示$n$个block的输出，$x_{\text{output }}$表示STDC模块的输出。使用concatenation操作融合$n$个block的特征。

下图表示stride=2版本的STDC模块：

对于stride=2版本的STDC模块，在$Block2$中进行下采样操作；为了在融合时保证feature map尺寸一致，对大尺寸的feature map使用stride=2、$3\times 3$的average pooling操作进行下采样。

STDC模块有2个特点：（1）随着网络加深，逐渐减少特征通道数，以减少计算量；（2）STDC的输出融合了多个block的输出feature map，包含多尺度信息。

下图表示由STDC模块组成的STDC网络：

该网络包含6个Stage，Stage1~Stage5中都对feature map进行了步长为2的下采样，Stage6输出预测结果。

为了减少计算量，Stage1和Stage2中只使用1个卷积层。Stage3~Stage5中每个Stage包含若干个STDC模块，其中第1个STDC模块包含下采样操作，其余STDC模块保持feature map尺寸不变。

以上图为框架，作者构建了2个STDC网络，分别命名为STDC1和STDC2，它们的结构如下表所示：

表中的$ConvX$表示“卷积+BN+ReLU”操作，Stage3~Stage5均由若干个STDC模块组成。上表中的$KSize$表示kernel尺寸，$S$表示步长，$R$表示重复次数，$C$表示输出通道数。

2.STDC分割网络架构

使用BiSeNet中的context path提取上下文信息，使用STDC网络作为backbone，构建STDC分割网络，如下图中图（a）所示：

网络中的Stage3、Stage4和Stage5均对feature map进行下采样操作，然后使用global average pooling提取全局上下文信息。使用U型结构上采样特征，并与Stage4、Stage5后面的ARM（Attention Refine Module，借鉴自BiSeNet）输出进行融合，融合结果作为FFM（Feature Fusion Module，借鉴自BiSeNet）的输入，为FFM提供高层次语义信息，FFM的另外一个输入为Stage3输出的特征，该特征为FFM提供低层次细节信息。FFM的输出特征通过$8\times$上采样进入Seg Head，Seg Head包括1个$3\times 3$的卷积、BN和ReLU操作，以及1个$1\times 1$卷积，最终输出$N$维度特征，$N$为分割类别数。

为了弥补去除BiSeNet中的Spatial分支造成的细节损失，作者在Stage3后面插入了Detail Head，使得Stage3能够学到细节信息。需要注意的是，Detail Head只在训练时使用，目的是让Stage3输出的feature map包含更多的细节特征，用于与Context分支的高层次语义特征融合。

作者使用上图中图（c）的结构得到Detail Head的ground truth，即对标签分别做步长为1、2、4的Laplacian Conv，分别对步长为2和步长为4的输出结果进行$2\times$和$4\times$的上采样，使用带有可学习参数的$1\times 1$卷积融合，通过0.1的阈值输出最终的2值ground truth。

上述操作生成的ground truth中，细节特征所占的像素远少于非细节所占的像素，因此Detail Head部分面临着类别不均衡问题。因此联合使用binary cross-entropy loss和dice loss，dice loss可以缓解类别不均衡问题，detail loss为：

$$L_{\text{detail }}\left(p_d,g_d\right)=L_{\text{dice }}\left(p_d,g_d\right)+L_{bce}\left(p_d,g_d\right)$$

上式中的$p_d\in {\mathbb{R}}^{H\times W}$表示预测结果，$g_d\in {\mathbb{R}}^{H\times W}$表示ground truth，$H$和$W$分别表示detail map的高度和宽度；$L_{\text{bce}}$表示binary cross-entropy loss，$L_{dice}$表示dice loss，$L_{dice}$表达式为：

$$L_{dice}\left(p_d,g_d\right)=1-\frac{2{\sum }_i^{H\times W}{p}_d^i{g}_d^i+ϵ}{{\sum }_i^{H\times W}{\left(p_d^i\right)}^2+{\sum }_i^{H\times W}{\left(g_d^i\right)}^2+ϵ}$$

上式中$ϵ$的存在是为了防止分母为0，$ϵ=1$。

上图中图（b）的Detail Head包括1个$3\times 3$的卷积、BN和ReLU操作，以及1个$1\times 1$卷积。

STDC分割网络和BiSeNet在总体结构上的差异如下图所示：

STDC分割网络使用Detail Guidance替代BiSeNet中的Spatial分支，在保留细节特征的同时减少了计算量。

3.实验结果

3.1 STDC分类网络性能

在ImageNet数据集上训练STDC分类网络，并在验证集上测试，结果如下表所示：

从上表中可以看出，STDC系列网络兼顾性能和计算速度。

3.2 STDC分割网络性能

在Cityscapes数据集上训练STDC分割网络，分别使用$512\times 1024$和$768\times 1536$分辨率的图片作为网络输入，将输出结果上采样至$1024\times 2048$。在测试集上测试的模型使用了训练集和验证集进行训练，性能和速度如下表所示：

上表中的Seg50表示输入图片分辨率为$512\times 1024$，Seg75表示输入图片分辨率为$768\times 1536$。

从上表中可以看出，STDC分割网络在保持高推理速度的同时有着更高的精度。

一个更直观的比较如下图所示：

上表中的性能是在Cityscapes测试集上的结果，FPS在1080Ti GPU上使用batch为1时测得。

3.3 Detail Guidance对细节特征的影响

下图直观地表示了使用Detail Guidance与否对Stage3特征和分割结果的影响：

图（a）表示输入图像，图（b）和图（d）分别表示未使用Detail Guidance时Stage3和网络预测结果，图（c）和图（e）表示使用了Detail Guidance时Stage3和网络预测结果，图（f）表示ground truth。

从图中可以看出，使用了Detail Guidance后，Stage3的feature map细节更丰富，分割结果中的细节也更好。

4.总结

作者提出了STDC轻量化backbone用于分割任务；使用Detail Guidance替代BiSeNet中的Spatial分支，保留低层次细节特征的同时有效地减少了网络计算量。

推荐阅读：

SAC:简单高效的域自适应语义分割自监督方法（CVPR2021语义分割）

引用超7000次的SegNet:Encoder-Decoder架构的端到端语义分割网络

FCN：使用端到端CNN进行语义分割的开山之作

如果你对计算机视觉中的目标检测、跟踪、分割、轻量化神经网络感兴趣，欢迎关注公众号一起学习交流~