点跟踪论文—RAFT: Recurrent All-Pairs Field Transforms for Optical Flow-递归的全对场光流变换

点目标跟踪论文—RAFT: Recurrent All-Pairs Field Transforms for Optical Flow-递归的全对场光流变换

读论文RAFT密集光流跟踪的笔记

RAFT是一种新的光流深度网络结构，由于需要基于点去做目标的跟踪，因此也是阅读了像素级别跟踪的一篇ECCV 2020的经典论文 ——RAFT,递归的全对场光流变换,使用密集的光流来对小而快速的物体进行跟踪。

• 作者：Zachary Teed and Jia Deng
• 发表：ECCV 2020 best paper

论文的难点与核心点

1. 不像之前的coarse-to-fine类的方法，RAFT在计算时始终保持同一分辨率，而coarse-to-fine则是对多尺度预测，逐步细化的方式
2. update operator是轻量的和循环的，而其他的算法则只能是循环几次，无法长时间循环。
3. 一个新的update operator，由卷积GRU组成，可look up 生成的4D相关信息。

是论文的难点也是自己理解不太好的地方。

背景信息

密集光流估计为每个像素分配一个二维光流向量，描述其在时间间隔内的水平和垂直位移。在稀疏光流中，该向量只分配给与边角等强特征相对应的像素。

水平梯度 Iₓ 和垂直梯度 Iᵧ 可用索贝尔算子近似，时间梯度 Iₜ 已知，因为我们有 t 和 t+1 时间的图像。方程有两个未知数 u 和 v，分别是时间 dt 上的水平位移和垂直位移。单个方程中的两个未知数使其成为一个未决问题，人们曾多次尝试求解 u 和 v。RAFT 是一种估算 u 和 v 的深度学习方法，但它实际上比根据两个框架预测流量更复杂。它是为精确估计光流场而精心设计的。

取得的成果：

在KITTI上，RAFT取得了5.10%的F1-all误差，比已公布的最佳结果 (6.10%)降低了16%。在Sintel（final pass）上，RAFT获得了2.855 EPE误差，比已发布的最佳结果（4.098像素）减少了30%的误差。此外，RAFT具有很强的跨数据集泛化能力，以及在推理时间、训练速度和参数量方面具有很高效率。

a new end-to-end trainable model for optical flow

摘要与整体概括

摘要

摘要的核心总结：

1. RAFT 逐像素提取特征，为所有像素对构建多尺度4D相关体，并通过循环单元在相关体上进行查找，以迭代更新光流场。

2. 是一种新的光流深度网络架构。

在学习完成论文之后总结来说：其中的两个关键的词包含了光流跟踪最重要的两个过程信息。

• correlation: 是我们计算像素之间的全相关性和进行保持高分辨率不变的基础上进行多尺度金字塔构建的一个核心。

• lookup：是作者们为了简化一定的计算和损失，所提出的一种在coor上寻找特征点的一种方法。（难理解要结合看代码）。

整体概括

给定一对连续的RGB图像作为输入，估计稠密位移场（dense displacement field），位移场将中的每个像素映射到中的对应坐标。算法框架如下图所示：

我们讲述这篇论文从整体到细节首先要明确的是，学习的重点是整篇论文实现的一个思想。我们给出总结的整体的步骤和框架，结合整体的结构不断的细化。论文中给出了我们一定的总结。

REFT主要包括以下的三个部分组成。

1. 特征提取编码器与内容提取编码器。
2. 一个相关层得出4d的一个相关体。
3. 一个用来更新的操作符从相关体中更新光流的信息。

我们对这个过程进行总结在详细的进行描述。

• 先是一个网络 Net1（特征提取编码器）提取两张输入 I1,I2的特征，还有另一个网络 Net2(内容提取编码器) 再提取一次 I1 的特征，然后通过一个correlation layer接收 Net1的输出并建立两张图片的相似度向量矩阵。最后作者使用了自然语言处理中GRU的思想，把相似度向量，每一次迭代预测出的光流，以及 Net2的输出三者作为输入去迭代着更新光流。

• RAFT由三部分组成：

• (1)一个feature encoder提取两张输入图片 I1,I2在每个像素点上的特征。这里我们假设 I1,I2的尺寸是 H×W ，那么经过feature encoder之后得到的特征维度就是 H×W×D ；此外还有一个 context encoder提取 I1的特征，也就是图片的左下角。

• (2)一个 correlation layer负责把 I1,I2 的特征向量通过点乘的方式连接起来，那么最终输出的是一个 H×W×H×W的4d向量，此向量表示 I1每一个像素点与所有 I2像素点的相关度。然后作者也考虑到这样的表示可能比较稀疏，因此在这个输出之后做了四层的池化，并将每一层池化的输出连接起来做成了一个具有多尺度特征的相似性变量。

• (3)一个update operator，通过使用一个look up方法（查看 4D Correlation Voulumes的值）迭代着去更新光流。当然第三点需要下面的详细介绍。

The feature encoder extracts per-pixel features. The correlation layer
computes visual similarity between pixels. The update operator mimics the steps of an iterative optimization algorithm.

因为我们已经进行了整体的概括了，有了一定的印象，我们按照论文的结构展开

Approach

方法部分是官方的一个整体的介绍。

给一组连续的RGB图像 I1 I2 we estimate a dense displacement
field (f1, f2)去估计一个密集的位移场 f1,f2.每一个光流（u,v）要map I2.同时总结了三个核心的部分

• 特征提取
• 相似度计算
• 迭代更新

$$\left(u^{\prime },v^{\prime }\right)=\left(u+f^1(u),v+f^2(v)\right)$$

分辨率描述

在建立完成相关量之后也就是完成8倍的下采样操作通过backebone网络提取出来特征图之后。作者给出了一副描述图

对于I1中的特征向量，我们取其与I2中所有对的内积，生成一个4D W×H×W×H体 (I2中的每个像素生成一个2D 响应图)。使用卷积核大小为1、2、4、8的平均池化对相关体池化。

这组相关性张量同时包含了大位移和小位移的信息；但是也保持前两个维度(维度)，因此保存了高分辨率的信息，从而可以恢复小的快速移动物体的运动。

这里自己直观的解读一下的话其实就是，每一个点我们代表的是一个像素的分辨率，我们只对最后的两个维度进行下采样的操作。导致每幅图的最后的两个维度发生变换，内部的每一组包含的像素点发生变换，但我们外部的尺寸保持不变使得内部像素点的分辨率其实也是保持不变的会一直维持在8倍的分辨率上。

保持高分辨率从而就克服了coarse-to-fine类的方法—有粗到细使用不同的分辨率进行采样。

Feature Extraction

核心总结：

1. 使用卷积网络进行特征的提取
2. 最终通过2倍4倍完成8倍的下采样操作 D=256
3. 由6个残差块组成。整个特征提取只进行一次。
4. 唯一的区别是特征编码器使用实例标准化，而上下文编码器使用批量标准化。

backbone主干网络细节

相比于其他的复杂的神经网络来说，这里的backbone部分相对比较简单，很大的程度上参考了RestNet50这种结构。

layer层使用的和YOLO中常用的botteneck模块类似，这里我们称为残差块，在代码中每两个残差块为一组。1不进行下采样 2 3在第一个卷积的部分进行下采样的操作。（下采样连接）

对于两幅图 I1和 I2都需要提取特征，该网络称之为 Feature Encoder

Computing Visual Similarity

Computing Visual Similarity计算视觉相似度

通过在所有输入图像对之间构造一个correlation volume(下称为相关性张量)来计算视觉相似性 -印象中代码里面最终对应的是coor

给定抽取得到图像特征： 通过对所有的特征向量对进行点积得到 相关性张量。记相关性张量为 C

$$g_{\theta }\left(I_1\right)\in {\mathbb{R}}^{H\times W\times D}\text{ and }{g}_{\theta }\left(I_2\right)\in {\mathbb{R}}^{H\times W\times D}$$

$$\mathbf{C}\left(g_{\theta }\left(I_1\right),g_{\theta }\left(I_2\right)\right)\in {\mathbb{R}}^{H\times W\times H\times W},C_{ijkl}=\sum _h{g}_{\theta }{\left(I_1\right)}_{ijh}\cdot g_{\theta }{\left(I_2\right)}_{klh}$$

Correlation Layer模块

这里我们得到了 I1对 I2上的多尺度4D Correlation Voulumes.

我们这个模块单独的分离出来，方便进行简单的说明。

我们从最终的一个公式中其实也是可以看出来的，4d向量空间最终融合的是通道数h嘛。

如果我们不看下面的代码来看这一个过程信息。

@staticmethod
    def corr(fmap1, fmap2):
        batch, dim, ht, wd = fmap1.shape
        fmap1 = fmap1.view(batch, dim, ht*wd) # 展平一个维度
        fmap2 = fmap2.view(batch, dim, ht*wd) 

        corr = torch.matmul(fmap1.transpose(1,2), fmap2) #转置相乘
        corr = corr.view(batch, ht, wd, 1, ht, wd)
        return corr  / torch.sqrt(torch.tensor(dim).float()) # coor 互相关运算得到的矩阵

代码是将h和w展平之后来进行操作的。—转置相乘。最后在进行展开合并维度相关的信息。

我们如果按照3通道图片自己来想这一个过程的话，其实就是三个通道的图像相同位置的像素同时做一个点积的操作。在依次的相加得到一个值

我们每一个像素值要做i2的大小 H x W然后i1中有 H x W个像素需要做相关性运算从而就是这个结果了。

从这两个图也可以看出下面会有一个保持分辨率的情况下进行构建金字塔的过程。

Correlation pyramid模块

我们得到 H×W 的向量之后，作者觉得这样比较稀疏，因为 I1不可能与 I2所有的像素点相关，所以作者又将这个向量进行了四层池化。

这样一个相关信息张量C​是非常大的, 因此在C的最后两个维度上进行汇合来降低维度大小，每个​的维度为保持前两个维度不变, 这种相关信息张量可以保证同时捕捉到较大和较小的像素位移。

这组相关性张量同时包含了大位移和小位移的信息；但是也保持前两个维度(维度)，因此保存了高分辨率的信息，从而可以恢复小的快速移动物体的运动。

这个叠加在代码中是放在一个out_pyramid = []变量中来进行返回实现的。

 self.corr_pyramid.append(corr) # 后两个维度下采样构建金字塔
 self.corr_pyramid.append(corr) # 后两个维度下采样构建金字塔
        for i in range(self.num_levels-1):
            corr = F.avg_pool2d(corr, 2, stride=2) # 2倍平均池化
            self.corr_pyramid.append(corr)

Look up模块

Look up的部分个人感觉是最难理解的一个部分（原因在于它描述的过程比较抽象。）具体的更多细节的实现还是参考代码中的具体过程。

We define a lookup operator LC which generates a feature map by indexing from the correlation pyramid.—我们定义了一个查找算子 LC，它通过从相关金字塔进行索引来生成特征图。

上一步构建了四层的Correlation Pyramid，这里要根据像素去查找这个Correlation Pyramid中的对应特征。如果对I1中的每个点的向量都要去I2中所有向量找对应点的话，需要的cost太大了，所以论文中设置了一个lookup的参数，即只对该位置附近位置的点做判断

$$\mathcal{N}{\left({\mathbf{x}}^{\prime }\right)}_r=\left\{{\mathbf{x}}^{\prime }+\mathbf{dx}\mid \mathbf{dx}\in {\mathbb{Z}}^2,\Vert \mathbf{dx}{\Vert }_1\le r\right\}$$

r超参数是超参数，有点类似于圆的半径，dx是整数，通过这个公式把x’附近的值拿到，同时这个操作会在每一层的金字塔上取值，最后将这些得到的值串联成一个向量。这个向量也就是Lookup的输出。总结一下就是光流建立了I的像素点到I2像素点的映射，然后使用对应的I2点的坐标，在对应的相似性向量的金字塔上采样得到一个输出向量。那么大胆猜测一下，对于快速移动的物体，r设置的偏大一些，效果应该更好；对于移动较慢的无题，r设置的应该偏小一些。

我自己在粗略的看代码的时候也是感觉这个Look up模型更像是一个采样的模块。是在相似性向量的金字塔通过这种邻近范围的方式查找符合计算条件的点（输入后面GRU模块进行计算，应该也是因为维度过高导致计算量过大的原因吧。）

代码中是通过grid_sample（）函数之间完成这一个过程的。这个r类似半径但实际上是一个线性的间隔值

bilinear_sampler
F.grid_sample（）

        for i in range(self.num_levels):
            corr = self.corr_pyramid[i]
            dx = torch.linspace(-r, r, 2*r+1, device=coords.device) # (2r+1) x方向的相对位置查找范围 -r,-r+1,...,r 从-r到r的线性间隔的值，并且包括端点。2*r+1是生成这些值的数量
            dy = torch.linspace(-r, r, 2*r+1, device=coords.device) # # (2r+1) y方向的相对位置查找范围
            delta = torch.stack(torch.meshgrid(dy, dx), axis=-1) # 查找窗 (2r+1,2r+1,2) 一维张量dy和dx的笛卡尔积torch.meshgrid返回的两个二维张量堆叠起来，形成一个三维张量。axis=-1参数指定了堆叠的轴，这里是最后一个轴(dy_i, dx_i)对表示二维空间中的一个位移向量

            centroid_lvl = coords.reshape(batch*h1*w1, 1, 1, 2) / 2**i # 某尺度下的坐标
            delta_lvl = delta.view(1, 2*r+1, 2*r+1, 2)
            coords_lvl = centroid_lvl + delta_lvl # 可以形象理解为：对于 bhw 这么多待查找的点，每一个点需要搜索 (2r+1)*(2r+1) 邻域范围内的其他点，每个点包含 x 和 y 两个坐标值

            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()

这里我们返回的corr变量其实也就是经过采样之后的一个变量值。作为下一部分的一个输出。

 corr = corr_fn(coords1) # 核心index correlation volume 从相关性查找表中获取当前坐标的对应特征

Efficient Computation for High Resolution Images（可选）

具体的公式含义细节我自己也不太明白不做过多的解读了。

原因在于这一个部分，对应代码中的CUDA变成部分对应C++模块是以是一个在训练中可以选择的计算方式，自己不太会CUDA的这一个部分

我自己读论文的话，其实他的思想就是将下采样的操作融合进相关性的计算里面以减少参数量来实现简化计算。从O(n2)到O（mn）

$${\mathbf{C}}_{ijkl}^m=\frac{1}{2^{2m}}\sum _p^{2^m}\sum _q^{2^m}⟨g_{i,j}^{(1)},g_{2^{m}k+p,2^{m}l+q}^{(2)}⟩=⟨g_{i,j}^{(1)},\frac{1}{2^{2m}}\left(\sum _p^{2^m}\sum _q^{2^m}{g}_{2^{m}k+p,2^{m}l+q}^{(2)}\right)⟩$$

Iterative Updates

Our update operator estimates a sequence of flow estimates {f1, …, fN} from an initial starting point f0 = 0

看出初始化的光流为0

 # 初始化光流的坐标信息，coords0 为初始时刻的坐标，coords1 为当前迭代的坐标，此处两坐标数值相等
        coords0, coords1 = self.initialize_flow(image1)
         flow = coords1 - coords0 # 初始值为0

输入：当前光流，以及从金字塔中提取的对应的相关特征，context。所以输入是相关特征，光流以及上下文特征。

 with autocast(enabled=self.args.mixed_precision):
                net, up_mask, delta_flow = self.update_block(net, inp, corr, flow)

            # F(t+1) = F(t) + \Delta(t)
            coords1 = coords1 + delta_flow

带有卷积的GRU模块（update block模块）

更新算子的核心组件是一个基于GRU的gated activation unit，其中的全连接层使用卷积替换：

其中 x t 是前面定义的光流、相关特征、context特征的拼接。论文还实验了一个可分离的ConvGRU单元，其中用两个GRU替换3×3卷积: 一个用1×5卷积，一个用5×1卷积，以便在不显著增加模型大小的情况下增加感受野。

整体结果的一个定义。

self.args = args
        self.encoder = BasicMotionEncoder(args)
        self.gru = SepConvGRU(hidden_dim=hidden_dim, input_dim=128+hidden_dim)
        self.flow_head = FlowHead(hidden_dim, hidden_dim=256)

对应更为细节的具体描述需要断点调试前向传播的部分在文章中就不过多的进行展开说明了。

 def __init__(self, hidden_dim=128, input_dim=192+128):
        super(SepConvGRU, self).__init__()
        self.convz1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
        self.convr1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))
        self.convq1 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (1,5), padding=(0,2))

        self.convz2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
        self.convr2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))
        self.convq2 = nn.Conv2d(hidden_dim+input_dim, hidden_dim, (5,1), padding=(2,0))


    def forward(self, h, x):
        # horizontal
        hx = torch.cat([h, x], dim=1)
        z = torch.sigmoid(self.convz1(hx))
        r = torch.sigmoid(self.convr1(hx))
        q = torch.tanh(self.convq1(torch.cat([r*h, x], dim=1)))        
        h = (1-z) * h + z * q

其中的全连接层使用卷积替换

光流预测

self.flow_head = FlowHead(hidden_dim, hidden_dim=256)
光流头通过两个全连接层进行输出。

将GRU输出的隐藏状态经过两个卷积层来预测光流的更新 Δf 。输出的光流的分辨率是输入图像的1/8。

​ 在训练和评估过程中，对预测的光流场进行上采样，以匹配ground-truth的分辨率。

class FlowHead(nn.Module):
    def __init__(self, input_dim=128, hidden_dim=256):
        super(FlowHead, self).__init__()
        self.conv1 = nn.Conv2d(input_dim, hidden_dim, 3, padding=1)
        self.conv2 = nn.Conv2d(hidden_dim, 2, 3, padding=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        return self.conv2(self.relu(self.conv1(x)))

上采样

RAFT 的作者对双线性和凸上采样进行了实验，发现凸上采样可以显着提高性能。

凸上采样将每个精细像素估计为其相邻 3x3 粗像素网格的凸组合

$$\mathcal{L}=\sum _{i=1}^N{\gamma }^{N-i}{\Vert {\mathbf{f}}_{gt}-{\mathbf{f}}_i\Vert }_1$$