论文阅读-RaftStereo

1 概述

在双目立体匹配中，基于迭代的模型是一种比较主流的方法，而其鼻祖就是本文要讲的RaftStereo。
先来说下什么是双目立体匹配。给定极线矫正后的左图和右图$(I_L,I_R)$，双目立体匹配的目的就是估计一张视差图$d$，$d$中每个像素位置的视差值$d_{xy}$表示在左图$I_L$的$(x,y)$位置的像素与右图$I_R$中的$(x-d_{xy},y)$位置的像素在真实3D世界中对应于同一个物理点。换句话说，$d$表示了左图和右图在水平方向上的匹配关系。
有了视差图$d$，左右相机的焦距$f_x$和基线$B$，我们就可以通过式(1-1)得到左图的深度图$D$，这部分具体可以参考这篇。
$$D=\frac{Bf}{d}\text{\hspace{1em}(1-1)}$$
RaftStereo的输入是左右图$(I_L,I_R)$，输出是视差$d$。由于视差$d$本身与相机参数无关，因此该方式训练得到的模型可以应用于不同参数的相机。
RaftStereo参考了Raft，由特征抽取器，相关金字塔和基于GRU的迭代更新算子组成，可见图1-1。

2 模块说明

2.1 特征抽取器

如图1-1所示，特征抽取器指的就是encoder，共有两种encoder，一个叫做Feature Encoder，另一个叫做Context Encoder。
Feature Encoder同时作用于左、右图像，将每幅图像映射为密集特征图，进而构建相关体。该网络由一系列残差块和下采样层组成，根据实验中使用的下采样层数不同，可生成分辨率为输入图像1/4或1/8且包含256个通道的特征图。在Feature Encoder采用了Instance Normalization。
Context Encoder的架构与特征编码器完全相同，不同之处在于使用Batch Normalization替代了Instance Normalization，并且仅在左图上应用Context Encoder。上下文特征用于初始化更新算子的隐藏状态，并在每次迭代更新算子时输入GRU中。
关于feature encoder使用IN，context encoder使用BN的疑问。有人在github的issue上面也问了这个问题，可见https://github.com/princeton-vl/RAFT-Stereo/issues/17。
feature encoder中的IN比BN更有意义，因为作者希望仅从单个图像中导出每组相关特征。context encoder也可以用IN，但是BN更考虑整个数据集的均值和方差。
有一个细节，作者的代码中对BN进行了freeze_bn的操作，对应的代码是m.eval()这只是将均值和方差限制为[0, 1]，对于 $\gamma$ 和 $\beta$ 并没有进行限制， $\gamma$ 和 $\beta$ 还是会进行梯度更新的。回顾一下BN的公式：
$$y=\frac{x-\mathrm{E}[x]}{\sqrt{\mathrm{Var}[x]+ϵ}}\ast \gamma +\beta$$

2.2 相关金字塔

（1）相关体
采用特征向量之间的点积作为视觉相似度的衡量指标，对具有相同y坐标的像素进行相关计算。给定分别从$I_L$和$I_R$提取的特征图$f,g\in R^{H\times W\times D}$，将内积运算限制在具有相同第一索引的特征向量之间进行，可得到相关体，如下式(2-1)所示。
$$C_{ijk}=\sum _h{f}_{ijh}\cdot g_{ikh},C\in R^{H\times W\times W}\text{\hspace{1em}(2-1)}$$
式(2-1)看不明白的话，可以先不看$i$，也就是我们只看$y$方向确定的情况下，同一行的相关体计算方法。很容易看出$C_{jk}$就表示左图$j$位置与右图$k$位置之间的相关性，这个相关性就是长度为$D$的向量$f_{ij}$和$g_{ik}$的点积。
相关体的计算可以通过单个矩阵乘法来高效实现，这可以在GPU上轻松计算，并且只占用总运行时间的一小部分。
由于模型输入是极线校正后的图像，因此可以假设所有视差均为正值。因此，相关体的计算实际上只需针对正视差进行。然而，完整体计算的优势在于其运算可通过矩阵乘法实现，而该运算经过高度优化。这简化了整体架构设计，使作者能够采用通用运算而非定制GPU内核。
（2）相关金字塔
通过重复对最后一维进行平均池化，构建了一个四层结构的相关体金字塔。金字塔的第$k$层是通过使用核尺寸为2、步长为2的一维平均池化操作得到，从第$k$层的体积生成的相关体$C_{k+1}$，其维度为$H\times W\times W/2^k$。虽然金字塔的每一层都具有更大的感受野，但通过仅对最后一维进行池化处理，保留了原始图像中的高分辨率信息，从而能够恢复极其精细的结构特征。
（3）相关查找
为了在相关金字塔中建立索引，作者定义了一个类似于Raft中定义的查找算子$L_C$。给定当前视差估计值$d$，作者会围绕该估计值构建一个带有整数偏移量的一维网格（如图2-1所示）。这个网格用于从相关金字塔的每一层进行索引操作。

由于网格的值都是整数，在每个相关体取索引时使用了线性插值。取回的值最终被拼接到一个特征图中。这部分的代码可以参考下面的x0变量的计算方式。

class CorrBlock1D:
	def __call__(self, coords):
		r = self.radius
		coords = coords[:, :1].permute(0, 2, 3, 1)
		batch, h1, w1, _ = coords.shape
	
		out_pyramid = []
		for i in range(self.num_levels):
			corr = self.corr_pyramid[i]
			dx = torch.linspace(-r, r, 2*r+1)
			dx = dx.view(2*r+1, 1).to(coords.device)
			x0 = dx + coords.reshape(batch*h1*w1, 1, 1, 1) / 2**i
			y0 = torch.zeros_like(x0)
	
			coords_lvl = torch.cat([x0,y0], dim=-1)
			corr = bilinear_sampler(corr, coords_lvl)
			corr = corr.view(batch, h1, w1, -1)
			out_pyramid.append(corr)
	
		out = torch.cat(out_pyramid, dim=-1)
		return out.permute(0, 3, 1, 2).contiguous().float()

2.3 多级更新算子

从初始点$d_0=0$出发，预测一系列视差场 { d 1 , . . . , d N } \{ d_1, ..., d_N \} {d1​,...,dN​}。在每次迭代中，利用当前视差估计值对相关体进行索引，生成一组相关特征。这些特征会依次通过两个卷积层处理，同时当前视差估计值也会经过两次卷积层处理。随后将相关特征、视差特征和上下文特征拼接后输入GRU，由GRU更新隐藏状态。最终利用更新后的隐藏状态来预测视差的动态变化。
（1）多个隐藏状态
原始的Raft算法完全采用固定高分辨率进行更新。这种方法存在一个明显缺陷：感受野面积会随着GRU更新次数的增加而缓慢增长，这在处理具有大面积无纹理区域且局部信息匮乏的场景时容易引发问题。为解决这一难题，提出了一种多分辨率更新算子，该算子可同时作用于1/8、1/16和1/32分辨率的特征图。实验结果表明，采用这种多分辨率更新算子能显著提升模型的泛化性能。
如图2-2所示，GRU通过相互使用对方的隐藏状态作为输入进行交叉连接。相关查找和最终的视差更新由分辨率最高的GRU完成。

（2）上采样
预测的视差场分辨率是输入图像的1/4或1/8。为了输出全分辨率视差图，采用了与Raft相同的凸上采样方法。RAFT-Stereo将全分辨率视差值视为其粗分辨率邻近点3x3网格的凸组合结果，而凸组合权重则由最高分辨率GRU进行预测。
上采样部分的代码为

def upsample_flow(self, flow, mask):
	""" Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
	N, D, H, W = flow.shape
	factor = 2 ** self.args.n_downsample
	mask = mask.view(N, 1, 9, factor, factor, H, W)
	mask = torch.softmax(mask, dim=2)

	up_flow = F.unfold(factor * flow, [3,3], padding=1)
	up_flow = up_flow.view(N, D, 9, 1, 1, H, W)

	up_flow = torch.sum(mask * up_flow, dim=2)
	up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
	return up_flow.reshape(N, D, factor*H, factor*W)

其中的mask就是最高分辨率GRU预测得到的凸组合权重。

2.4 Slow-Fast GRU

将GRU更新到1/8分辨率的隐藏状态，其浮点运算量大约是更新1/16分辨率隐藏状态的四倍。为了利用这一特性提升推理速度，作者对RAFT-Stereo模型进行了优化，每当更新1/8分辨率隐藏状态时，就会同步多次更新1/16和1/32分辨率的隐藏状态。在KITTI分辨率图像上进行32次GRU更新后，这一改进使RAFTStereo的运行时间从0.132秒缩短至0.05秒，效率提升达52%。
这种改进使RAFT-Stereo在实时立体视觉方面具有竞争力的性能，而运行的方法快一个数量级，如下表2-1的Slow-Fast行所示。

作者发现，通过增加低分辨率GRU的迭代次数并减少高分辨率GRU的迭代次数，在略微降低精度的情况下，RAFT-Stereo的运行时间显著缩短。Slow-Fast版本的RaftStereo分别将最低、中等和最高分辨率的隐藏状态更新30次、20次和10次，而Regular版本则将每个隐藏状态更新32次。无论是Slow-Fast还是Regular版本，都使用相同的模型权重。Slow-Fast-GRU这部分的代码如下所示

for itr in range(iters):
    coords1 = coords1.detach()
    corr = corr_fn(coords1) # index correlation volume
    flow = coords1 - coords0
    with autocast(enabled=self.args.mixed_precision):
		if self.args.n_gru_layers == 3 and self.args.slow_fast_gru: # Update low-res GRU
			net_list = self.update_block(net_list, inp_list, iter32=True, iter16=False, iter08=False, update=False)
		if self.args.n_gru_layers >= 2 and self.args.slow_fast_gru:# Update low-res GRU and mid-res GRU
			net_list = self.update_block(net_list, inp_list, iter32=self.args.n_gru_layers==3, iter16=True, iter08=False, update=False)
		net_list, up_mask, delta_flow = self.update_block(net_list, inp_list, corr, flow, iter32=self.args.n_gru_layers==3, iter16=self.args.n_gru_layers>=2)

从代码中不难看出，self.args.slow_fast_gru=True时，计算量反而会增加，并没有起到提速的效果。这是为啥呢？
使用Slow-Fast-GRU时的做法是将总的迭代次数减少，32->7，但是这样的效果会变差。但现在模型的做法都是在Slow-Fast-GRU为True的情况下，同时设置迭代次数为32，这其实更慢了。在github issue上面其他人的提问和作者解答可见https://github.com/princeton-vl/RAFT-Stereo/issues/25。

2.5 监督

作者在预测序列 { d 1 , . . . , d N } \{ d_1, ..., d_N \} {d1​,...,dN​}中，通过指数级递增的权重对预测视差与真实视差之间的$L1$距离进行监督。给定真实视差$d_{gt}$时，损失函数定义为
$$L=\sum _{i=1}^N{\gamma }^{N-i}||d_{gt}-d_i|{|}_1\text{\hspace{1em}(2-2)}$$
其中，$\gamma =0.9$。

3 效果

作者对比了RaftStereo与其他双目深度估计算法的泛化能力，这些算法均在虚拟场景数据集Sceneflow上进行训练，并在真实场景数据集上进行测试，可见表3-1。表中的数值表示误差，即误差超过指定阈值的像素百分比。作者采用标准评估阈值是KITTI为3像素，Middlebury为2像素，ETH3D为1像素。

不同模型之间的可视化效果对比可见图3-1，可见RaftStereo对于细结构的物体效果更好。