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模块的结构如下图所示：
图中的 C o n v X ConvX ConvX表示“卷积+BN+ReLU”操作， M M M表示输入特征通道数， N N N表示输出特征通道数。
在STDC模块中，第1个block的卷积核尺寸为 1 × 1 1 \times 1 1×1，其余block的卷积核尺寸为 3 × 3 3 \times 3 3×3。
若STDC模块的最终输出通道数为 N N N，除最后一个block外，该模块内第 i i i个block的输出通道数为 N / 2 i N / 2^{i} N/2i；最后一个block的输出特征通道数与倒数第二个block保持一致。
与传统的backbone不同的是，STDC模块中深层的特征通道数少，浅层的特征通道数多。作者认为，浅层需要更多通道的特征编码细节信息；深层更关注高层次语义信息，过多的特征通道数量会导致信息冗余。
STDC模块最终的输出为各block输出特征的融合，即：
x output = F ( x 1 , x 2 , … , x n ) x_{\text {output }}=F\left(x_{1}, x_{2}, \ldots, x_{n}\right) xoutput =F(x1,x2,…,xn)
上式中的 F F F表示融合函数， x 1 , x 2 , … , x n x_{1}, x_{2}, \ldots, x_{n} x1,x2,…,xn表示 n n n个block的输出， x output x_{\text {output }} xoutput 表示STDC模块的输出。使用concatenation操作融合 n n n个block的特征。
下图表示stride=2版本的STDC模块：
对于stride=2版本的STDC模块，在 B l o c k 2 Block2 Block2中进行下采样操作；为了在融合时保证feature map尺寸一致，对大尺寸的feature map使用stride=2、 3 × 3 3 \times 3 3×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，它们的结构如下表所示：
表中的 C o n v X ConvX ConvX表示“卷积+BN+ReLU”操作，Stage3~Stage5均由若干个STDC模块组成。上表中的 K S i z e KSize KSize表示kernel尺寸， S S S表示步长， R R R表示重复次数， C C 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 × 8 \times 8×上采样进入Seg Head，Seg Head包括1个 3 × 3 3 \times 3 3×3的卷积、BN和ReLU操作，以及1个 1 × 1 1 \times 1 1×1卷积，最终输出 N N N维度特征， N 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 × 2 \times 2×和 4 × 4 \times 4×的上采样，使用带有可学习参数的 1 × 1 1 \times 1 1×1卷积融合，通过0.1的阈值输出最终的2值ground truth。
上述操作生成的ground truth中，细节特征所占的像素远少于非细节所占的像素，因此Detail Head部分面临着类别不均衡问题。因此联合使用binary cross-entropy loss和dice loss，dice loss可以缓解类别不均衡问题，detail loss为：
L detail ( p d , g d ) = L dice ( p d , g d ) + L b c e ( p d , g d ) L_{\text {detail }}\left(p_{d}, g_{d}\right)=L_{\text {dice }}\left(p_{d}, g_{d}\right)+L_{b c e}\left(p_{d}, g_{d}\right) Ldetail (pd,gd)=Ldice (pd,gd)+Lbce(pd,gd)
上式中的 p d ∈ R H × W p_{d} \in \mathbb{R}^{H \times W} pd∈RH×W表示预测结果， g d ∈ R H × W g_{d} \in \mathbb{R}^{H \times W} gd∈RH×W表示ground truth， H H H和 W W W分别表示detail map的高度和宽度； L bce L_{\text {bce}} Lbce表示binary cross-entropy loss， L d i c e L_{dice} Ldice表示dice loss， L d i c e L_{dice} Ldice表达式为：
L d i c e ( p d , g d ) = 1 − 2 ∑ i H × W p d i g d i + ϵ ∑ i H × W ( p d i ) 2 + ∑ i H × W ( g d i ) 2 + ϵ L_{d i c e}\left(p_{d}, g_{d}\right)=1-\frac{2 \sum_{i}^{H \times W} p_{d}^{i} g_{d}^{i}+\epsilon}{\sum_{i}^{H \times W}\left(p_{d}^{i}\right)^{2}+\sum_{i}^{H \times W}\left(g_{d}^{i}\right)^{2}+\epsilon} Ldice(pd,gd)=1−∑iH×W(pdi)2+∑iH×W(gdi)2+ϵ2∑iH×Wpdigdi+ϵ
上式中 ϵ \epsilon ϵ的存在是为了防止分母为0， ϵ = 1 \epsilon=1 ϵ=1。
上图中图（b）的Detail Head包括1个 3 × 3 3 \times 3 3×3的卷积、BN和ReLU操作，以及1个 1 × 1 1 \times 1 1×1卷积。
STDC分割网络和BiSeNet在总体结构上的差异如下图所示：
STDC分割网络使用Detail Guidance替代BiSeNet中的Spatial分支，在保留细节特征的同时减少了计算量。
3.实验结果
3.1 STDC分类网络性能
在ImageNet数据集上训练STDC分类网络，并在验证集上测试，结果如下表所示：
从上表中可以看出，STDC系列网络兼顾性能和计算速度。
3.2 STDC分割网络性能
在Cityscapes数据集上训练STDC分割网络，分别使用 512 × 1024 512 \times 1024 512×1024和 768 × 1536 768 \times 1536 768×1536分辨率的图片作为网络输入，将输出结果上采样至 1024 × 2048 1024 \times 2048 1024×2048。在测试集上测试的模型使用了训练集和验证集进行训练，性能和速度如下表所示：
上表中的Seg50表示输入图片分辨率为 512 × 1024 512 \times 1024 512×1024，Seg75表示输入图片分辨率为 768 × 1536 768 \times 1536 768×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进行语义分割的开山之作
如果你对计算机视觉中的目标检测、跟踪、分割、轻量化神经网络感兴趣，欢迎关注公众号一起学习交流~
