光流DL——1.Learning Optical Flow from a Few Matches(2021)

Abstract

最先进的光流估计的神经网络模型需要一个高分辨率的密集相关体来表示每像素的位移。虽然密集的相关量为准确的估计提供了丰富的信息，但其大量的计算和内存使用阻碍了模型的有效训练和部署。在本文中，我们证明了密集相关体积表示是冗余的，并且在其中只有一小部分元素时就可以实现精确的流估计。在此基础上，我们提出了一种替代的位移表示，即稀疏相关体，它是通过计算一个特征映射中每个特征向量在另一个特征映射中的k个最接近匹配，并存储在稀疏数据结构中。实验表明，与以往具有密集相关体积的方法相比，该方法可以显著降低计算成本和内存使用，同时保持了较高的准确性。

1. Introduction

光流估计是计算机视觉[11]中的一个经典问题。它的目的是寻找两幅图像之间的像素级对应关系。传统上，它被表述为一个由连续[4,11,32]或离散[22,7,31]优化来解决的优化问题。随着深度学习的发展，光流估计已经被表述为一个学习问题，其中从神经网络直接回归成为一种常见的方法[9,16]。

在密集对应问题中，一个流行的表示是相关（成本）体积（the correlation (cost) volume,），由Hosni等人[12]首先提出。相关体积给出了每个像素位移的明确表示，并证明了它们在立体匹配[18]和光流[9,27]学习问题中的广泛应用。与搜索空间沿扫描线的立体匹配问题相反，光流问题具有二维搜索空间，这给直接处理4维卷带来了两个挑战：直接处理4维卷时的大内存消耗和高计算成本。

为了降低内存和计算成本，现有的方法[27,34,33,13,36]首先建立一个特征金字塔，在粗分辨率下计算相关体，然后基于上采样流逐步扭曲上层特征图，并在有限的搜索范围内构造一个局部相关体。在之前的工作[4,32,20]中观察到的一个值得注意的问题是，从粗到细的框架无法解决流动位移大于流动结构的情况，即著名的小物体快速移动问题。

最近的方法，Devon和RAFT[20,29]提出在第二张图像中使用直接搜索来消除扭曲的需要。RAFT特别证明了首先构造一个全对相关体，并直接以单一分辨率处理它，而不是以粗到细的方式的好处（就是精度最高）。然而，全对相关体积需要两个特征图之间的成对点积。因此，时间和空间的复杂度都是$O(N^2)$，其中N是一个图像的像素数。需要一个小的N来减少内存消耗，因此RAFT只能使用1/8分辨率的特征图。一个低分辨率的特征图不能完全表示一个图像的细节。我们想知道是否有一种方法来构建一个相关量与所有对搜索范围，但不超过最大GPU内存。因此，我们认为不必存储所有成对相关性，并假设只存储每个像素的顶部-k个相关性可能就足够了。

我们的直觉是，一幅图像中的特征向量在另一幅图像中只有少数具有高相关性的特征向量来匹配。因此，在密集的相关体中可能存在较大的冗余，其中小的相关性对预测没有贡献。图1说明了密集相关体积和稀疏相关体积之间的比较。

我们提出了一个稀疏相关体表示$(SparseCorrelationVolume)$，其中只有每个像素的顶部k个相关性存储在一个由{值，坐标}对定义的稀疏数据结构中。在本文中，我们演示了如何使用稀疏相关体积表示来解决光流问题。我们提出了一种在光流学习框架中构造和处理这种稀疏相关体的方法。我们证明，即使只存储了一小部分元素，我们的结果仍然可以与以前使用密集相关体积的工作[29]相比较。最后，我们证明了稀疏方法允许构建一个高分辨率的相关体，它可以比以前的方法更准确地预测精细结构的运动。

2. Method

设$I_1,I_2:{\mathbb{Z}}^2\to {\mathbb{R}}^3$是两张RGB图像。问题是估计一个密集的流场$f:{\mathbb{Z}}^2\to {\mathbb{R}}^2$，它将每个像素坐标x映射到一个位移向量f(x).

在现代深度学习光流方法中，首先应用特征提取网络从图像对$F_1,F_2:{\mathbb{Z}}^2\to {\mathbb{R}}^c$中提取特征映射，其中c为信道数。相关体积$C:Z^4\to R$是通过计算成对特征向量之间的内积而形成的$$C(\mathbf{x},\mathbf{d})=F_1(\mathbf{x})\cdot F_2(\mathbf{x}+\mathbf{d})\text{\hspace{1em}(1)}$$输出是一个四维张量，它可以表示为一个集合$$\mathcal{C}=\{C(\mathbf{x},\mathbf{d})\mid \mathbf{x}\in \mathcal{X},\mathbf{d}\in \mathcal{D}\}\text{\hspace{1em}(2)}$$其中$\mathcal{X}=[0,h)\times [0,w)\cap {\mathbb{Z}}^2$是特征图F1的定义域且$|\mathcal{X}|=hw$，其中，h和w分别表示F1的高度和宽度，位移集D被定义为$\mathcal{D}=[-d,d{]}^2\cap {\mathbb{Z}}^2$，其中D表示沿x或y方向的最大位移，$|\mathcal{D}|=(2d+1{)}^2$。因此，相关体积C包含$hw(2d+1{)}^2$个元素。

为了减少相关体积的大小，以前的方法使用从粗到细和扭曲的方法来约束d的大小。为了准确地处理大的位移，RAFT构造了一个全对相关体，其中位移范围包含整个特征图。排除范围外的匹配，RAFT的全对相关体积包含$N^2$个元素，其中N=hw. 因此，需要较低分辨率的特征图来约束N。在本工作中，我们证明了全对相关体积实际上可以是一个稀疏张量，其中只有一小部分的值被存储和处理。我们证明，我们可以有效地将空间复杂度从$O(N^2)$降低到$O(NK)$，而性能仅略有下降，其中K给出了我们想要保持的匹配数量。其主要思想如图2所示。

2.1. Sparse Correlation Volume

对于每个$x\in \mathcal{X}$，我们定义一个集合$$S_{\mathbf{x}}^{(k)}=\underset{S\subset \mathcal{D},|S|=k}{\mathrm{argmax}}\sum _{\mathbf{d}\in S}C(\mathbf{x},\mathbf{d})\text{\hspace{1em}(3)}$$给出了最大相关性的k个位移。相关体积现在可以表示为一个四维稀疏张量$$\tilde{\mathcal{C}}=\left\{C(\mathbf{x},\mathbf{d})\mid \mathbf{d}\in S_{\mathbf{x}}^{(k)},\mathbf{x}\in \mathcal{X}\right\}\text{\hspace{1em}(4)}$$这个稀疏的相关体积包含了hwk个元素，而不是原始的与$h^2{w}^2$元素的密集相关体积。常数k通常是一个很小的数字(例如k=8)。

我们现在展示如何构造稀疏相关体积并从中估计光流。我们的网络架构如图3所示。

(1) 构建四维稀疏相关体，首先用KNN计算一组具有top-k相关性的位移$d_0\in \tilde{\mathcal{K}}$；然后取F1中每个特征向量与F2中k个对应的特征向量的点积；虚线箭头表示没有梯度流的路径，而实线箭头表示有梯度流的路径。
(2) 在每次迭代中，通过减去剩余流量${\mathbf{d}}_i-\Delta {\mathbf{f}}_i$来更新位移向量，以更新四维相关量。采用多尺度位移编码器将4D稀疏相关量编码为二维致密运动张量。
应用一个GRU更新块来预测下一次迭代的剩余流量$∆f_{i+1}$。GRU块还输入表示当前迭代的隐藏状态向量的$h_i,F_c,f_i$，上下文网络提取的特征图和当前光流估计，并输出下一次迭代的隐藏状态向量和残余流。

2.2. k-Nearest Neighbours

我们首先使用两个权重共享特征提取网络从输入图像中提取1/4分辨率的特征图。我们的特征提取网络由6个残差块组成，特征通道数为256个。为了构造稀疏相关体，我们使用一个近邻(kNN)模块[17]来计算F1中每个特征向量具有k个最大相关得分的一组指标。稀疏相关量的计算方法是取F1中每个特征向量与F2中的指标给出的前k个特征向量之间的点积。在反向传播过程中，梯度只传播到kNN模块选择的k个特征向量。

2.3. Displacements Update

我们采用了一种整体的迭代残差细化方法。如之前的工作[15,29]所示，估计残差流比直接回归[16,9]可以有效地减少搜索空间，预测更好的结果。不是直接预测光流$f$，而是预测剩余流$∆f_{i+1}$，并用于更新当前流量估计$f_{i+1}=f_i+∆f_{i+1}$.

在每一步中，根据当前的光流估计，将F1中的像素$\mathbf{x}$映射到F2中的${\mathbf{x}}_i$. 在第2.1节中描述的稀疏相关量$\tilde{C}$可以看作是在i=0处的初始估计$f_0=0$。当坐标${\mathbf{x}}_i$更新到${\mathbf{x}}_{i+1}={\mathbf{x}}_i+∆f_i$时，$\tilde{C}$中的相对位移也应相应地更新。为此，我们通过在每一步$C_i\left(\mathbf{x},{\mathbf{d}}_i\right)=C_{i+1}\left(\mathbf{x},{\mathbf{d}}_i-\Delta {\mathbf{f}}_i\right)$从$d_i$中减去最近的$∆f_i$来移动稀疏相关张量的坐标，如图4a所示。注意，这里我们允许${\mathbf{d}}_i-\Delta {\mathbf{f}}_i$作为浮动点。同样重要的是，内积在开始时只计算一次，因为在每个步骤中，只有相关坐标发生变化，而相关值保持不变。

2.4. Multi-scale Displacement Encoder

任何稀疏方法经常提出的一个问题是如何处理稀疏张量，因为一个正常密集的$h\times w\times c$张量的规律性丢失了。可以使用稀疏卷积[8]，但是，我们将在这里提出一个更简单和更有效的方法。

一个密集的全对相关体的维数为$h\times w\times h\times w$，我们将其简化为一个具有$h\times w\times k$元素的稀疏张量，其中只有每个像素的top-k相关被保存。我们可以看到，前两个维度仍然是密集的，而已经变得稀疏的是第三个和第四个维度。这里的目标是为每个像素的k个元素进行编码，并形成一个密集的$h\times w\times c$张量，稍后可用于预测$∆f_{i+1}$.

遵循之前的工作[29]，我们提出创造多尺度稀疏张量和在不同分辨率下具有固定半径的局部采样位移。较粗的分辨率提供更大的背景，而更细的分辨率提供更准确的位移。然后将每层的稀疏张量转换为密集张量，并将它们连接成一个单一的二维张量。这一点如图4b所示。

在每次迭代i中，对于每个像素x，我们从一组top-k个相关位置$S_{\mathbf{x}}^{(k)}$开始。因此，集合$\left\{(\mathbf{d},C(\mathbf{x},\mathbf{d}))\mid \mathbf{d}\in S_{\mathbf{x}}^{(k)}\right\}$记录使用kNN算法获得的像素x及其位置的前k个相关值。我们将坐标除以（1,2,4,8,16），构造一个五级稀疏相关体积金字塔，并表示$l$ 级的比例位移，用$d^l=(d_i-∆f_i)/2^{l-1}$更新当前的$∆f_i$。

此外，我们用$C^l\left(\mathbf{x},{\mathbf{d}}^l\right)=C(\mathbf{x},\mathbf{d}),d\in S_{\mathbf{x}}^{(k)}$来表示l层处的相关值。在每一层，我们用一个常数半径r 约束位移，并在l层定义有窗口的相关值集，$$\left\{\left({\mathbf{d}}^l,C^l\left(\mathbf{x},{\mathbf{d}}^l\right)\right)\mid {\Vert {\mathbf{d}}^l\Vert }_{\infty }\le r,\mathbf{d}\in S_{\mathbf{x}}^{(k)}\right\}\text{\hspace{1em}(5)}$$由于坐标$d^l$ 不一定是整数，我们需要重新采样到整数坐标，以增强相关性的稀疏张量。我们提出了一种我们称之为“双线性溅射(bilinear splatting)”的方法。相关值被双向分布到四个最近的整数网格上。例如，在位置$d^l$ 处的相关量$C^l(x,d^l)$被传播到四个相邻的整数点中的每一个，用$[d^l]=(d_x,d_y)$表示，按照$$C^l\left(\mathbf{x},\left[{\mathbf{d}}^l\right]\right)=\left(1-\left|d_x^l-d_x\right|\right)\left(1-\left|d_y^l-d_y\right|\right)C^l\left(\mathbf{x},{\mathbf{d}}^l\right)$$然后将这些值相加为相关性集合（5），并将每个层的稀疏张量转换为密集张量，重塑并连接形成一个5维$(2r+1{)}^2$的2维密集张量，其中5是金字塔层的数量。

我们在这里介绍的方法不需要学习。它仅仅是稀疏张量和密集张量之间的转换，因此比稀疏卷积更简单。

2.5. GRU Update Block

这个二维密集张量中的每个向量编码位置信息以及k匹配的相关值。我们将二维运动张量与上下文特征和当前的流量估计连接起来，并将其通过一个门控循环单元(GRU)更新块。GRU更新块估计剩余流$∆f_{i+1}$，用于在下一步移动相关量坐标。

3. Experiments

3.1. Implementation details

Network details 我们首先提取了具有256个通道的四分之一分辨率的特征图。我们的特征提取网络包含6个残差块。当将特征映射传递给kNN时，我们设置了k=8。即，对于每个特征向量，返回给出最大内积的前8个特征向量的索引。GRU更新块以当前的光流估计值以及上下文特征图作为输入。上下文特征图由一个具有128个通道的单独网络提取。GRU更新块还更新了一个128维的隐藏状态特征向量。在训练期间，GRU迭代8次，而在RAFT[29]中迭代12次。

Training schedule 在之前的工作之后，我们首先在[9]的模型上进行批量大小为6的120k迭代，然后在[21]上进行另一个批量大小为4的120k迭代。然后，我们对Sintel[6]、KITTI2015[23]和HD1K[19]进行120k迭代进行Sintel评估，对KITTI2015[23]进行50k迭代进行KITTI评估。我们使用大小为4的批处理大小来进行微调。我们在两个2080Tigpu上训练我们的模型。烧蚀实验是在一个特斯拉P100GPU上进行的。我们使用PyTorch库[24]实现了它。

Loss function 与RAFT[29]类似，我们采用了一个循环网络架构，其中预测了一系列残余流∆fi。每个步骤的光流预测可以表示为$f_{i+1}=f_i+∆f_{i+1}$，初始值为$f_0=0,∆f_0=0$. 我们将损耗函数应用于光流预测的序列。给定真实光流$f_{gt}$和每步$f_i$的预测光流，将损失函数定义为$$L=\sum _{i=1}^N{\gamma }^{N-i}{\Vert {\mathbf{f}}_i-{\mathbf{f}}_{\mathrm{g}\mathrm{t}}\Vert }_1$$

kNN 我们使用faiss库[17]在gpu上运行kNN。目前，我们正在应用蛮力精确搜索方法，因为我们的问题规模仍然被认为很小。faiss库提供了优化的k-选择例程来加快计算速度，有关更多细节，我们请读者参阅原始文章[17]。

3.2. Results

我们在表1中展示了与现有工作的定量比较。在双视图的情况下，我们在Sintel干净数据集上取得了最先进的结果，比RAFT[29]获得了11.3%的改进（1.94→1.72）