FastFCN: 对语义分割中扩张卷积的重新思考
Deepwise AI Lab
2019.03.28
翻译时间:2019.04.09
摘要
现代语义分割的方法通常会在主干网络中使用扩张卷积来抽取高分辨率的特征图,使模型的计算复杂度和存储占用增大。为了替代耗时耗空间的扩张卷积,我们将抽取高分辨率特征图的任务构建为一个联合上采样问题,提出了一种新的联合上采样模块,称为***Joint Pyramid Upsampling (JPU)***。我们使用JPU的方法在没有性能损失的前提下降低了三倍的计算复杂度。实验显示了JPU优于其他上采样模块,且可以插入进现有的模型中降低计算复杂度并提升性能。通过将扩张卷积替换为JPU,我们的方法在Pascal Context数据集上达到sota水平,mIoU为51.3%,在ADE20K数据集上获得了0.5584的最终得分,而且速度快了三倍。代码已经开源 https://github.com/wuhuikai/FastFCN
引言
语义分割是一种计算机视觉中的基本任务之一,目标是将图像中的每个像素都分配一个语义标签。现代的方法通常使用全卷积网络(FCN)来解决这个问题,并取得了极大的成功。
传统的FCN由Long等人提出,从卷积神经网络中转变过来,用来进行图像分类。这种设计根据步长逐次对输入图像进行卷积或池化的的下采样,最后得到了一个低分辨率的特征图。尽管最后的特征图中编码了丰富的语义信息,但是精细化的图像结构信息被丢失了,这会使图像边界的预测不准。如图1a所示,传统的FCN通常会下采样5倍到最后的特征图,也就是减少了32倍的空间分辨率。
为了获得高分辨率的最终特征图,[3,28,18,30,27]采用了传统的FCN作为编码器来捕获高层的语义信息,接一个解码器结合编码器生成的多层特征图来逐步恢复空间信息。如图1b所示,我们把这样的方法叫做EncoderDecoder,其中解码器输出的预测图是高分辨率的。
另外,DeepLab移除了FCN的最后两层下采样操作,并且用扩张卷积来保证感受野和FCN相同。根据DeepLab,[38,6,36]采用了一种多尺度的语义模块在最终特征图的位置上,比EncoderDecoder方法更加出色。如图1c所示,在DilatedFCN中,最终特征图的空间分辨率比传统的FCN大了4倍,进而保留了更多的结构和位置信息。
扩张卷积在保留最终特征图的空间分辨率上非常重要,使其性能优于EncoderDecoder的大多数方法。然而,扩张卷积会带来额外的计算复杂度和内存占用,限制了其在许多实时场合上的应用。以ResNet-101为例,相比于传统FCN,DilatedFCN有23个残差模块(69层卷积),需要4倍的额外计算开销和内存使用,3个残差模块(9层卷积)需要16倍额外的资源。
我们的目的是要解决上述的这些因扩张卷积而引起的问题,提出了新的联合上采样模块来代替扩张卷积,称为JPU。我们的方法以传统的FCN作为骨架,使用JPU在上采样低分辨率的最终特征图(output stride=32,OS=32),获得了一个高分辨率的特征图(OS=8),这样,计算时间和内存消耗都大大减少,并且没有性能上的损失。我们将这个优势归功于JPU实现的多尺度特征图的融合。
为了验证方法的有效性,我们首先进行了系统性的实验,展示了JPU能够在几种流型的方法中替代扩张卷积,并且不降低性能。接着,我们在各个语义分割基准上测试了所提方法,结果显示我们的方法达到网络sota性能,运行时间缩短了3倍。具体来说,我们优于Pascal Context上所有的baseline。【具体指标】
总结如下,我们的贡献有三点:
- 提出了一种减小计算量和存储空间的联合上采样模块,JPU;
- 基于JPU,在提升3倍速度的基础上,获得了更好的性能;
- 在Pascal Context和ADE20K上获得了sota效果。
相关工作
在这部分,我们首先给出了语义分割各种方法的总览,可以分为两个大方向,接下来介绍了上采样的相关工作。
语义分割
FCN在语义分割上获得了巨大的成功,接着FCN的工作,有两大方向,一是DilatedFCN,一是EncoderDecoder。DilatedFCN使用扩张卷积在保证感受野,并且获得多尺度的高层特征图。EncoderDecoder采用编码器抽取多层特征图,送入解码器得到最后的预测。
DilatedFCN
为了在较高分辨率的最终特征图上获得多尺度的语义信息,PSPNet在多网络尺度上进行了池化操作,DeepLabV3采用了不同扩张率的ASPP,EncNet采用了语义编码模块(Context Encoding Module)来获得全局语义信息。我们的方法提出了JPU来替代DilatedFCN中的扩张卷积,在没有性能损失的条件下极大降低了计算量。
EncoderDecoder
为了逐步恢复空间信息,[28]提出了skip连接来构建UNet,结合了编码器的特征和对应解码器的激活。[18]提出了一种多路径的精调网络,显式地采用了下采样阶段的全部信息。DeepLabV3+结合了DilatedFCN和EncoderDecoder的优点,把DeepLabV3作为编码器。我们的方法是一种DeepLabV3+的补充,可以降低DeepLabV3的计算量,并且没有性能损失。
上采样
在我们的方法中,我们提出了一种上采样的模块将高分辨率的特征图作为指导,这和联合上采样(joint upsampling)和依赖数据的上采样(data-dependent upsampling)非常相关。
Joint Upsampling
在图像处理的过程中,联合上采样将指导图像作为先验,将指导图像中的信息转移到目标图像上。[17]构建了一种基于CNN的联合滤波器,可以学着恢复指导图像中的结构细节。[31]提出了一种端到端的可训练引导滤波模块,有条件地上采样低分辨率图像。我们的方法和上述方法相关,但是JPU可以处理含很多通道数的图像,而[17,31]只能处理3通道的图像,这样就不能捕获到高维特征图的复杂关联。而且我们的动机和目的是完全不同的。
Data-Dependent Upsampling
DUpsampling也和我们的方法有关,其利用了语义标签空间的冗余,能够从低分辨率的CNN输出中恢复像素级的预测。相比我们的方法DU对标签空间有很强的依赖,对于更大更复杂的标签空间的泛化能力弱。
方法
在这一部分,我们首先介绍最流行的语义分割模型,DilatedFCN。接着我们用JPU改造DilatedFCN。最后我们具体讨论JPU的细节,在此之前会简要介绍扩张卷积和步幅卷积。
DilatedFCN
为了在语义分割中用到深度神经网络,Long等人将CNN转换成全卷积的图像分类器。以ResNet-101为例,传统的CNN包含5个卷积阶段、1个全局平均池化、一个全连接层。为了构建FCN,全局平均池化和全连接层被替换为了卷积层,用来生成最后的标签图,如图1a所示。在两个连续的卷积模块之间会用到步幅卷积和/或空间池化层会被用到,使特征图在空间分辨率上减小。
FCN的空间分辨率减小了32倍,使预测位置和细节不准。为了获得高分辨率的最终特征图,DeepLab移除了两个特征图之间的下采样操作,如图1c,并用扩张卷积替代,使得感受野不变,因为叫做DilatedFCN。最后一层特征图下采样了8倍,保留了更多的位置和细节信息。接着DeepLab的工作,[38,6]提出了一种多尺度的语义模块来从最终特征图中获得语义信息,在语义分割上效果出色。
我们方法的框架
为了获得高分辨率的最终特征图,DilatedFCN中的方法移除了最后两个下采样操作,这样带来了计算负担和存储负担。在本文中,我们寻找一种DilatedFCN的替代品,来减小计算量和空间要求。同时,我们也希望没有性能损失。
为了实现这个目的,我们首先用传统的卷积层题海所有的扩张卷积,如图2,这样我们的主干和FCN是一致的,这样从Conv1到Conv5的5个特征图的空间分辨率依次降低2倍。为了获得类似于DilatedFCN的最终特征图,我们提出了一种新的联合上采样模块,JPU,它将最后三层特征图(Conv3-5)输入。之后经过一个多尺度语义模块(PSP/ASPP)或者一个全局语义模块(Encoding)来产生最终的预测输出。
对比DilatedFCN,我们的方法快了4倍,内存占用减小了。因此我们的方法比DilatedFCN又快又省空间。
JPU
JPU生成特征图的方法和DilatedFCN骨架的最终特征图激活相似,可以被形式化为联合上采样,最终被CNN所用。
背景
联合上采样(Joint Upsampling)
给定一个低分辨率的目标图像和一个高分辨率的指导图像,联合上采样旨在从指导图像中转换出结构和细节来生成高分辨率的目标图像。一般而言,低分辨率的目标图像
y
l
y_l
yl由低分辨率的指导图像
x
l
x_l
xl通过变换
f
(
⋅
)
f(\cdot)
f(⋅)生成,也就是说:
y
l
=
f
(
x
l
)
y_l=f(x_l)
yl=f(xl)。给定
x
l
x_l
xl和
y
l
y_l
yl,我们需要获得变换
f
^
(
⋅
)
\hat{f}(\cdot)
f^(⋅)来趋近
f
(
⋅
)
f(\cdot)
f(⋅),其中
f
^
(
⋅
)
\hat{f}(\cdot)
f^(⋅)的计算复杂度远低于
f
(
⋅
)
f(\cdot)
f(⋅)。举例而言,如果
f
(
⋅
)
f(\cdot)
f(⋅)是多层感知机MLP,那么
f
^
(
⋅
)
\hat{f}(\cdot)
f^(⋅)可以被简化为线性变换。高分辨率的目标图像
y
h
y_h
yh就通过
f
^
(
⋅
)
\hat{f}(\cdot)
f^(⋅)作用在高分辨率的指导图像
x
h
x_h
xh来获得。严格来说,给定
x
l
x_l
xl、
y
l
y_l
yl和
x
h
x_h
xh,联合上采样定义如下:
y
h
=
f
^
(
x
h
)
,
w
h
e
r
e
f
^
(
⋅
)
=
arg
min
h
(
⋅
)
∈
H
∣
∣
y
l
−
h
(
x
l
)
∣
∣
y_h=\hat{f}(x_h), \ where\ \hat{f}(\cdot)=\mathop{\arg\min_{h(\cdot)\in\mathcal{H}}}||y_l-h(x_l)||
yh=f^(xh), where f^(⋅)=argh(⋅)∈Hmin∣∣yl−h(xl)∣∣
其中
H
\mathcal{H}
H是所有可能的变换函数的集合,
∣
∣
⋅
∣
∣
||\cdot||
∣∣⋅∣∣是预定义的距离度量。
扩张卷积
扩张卷积由DeepLab提出,为了在保证感受野的前提下获得更高的空间分辨率。图3a展示了1D(扩张率=2)的扩张卷积,可以被分成以下三步:(1) 根据索引的奇偶性将输入特征图 f i n f_{in} fin分成 f i n 0 f^0_{in} fin0和 f i n 1 f^1_{in} fin1两组,(2) 将每个特征图经过同一个卷积层,得到 f o u t 0 f^0_{out} fout0和 f o u t 1 f^1_{out} fout1,(3) 交错合并两个图生成 f o u t f_{out} fout。
步幅卷积
步幅卷积使输入特征图转换为空间分辨率降低的输出特征图,等效于图3b所示的两步:(1) 用常规的卷积处理输入特征图
f
i
n
f_{in}
fin得到中间特征图
f
m
f_m
fm,(2) 移除奇数索引的元素得到
f
o
u
t
f_{out}
fout。
重新构造Joint Upsampling
我们方法和DilatedFCN不同的地方在最后的两层。以第4层为例,DilatedFCN中,输入特征图先被常规的卷积处理,然后再经过一系列的扩张卷积(d=2)。而我们的方法先用stride=2的卷积处理输入特征图,然后采用几个常规的卷积来生成输出特征图。
严格来说,给定一个输入特征图
x
x
x,DilatedFCN的输出特征图
y
d
y_d
yd获得方式如下:
y
d
=
x
→
C
r
→
C
d
→
.
.
.
→
C
d
=
x
→
C
r
→
S
C
r
M
→
.
.
.
→
S
C
r
M
=
x
→
C
r
→
S
→
C
r
→
.
.
.
→
C
r
→
M
=
y
m
→
S
→
C
r
→
.
.
.
→
C
r
→
M
=
{
y
m
0
,
y
m
1
}
→
C
r
n
→
M
\begin{aligned} y_d &= x \rightarrow C_r \rightarrow C_d \rightarrow ...\rightarrow C_d\\ &= x \rightarrow C_r \rightarrow SC_rM \rightarrow ...\rightarrow SC_rM\\ &= x \rightarrow C_r \rightarrow S \rightarrow C_r \rightarrow ...\rightarrow C_r \rightarrow M\\ &= y_m\rightarrow S \rightarrow C_r \rightarrow... \rightarrow C_r \rightarrow M\\ &= \{y^0_m,y^1_m\} \rightarrow C^n_r \rightarrow M \end{aligned}
yd=x→Cr→Cd→...→Cd=x→Cr→SCrM→...→SCrM=x→Cr→S→Cr→...→Cr→M=ym→S→Cr→...→Cr→M={ym0,ym1}→Crn→M
在我们的方法中,特征图
y
s
y_s
ys如下生成:
y
s
=
x
→
C
s
→
C
r
→
.
.
.
→
C
r
=
x
→
C
r
→
R
→
C
r
→
.
.
.
→
C
r
=
y
m
→
R
→
C
r
n
=
y
m
0
→
C
r
n
\begin{aligned} y_s &= x \rightarrow Cs \rightarrow C_r \rightarrow ...\rightarrow C_r\\ &= x \rightarrow C_r \rightarrow R \rightarrow C_r \rightarrow ...\rightarrow C_r\\ &= y_m \rightarrow R \rightarrow C^n_r=y^0_m \rightarrow C^n_r \end{aligned}
ys=x→Cs→Cr→...→Cr=x→Cr→R→Cr→...→Cr=ym→R→Crn=ym0→Crn
C
r
C_r
Cr、
C
d
C_d
Cd和
C
s
C_s
Cs分别代表了常规、扩张、步幅卷积,
C
r
n
C^n_r
Crn代表了
n
n
n层常规卷积。
S
S
S、
M
M
M和
R
R
R分别代表了图3中的分离(split)、合并(merge)、缩小(reduce)操作,其中
S
S
S、
M
M
M可以被去掉。值得一提的是,上两个公式是1D简化条件下的,当然在2D下同样也可以得到相同的结论。
上述提到的公式显示了
y
s
y_s
ys和
y
d
y_d
yd可以通过不同输入(
y
m
0
y^0_m
ym0和
y
m
y_m
ym,前者从后者下采样得到)经过同一个函数
C
r
n
C^n_r
Crn后获得。因此,给定
x
x
x和
y
s
y_s
ys,特征图
y
y
y可以得到:
y
=
{
y
m
0
,
y
m
1
}
→
h
^
→
M
w
h
e
r
e
h
^
=
arg
min
h
∈
H
∣
∣
y
s
−
h
(
y
m
0
)
∣
∣
y=\{y^0_m,y^1_m\} \rightarrow \hat{h} \rightarrow M\\ where\ \hat{h}=\mathop{\arg\min_{h\in\mathcal{H}}}||y_s-h(y^0_m)||
y={ym0,ym1}→h^→Mwhere h^=argh∈Hmin∣∣ys−h(ym0)∣∣
和公式1中定义的联合上采样相同。当然对于Conv5也可以获得相似的结论。
用CNN解决问题
公式4是一个优化问题,需要花很多次的梯度下降迭代收敛。取而代之的是,我们提出了一种用CNN模块的近似优化方法。首先我们如公式4所示,给定 x x x,需要生成 y m y_m ym,然后特征图 y m 0 y^0_m ym0和 y s y_s ys聚集起来学习映射 h ^ \hat{h} h^。最后,聚集的特征图经过一个卷积模块输出最终的预测 y y y。
根据上文的分析,我们设计了如图4的JPU模块。具体来说,每个输入特征图首先先经过一个常规卷积模块(图4a),作用是(1) 将 x x x变成 y m y_m ym,(2) 把 f m f_m fm降维进嵌入空间中。结果是,所有输入特征都被映射到同一个空间,便于融合和降低计算复杂度。
接着,生成的特征图经过上采样和合并操作,得到
y
c
y_c
yc(图4b)。四个分离的不同扩张率的扩张卷积并行抽取
y
c
y_c
yc特征,其中不同的扩张率为不同的函数。具体来说,扩张率为1的卷积捕获
y
m
0
y^0_m
ym0和
y
m
y_m
ym的剩余部分的关联,如图5蓝色框所示。扩张率为2、4、8的卷积用于学习将
y
m
0
y^0_m
ym0转换为
y
s
y_s
ys的映射
h
^
\hat{h}
h^,如图5绿框所示。因此JPU可以从多尺度特征图中抽取出多尺度的语义信息,达到更好的性能。这和ASPP有明显不同,后者仅挖掘最后一层特征图的信息。
图5:扩张率为1的卷积关注
y
m
0
y^0_m
ym0和
y
m
y_m
ym的剩余部分的关联,扩张率2的卷积关注
y
m
0
y^0_m
ym0和
y
s
y_s
ys。
抽取出的特征编码了从 y m 0 y^0_m ym0到 y s y_s ys和 y m 0 y^0_m ym0到 y m y_m ym其他部分的映射。接下来使用一个常规卷积模块将特征转换为最终的预测图(图4c)。
值得一提的是,所提出的JPU模块同时解决了两个非常相关的联合上采样问题:(1) 基于Conv3的上采样到Conv4,(2) 在Conv4指导下放大了Conv5。