论文阅读-CompletionFormer

1 背景

在给定稀疏深度和相应的RGB图像的情况下，深度补全的目的是将稀疏测量值在空间上传播到整个图像中，得到密集的深度预测。尽管基于深度学习的深度补全方法取得了巨大的进步，但卷积层或图模型的局部性使得网络很难对像素之间的远程关系进行建模。Transformer虽然有全局感受野，可以对远程关系建模，但是其局部特征建模效果并不入CNN。
本文提出了联合卷积注意力和变压器块(Joint Convolutional Attention and Transformer block, JCAT)，它将卷积注意层和视觉变压器深度耦合为一个块，作为构建金字塔结构深度补全模型的基本单元。这种混合架构自然有利于在单个模型中使用卷积的局部连接和Transformer的全局上下文。
因此，本文的CompletionFormer方法在室外KITTI深度补全基准测试和室内NYUv2数据集上优于最先进的基于CNN的方法，与纯基于transformer的方法相比，效率更高(近1/3的FLOPs)。
总的来说，有两点创新：
（1）提出了JCAT(Joint Convolutional Attention and Transformer)模块，结合CNN的局部优势和Transformer的全局优势，有利于多模态信息(RGB和深度图)的交互和融合。
（2）基于JCAT模块搭建了CompletionFormer网络，效果优于纯CNN模型，效率优于纯Transformer模型。

2 方法

2.1 总体结构

图2-1是整个模型的结构示意，在获得深度和RGB图像embedding之后，由JCAT块构造的主干用于多尺度的特征提取，并且解码器提供用于初始深度预测的全分辨率特征。最后，为了从稀疏输入中保持准确的深度，作者用SPN来细化初始估计。

图2-1 CompletionFormer整体结构

2.2 RGB和深度图嵌入

对于深度补全，在早期进行多模态的融合有几个优点：
（1）它使每个像素的特征向量同时具有RGB和深度信息，使得具有无效深度的像素仍然有机会根据外观相似性通过可靠的深度值来校正；
（2）对于后续网络，仅需要一个分支，相比两个分支，整个网络的效率更高。
因此，作者首先使用两个单独的卷积来对输入稀疏深度图$S$和RGB图像$I$进行编码。输出被拼接并由另一个卷积层进一步处理，以获得包含来自两个源的内容的原始特征。

2.3 JCAT模块

如何在像素之间建立连接以从可靠的像素实现深度传播，同时避免不正确的像素这件事已经被广泛地研究了。最近，卷积层或是基于注意力的图传播就是是新这个目标的主流操作。虽然完全基于transformer的网洛也已被用于此目的，但与纯基于CNN的方法相比，它显示出更差的结果和更高的计算成本。考虑到这两种风格的操作的互补性质，深度补全任务需要这两种范式的优雅集成。另一方面，对于分类和目标检测任务，MPViT和CMT是两个代表性的最先进的自注意力和卷积相结合的网络，如图2-2（a）和（b）所示。通常，可以以并行或级联方式实现集成。

图2-2 MPViT和CMT结构

因此，受其设计的启发，在CompletionFormer内，作者提出了如图2-3（c）和（d）所示的接头设计。为了减少计算开销，同时获得高度准确的深度补全结果，本文的CompletionFormer只包含单个而不是多个耗时的基于Transformer的路径，如MPViT。此外，基于卷积的路径的表示能力因空间和通道注意力而增强。

图2-3 JCAT结构

具体来说，本文的编码器有五个阶段，允许在不同尺度的特征表示进行有效地沟通。在第一阶段，为了减少Transformer层引入的计算成本和内存开销，作者使用ResNet 34中的一系列BasicBlocks进行处理，并最终获得$\frac{1}{2}$分辨率的下采样特征图$F_1$。对于接下来的四个阶段，作者介绍了提出的JCAT模块作为框架设计的基本单元。
对于每个阶段i∈{2,3,4,5}i \in \{2,3,4,5\}i∈{2,3,4,5}，它由一个patch embedding模块和$L_i$个重复的JCAT块组成。Patch embedding模块首先将上一层$i-1$输出的特征$F_{i-1}$切分成尺寸为$2\times 2$的patches。作者使用$3\times 3$卷积层和步幅设置为2来实现它，因此它实际上将特征的分辨率减半，从而获得特征金字塔{F2,F3,F4,F5}\{F_2, F_3, F_4, F_5\}{F2​,F3​,F4​,F5​}，其分辨率相对于输入图像为{1/4,1/8,1/16,1/32}\{1/4,1/8,1/16,1/32\}{1/4,1/8,1/16,1/32}。此外，position embedding也包括在embedded patches中，并通过JCAT块传递。
总的来说，本文的JCAT块可以以并行或级联的方式组织，如图2-3（c）和（d）所示。Transformer层以高效的方式实现，如Pyramid Vision Transformer，其中包含具有多头机制的spatial-reduction attention（SRA）层和前馈层（FNN）。
给定从patch embedding模块或者上一个节点block得到的输入特征$F\in R^{H_i\times W_i\times C}$，首先使用layer normalization(LN)进行归一化，然后展开成向量$X\in R^{N\times C}$，其中$N=H_i\times W_i$。使用线性转换器$W^Q$，$W^K$和$W^V\in R^{C\times C}$，$X$被映射到对应的$Q$，$K$，$V\in R^{N\times C}$。自注意力可以表示为
$$Attention(Q,K,V)=Softmax(\frac{Q{K}^T}{\sqrt{C_{head}}})V\text{\hspace{1em}(2-1)}$$
其中，$C_{head}$表示每个SRA的attention head的通道数。
根据式2-1，整个输入空间中的每个记号与任何记号匹配，包括其自身。深度补全网络从两个方面受益于自注意机制：
（1）它将网络的接收域扩展到每个Transformer层的完整图像；
（2）由于已经用深度和RGB图像信息嵌入了每个标记，因此自注意机制不仅通过外观，而且通过点积运算通过深度来显式地比较每个像素的相似性。
因此，可以将可靠的深度信息广播到整个图像，使得能够校正错误像素。
作者通过通道和空间注意力提升卷积路径的表征能力。一方面，有助于对局部精确注意力进行建模，降低噪声。另一方面，由于卷积和Transformer之间存在语义鸿沟，使用注意力机制的建模能力使得该路径能够关注Transformer层提供的重要特征，同时抑制不必要的特征。最后，用3 × 3的卷积将两条路径融合并输入到下一个块或阶段。
以提出的JCAT块为基本单元，作者构建具有重复配置的阶段2-5。如表2-1所示，作者将CompletionFormer中的4个阶段从微小、小型扩展到基本规模。

表2-1 不同规模的CompletionFormer

作者对比了使用图2-3中（c）和（d）两种结构的JCAT，（c）的效果会更好一些，如表2-2所示。

表2-2 CompletionFormer不同设置效果对比

2.4 解码器

在解码器中，来自每个编码层的输出被级联，并且经由skip connections由对应的解码层进一步处理。为了更好地适应不同的尺度特征，来自先前解码器层的特征通过反卷积层被上采样到当前尺度，并且卷积注意力机制也被利用来加强通道和空间维度中的特征融合。最后，来自解码器的融合结果与来自阶段1的特征级联，并被喂到预测头的第一卷积层。它的输出与来自RGB和深度嵌入模块的原始特征连接并被发送到另一个卷积，该卷积负责初始深度预测$D^0$。

2.5 SPN模块

考虑到稀疏输入的准确深度值在经过U-Net后可能无法很好地保留，SPN已成为最终细化的标准操作。最近的工作主要集中在将SPN从fixed-local改进为nonlocal传播。在本文的实验中，作者观察到，通过本文增强的U-Net主干，网络能够提供良好的深度affinity，从而fixed-local或nonlocal用于空间传播获得几乎相同的精度。
由于CSPN++消耗更多的计算成本，作者采用非局部空间传播网络NLSPN进行进一步深度细化。具体可以看我关于NLSPN的博客，这里简单说下。
使得$D^t=(d_{u,v}^t)\in R^{H\times W}$表示在第$t$步经过空间传播更新的2D深度图，其中$d_{u,v}^t$表示在像素$(u,v)$的深度值。$(H\times W)$表示$D^t$的高和宽。$d_{u,v}^t$在第$t$步使用non-local的邻域$N_{u,v}^{NL}$进行传播的公式为
$$d_{u,v}^t=w_{u,v}(0,0)d_{u,v}^{t-1}+\sum _{(i,j)\in N_{u,v}^{NL},i\ne 0,j\ne 0}{w}_{u,v}(i,j)d_{i,j}^{t-1}\text{\hspace{1em}(2-2)}$$
其中，$w_{u,v}(i,j)\in (-1,1)$表示参考点和其邻居的affinity weight。
$$w_{u,v}(0,0)=1-\sum _{(i,j)\in N_{u,v}^{NL},i\ne 0,j\ne 0}{w}_{u,v}(i,j)\text{\hspace{1em}(2-3)}$$
$w_{u,v}(0,0)$表示原始深度$d_{u,v}^{t-1}$有多少会被保留。
此外，亲和度矩阵$w$也由解码器输出并且由来自解码器的预测置信度图调制，以防止不太置信的像素无视亲和度有多大传播到邻居中。经过$K$步空间传播，得到最终的细化深度图$D^K$。
SPN模块对比大部分的backbone而言都是有正提升的，可见表2-3所示。

表2-3 有无SPN效果对比表

2.6 损失函数

最终的损失为$L_1$和$L_2$的组合
$$L(\hat{D},D^{gt})=\frac{1}{|V|}\sum _{v\in V}(|{\hat{D}}_v-D_v^{gt}|+|{\hat{D}}_v-D_v^{gt}{|}^2)\text{\hspace{1em}(2-4)}$$
其中，$\hat{D}=D^K$，$V$是真值$D^{gt}$中的有效像素集合，$|V|$表示$V$的size。
$L_1$损失对误差"线性惩罚"，不会像$L_2$那样放大误差，效果更佳鲁棒，但是$L_1$在极小值点不可倒，很难收敛到极小值点。
$L_2$损失处处可导，但是对于异常值效果不稳定，导数会很大。
$L_1$和$L_2$组合使用可以实现优势互补。

3 效果

不同SOTA模型之间的效果对比，如下表3-1所示。

表3-1 不同SOTA模型效果对比

对比不同稀疏程度的深度图输入，不同模型之间的效果对比可见图3-1，其中Ours-ViT表示只使用了JCAT中的Transformer部分。

图3-1 不同稀疏程度深度图输入效果对比

不同模型之间的可视化效果对比，如图3-2所示。

图3-2 不同SOTA模型的可视化效果对比

4 结论

本文提出了一种单分支深度补全网络CompletionFormer，将卷积注意力和Transformer无缝集成到一个块中。大量的消融研究证明了本文的模型在输入稀疏时深度补全的有效性和高效性。这种新颖的设计在室内和室外数据集上产生了最先进的结果。目前，CompletionFormer的运行时间约为10FPS，进一步降低其运行时间以满足实时需求将是值得研究的部分。

参考资料

[1] CompletionFormer: Depth Completion with Convolutions and Vision Transformers