《RAFT-Stereo：Multilevel Recurrent Field Transforms for Stereo Matching》论文笔记

Posted 2021-11-07 m_buddy

参考代码：RAFT-Stereo

1. 概述

导读：RAFT算法是非常经典的立体匹配算法，在光流和立体匹配任务中有着广泛的运用。而这篇文章正是基于RAFT并将其运用到了立体匹配中，并且在如下的几个方面进行改进：
1）相比原生的RAFT算法钟重点关注X轴（W方向）的视差信息，前提是输入的图像对需要事先经过极线校准；
2）在GPU的update阶段使用stride为$[8,16,32]$的特征图进行运算，这样可以在迭代优化的同时更加增大网络感受野，从而增加对与大范围无/弱纹理区域的适应性；
通过上述两点的改进文章的立体匹配算法在现有的一些基准测试集上均取得了较为不错的结果，而且基于RAFT其泛化性能也是不错的。

文章的方法通过聚焦X维度视差，并且在update阶段采用多尺度优化的形式，极大增强了在极线对齐情况下网络的性能，文章的算法得到的效果见下图所示：

文章的算法在时间上的对比：

2. 方法设计

2.1 算法pipline

文章算法的pipeline见下图所示：

可以看到文章的pipeline主要由三个部分组成：左右视图特征抽取网络，使用左视图的内容特横抽取网络，基于Conv-GPU的迭代更新模块。

2.2 特征抽取模块

从图1中可以看到文章的特征抽取网络包含两个部分：

• 1）左右视图特征抽取：这里将左右视图输入同一网络得到它们的特征，需要注意的是这里使用了Instance Norm，参考：

# core/raft_stereo.py#L39
self.fnet = BasicEncoder(output_dim=256, norm_fn='instance', downsample=args.n_downsample)

• 2）context特征抽取：这里是对左视图进行特征抽取，这里使用的归一化为BatchNorm，参考：

# core/raft_stereo.py#L29
self.cnet = MultiBasicEncoder(output_dim=[args.hidden_dims, context_dims], norm_fn="batch", downsample=args.n_downsample)

2.3 基于Conv-GRU的迭代更新模块

2.3.1 相关性金字塔构建与X维度采样

回顾原生RAFT中是通过矩阵相乘，也就是torch.matmul()的形式实现相关性计算的，而在这篇文章中也是通过矩阵相乘的形式，只不过在后面对其中一个X维度进行归约得到，其公式为：
$$C_{ijk}=\sum _h{f}_{ijh}\cdot g_{ikh},C\in R^{H\ast W\ast W}$$
其代码实现可以参考：

# core/corr.py#L149
def corr(fmap1, fmap2):
    B, D, H, W1 = fmap1.shape
    _, _, _, W2 = fmap2.shape
    fmap1 = fmap1.view(B, D, H, W1)
    fmap2 = fmap2.view(B, D, H, W2)
    corr = torch.einsum('aijk,aijh->ajkh', fmap1, fmap2)
    corr = corr.reshape(B, H, W1, 1, W2).contiguous()
    return corr / torch.sqrt(torch.tensor(D).float())

通过上述的操作得到相关性矩阵，则在此基础上使用RAFT中的方式构建特征金字塔：

# core/corr.py#L122
self.corr_pyramid.append(corr)
for i in range(self.num_levels):
    corr = F.avg_pool2d(corr, [1,2], stride=[1,2])
    self.corr_pyramid.append(corr)

接下来就是在极线校准的左右视图上对X轴方向进行采样操作，采样的示意图可以参考下图：

对应的代码实现可以参考：

# core/corr.py#L133
for i in range(self.num_levels):
    corr = self.corr_pyramid[i]
    dx = torch.linspace(-r, r, 2*r+1)  # 在X维度（W方向）设置采样点
    dx = dx.view(1, 1, 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)  # 在X维度进行双线性插值采样
    corr = corr.view(batch, h1, w1, -1)
    out_pyramid.append(corr)

out = torch.cat(out_pyramid, dim=-1)

2.3.2 update过程中的多尺度Conv-GRU

原生的RAFT算法只在一个尺度上进行视差优化，因而可以从迭代的角度逐渐增加感受野，从而增强对无/弱纹理区域的适应性。对此，文章从多尺度的角度去构建多尺度Conv-GRU模块去显示增加更新模块的感受野，其更新示意图见下图所示：

其对应的实现可以参考：

# core/update.py#L104
self.gru08 = ConvGRU(hidden_dims[2], encoder_output_dim + hidden_dims[1] * (args.n_gru_layers > 1))
self.gru16 = ConvGRU(hidden_dims[1], hidden_dims[0] * (args.n_gru_layers == 3) + hidden_dims[2])
self.gru32 = ConvGRU(hidden_dims[0], hidden_dims[1])

对于上采样模块这里使用与RAFT一致的方式构建。需要注意的是使用大分辨率的特征图进行多尺度Conv-GRU迭代会带来较大的计算开销，对此文章引入了Slow-Fast GRU的概念，也就是大尺度特征处的GRU更新一次，对应小尺度特征处的GRU更新多次。

上述各个部分对性能的影响：

2.4 损失函数

这里的损失函数参考RAFT中的迭代算式函数形式，描述为：
$$L=\sum _{i=1}^N{\gamma }^{N-i}||d_{gt}-d_i|{|}_1,\gamma =0.9$$

3. 实验结果

KITTI-2015数据集：

ETH3D数据集：

Middlebury数据集：