RAFT

Teed 2020 · 论文

一句话总结 — 构建一个4D全对相关性体,并用一个权重共享的ConvGRU围绕当前估计值查询相关性,迭代细化单一的高分辨率光流场——ECCV 2020最佳论文,也是现代光流的代表性架构。

问题

主流的深度光流架构(PWC-Net及其同类)继承了经典的由粗到精金字塔:在低分辨率下估计光流,然后变形并细化。这种设计存在结构性的盲区——每一级的代价体只覆盖一个很小的搜索窗口,小而快速移动的物体在粗分辨率下会消失,金字塔早期阶段犯下的错误难以纠正,而多阶段级联往往需要超过100万次的训练迭代。此前的迭代细化方案没有在各次迭代之间共享权重(或者像IRR那样受限于其庞大的循环单元)。RAFT探究的问题是:如果网络预先计算所有像素对之间的匹配代价,并用一个轻量级的学习式优化器按需查询该体来细化单一的高分辨率光流场,会怎样?

方法与架构

三个阶段,均可微且端到端训练:

  1. 特征提取:一个编码器gθ:RH×W×3RH/8×W/8×Dg_\theta : \mathbb{R}^{H \times W \times 3} \mapsto \mathbb{R}^{H/8 \times W/8 \times D}(D=256D = 256,6个残差块)对两帧进行编码;一个结构相同的上下文网络hθh_\theta仅对I1I_1进行编码。两者每对帧只运行一次。
  2. 全对相关性:视觉相似度通过一次矩阵乘法为每一对像素预先计算好,

Cijkl=hgθ(I1)ijhgθ(I2)klh,CRH×W×H×WC_{ijkl} = \sum_h g_\theta(I_1)_{ijh} \cdot g_\theta(I_2)_{klh}, \qquad \mathbf{C} \in \mathbb{R}^{H \times W \times H \times W}

随后最后两个维度用核大小1、2、4、8做平均池化,构成金字塔{C1,C2,C3,C4}\{\mathbf{C}^1, \mathbf{C}^2, \mathbf{C}^3, \mathbf{C}^4\}。只对I2I_2的维度做池化,使I1I_1的维度保持在完整(1/8)分辨率——大位移和小位移都能被捕捉到,而不会丢失小而快速移动的物体。一个查询算子LCL_\mathbf{C}在当前对应关系x=x+f(x)\mathbf{x}' = \mathbf{x} + \mathbf{f}(\mathbf{x})周围的局部网格上对每一级做双线性采样,

N(x)r={x+dxdxZ2, dx1r}\mathcal{N}(\mathbf{x}')_r = \{ \mathbf{x}' + \mathbf{dx} \mid \mathbf{dx} \in \mathbb{Z}^2,\ \lVert \mathbf{dx} \rVert_1 \le r \}

在每一级kkN(x/2k)r\mathcal{N}(\mathbf{x}'/2^k)_r索引——恒定的半径在更粗的层级上对应更大的上下文范围(在k=4k=4时半径4对应原始分辨率下的256像素)。 3. 迭代更新:从f0=0\mathbf{f}_0 = \mathbf{0}出发,一个循环更新算子(仅270万参数,权重在所有迭代间共享)输入相关性查询结果、光流特征以及上下文特征xtx_t,并通过一个卷积GRU发出残差更新fk+1=fk+Δf\mathbf{f}_{k+1} = \mathbf{f}_k + \Delta\mathbf{f}:

zt=σ(Conv3×3([ht1,xt],Wz)),rt=σ(Conv3×3([ht1,xt],Wr))z_t = \sigma(\mathrm{Conv}_{3\times3}([h_{t-1}, x_t], W_z)), \qquad r_t = \sigma(\mathrm{Conv}_{3\times3}([h_{t-1}, x_t], W_r))

h~t=tanh(Conv3×3([rtht1,xt],Wh)),ht=(1zt)ht1+zth~t\tilde{h}_t = \tanh(\mathrm{Conv}_{3\times3}([r_t \odot h_{t-1}, x_t], W_h)), \qquad h_t = (1 - z_t) \odot h_{t-1} + z_t \odot \tilde{h}_t

该算子模拟了一个一阶优化器——但它不是使用泰勒线性化的数据项,而是学习提出下降方向;有界的激活函数促使其收敛到一个不动点,并且可以运行100次以上的迭代而不发散。光流在1/8分辨率下被预测,再通过在每个像素3x3粗邻域上学习到的凸组合进行上采样(权重通过softmax得到)。

监督覆盖了整个估计序列,权重随迭代次数指数增加:

L=i=1NγNifgtfi1,γ=0.8\mathcal{L} = \sum_{i=1}^{N} \gamma^{N-i} \lVert \mathbf{f}_{gt} - \mathbf{f}_i \rVert_1, \qquad \gamma = 0.8

训练依次在FlyingChairs、FlyingThings上进行,再在基准数据集上微调;在视频上,warm-start初始化会将前一帧的光流前向投影过去。

实验结果

对SLAM的意义

RAFT”相关性体+迭代循环细化”的方案成为了SLAM中学习式数据关联的主力工具:DROID-SLAM和DPVO本质上就是包裹着可微光束法平差层的RAFT式更新算子。其衍生方法(用于场景流的RAFT-3D,面向实时应用的SEA-RAFT)主导了光流基准测试,而它所推广的展开式学习优化模式如今已出现在稠密预测和SLAM系统的各个角落。

相关条目