一种基于Transform的神经网络架构

摘要

大家好，今天为大家带来的文章 A Transformer Architecture for Optical Flow 我们介绍了光流估计网络，称为FlowFormer，一种基于Transform的神经网络架构，用于学习光流。

FlowFormer化由图像对构建的4D cost volume，将成本编码到一个新的潜在空间中具有交替组转换器(AGT)层的成本内存中，并通过一个带有动态位置成本查询的循环Transform解码器对位置 cost queries进行解码。

在sinintel基准测试中，FlowFormer的平均终点误差(AEPE)分别为1.159和2.088，比已发布的最佳结果误差分别降低了16.5%和15.5%

主要工作与贡献

我们的贡献可以概括为四个方面。

1)我们提出了一种新的基于Transform的神经网络结构FlowFormer，用于光流量估计，它实现了最先进的流量估计性能。

2)设计了一种新颖的cost volume编码器，有效地将成本信息聚合为紧凑的潜在cost tokens。

3)我们提出了一种循环成本解码器，该解码器通过动态位置成本查询循环解码成本特征，迭代细化估计光流。

4)据我们所知，我们第一次验证imagenet预先训练的传输

方法

光流估计任务要求输出逐像素位移场f: R2→R2，将源图像的每个二维位置x∈R2 Is映射到目标图像It对应的二维位置p = x+f (x)。为了充分利用现有的视觉Transform体系结构，以及之前基于cnn的光流估计方法广泛使用的四维成本体积，我们提出了一种基于Transform的结构FlowFormer，它对四维成本体积进行编码和解码，以实现精确的光流估计。

在图1中，我们展示了FlowFormer的总体架构，它处理来自siamese特征的4D成本卷，包含两个主要组件: 1)cost volume编码器，将4D成本卷编码到潜在空间中形成cost memory; 2)成本记忆解码器，用于基于编码的成本记忆和上下文特征预测每像素位移场。



图1 FlowFormer的体系结构。

FlowFormer通过三个步骤估算光流量:1)根据图像特征构建4Dcost volume。2)成本卷编码器，将成本卷编码到成本内存中。3)循环Transform解码器，将具有源图像上下文特征的代价内存解码为流

1.1 构建 4D Cost Volume

骨干网用于从输入HI × WI × 3rgb图像中提取H ×W ×Df特征图，通常我们设置(H, W) = (HI /8, WI /8)。在提取出源图像和目标图像的特征图后，通过计算源图像和目标图像之间所有像素对的点积相似度，构造H × W × H × W 4D代价体。

1.2 Cost Volume Encoder 为了估计光流，需要基于4D代价体中编码的源目标视觉相似性来识别源像素在目标图像中的对应位置。构建的4D成本体积可以被视为一系列大小为H × W的2D成本图，每个成本图度量单个源像素和所有目标像素之间的视觉相似性。

我们将源像素x的代价映射表示为Mx∈RH×W。

在这样的成本图中找到对应的位置通常是具有挑战性的，因为在两个图像中可能存在重复的模式和非歧视性区域。当只考虑来自地图局部窗口的成本时，任务变得更加具有挑战性，就像以前基于cnn的光流估计方法所做的那样。即使是估计单个源像素的精确位移，考虑其上下文源像素的代价图也是有益的。

为了解决这一难题，我们提出了一种基于Transform的成本体积编码器，该编码器将整个成本体积编码到成本内存中。我们的成本卷编码器包括三个步骤:1)成本映射补丁化，2)成本补丁嵌入，3)成本内存编码。我们详细阐述这三个步骤如下。

1.2 地图patchification成本

我们根据已有的视觉转换器，对每个源像素x的代价映射Mx∈RH×W进行跨卷积拼接，得到代价补丁嵌入序列。具体来说，给定一个H ×Wcost图，我们首先在它的右侧和底部填充0，使其宽度和高度为8的倍数。然后，填充的代价映射通过三个stride-2卷积的堆栈和ReLU转换为特征映射Fx∈R⌈H/8⌉×⌈W/8⌉×Dp。特征图中的每个特征代表输入成本图中的一个8 × 8补丁。三种卷积的输出通道均为Dp / 4, Dp / 2 Dp。

1.3基于潜在摘要的补丁特征标记

尽管对每个源像素进行修补会得到一系列cost patch特征向量，但这种patch特征的数量仍然很大，影响了信息在不同源像素间传播的效率。实际上，成本图是高度冗余的，因为只有少数高成本是最有信息的。为了获得更紧凑的代价特征，我们进一步通过K个潜码字C∈RK×D总结了每个源像素x的patch特征Fx。

具体来说，latent codewords query每个源像素的cost补丁特征，通过点积注意机制将每个成本图进一步总结为K个D维的潜在向量。潜码字C∈RK×D随机初始化，通过反向传播进行更新，并在所有源像素之间共享。归纳Fx的潜在表示Tx得到为



在投影成本-补丁特征Fx以获得键Kx和值Vx之前，将补丁特征与位置嵌入序列PE∈R⌈H/8⌉×⌈W/8⌉×Dp进行拼接。给定一个2D位置p，我们将其编码为一个长度为Dp的位置嵌入，跟随COTR[27]。最后，通过对查询、键和值进行多头点积注意，可以将源像素x的代价图汇总为K个潜在表示Tx∈RK×D。

关注潜在cost空间。上述两个阶段将原始的4D成本体积转化为潜在的、紧凑的4D成本体积t。但是，直接对4D体积中的所有向量应用自注意仍然过于昂贵，因为计算成本随tokens数量的增加呈二次增长。 如图2所示，我们提出了一种交替分组转换层(AGT)，该层以两种相互正交的方式对标记进行分组，并在两组中交替应用注意，减少了注意成本，同时仍然能够在所有标记之间传播信息。



图2 Alternate-Group Transform层。

交替组转换层(AGT)将T中的token交替分组为包含K token的H × W组(Tx)和包含H × W token的K组(Ti)，并分别通过self-attention和ss self-attention[8]对组内的token进行编码。 对每个源像素进行第一次分组，即每个Tx∈RK×D组成一个组，在每个组内进行自我注意。



在每个源像素x的所有K个潜在令牌之间进行自注意后，更新的Tx通过前馈网络(FFN)进一步变换，然后重新组织，形成更新的4D成本体积t。自注意子层和FFN子层均采用Transform残差连接和层归一化的共同设计。这种自我注意操作在每个成本映射内传播信息，我们将其命名为内部成本映射的自我注意。

第二种方法根据K种不同的潜在表示，将所有潜在代价令牌T∈RH×W ×K×D分成K组。因此，每一组都有D维度的(H × W)标记，用于通过双胞胎[8]中提出的空间可分离的自我注意(SS-SelfAttention)在空间域中传播信息，

上述自我注意操作的参数在不同组之间共享，并按顺序进行操作，形成提出的交替组注意层。通过多次叠加交替组Transform层，潜在成本令牌可以有效地跨源像素和跨潜在表示交换信息，以更好地编码4D成本体积。通过这种方式，我们的成本体积编码器将H × W × H × W 4D成本体积转换为H × W × K长度为d的潜在标记。我们将最终的H × W × K标记称为成本存储器，用于光学解码。

1.4用于流量估计的成本记忆解码器

考虑到成本体积编码器编码的cost记忆，我们提出了一种用于光流预测的cost记忆解码器。由于输入图像的原始分辨率为HI × WI，我们在H × W分辨率下估计光流，然后使用可学习的凸上采样器[46]对预测的流进行上采样到原始分辨率。然而，与以往的视觉Transform寻找抽象的语义特征不同，光流估计需要从记忆中恢复密集的对应。

受RAFT[46]的启发，我们建议使用成本查询从cost内存中检索成本特征，并使用循环注意解码器层迭代优化流量预测。

1.5Cost memory aggregation:

为了预测H × W源像素的流，我们生成了一个(H × W)成本查询序列，每个(H × W)成本查询通过对成本内存的共同关注负责估计单个源像素的流。为了为源像素x生成成本查询Qx，我们首先计算其在目标图像中的对应位置，给定其当前估计流量f (x)，即p = x + f (x)。

然后，通过裁剪成本映射Mx上以p为中心的9×9本地窗口内的成本，我们检索一个局部9×9成本映射补丁qx = Crop9×9(Mx, p)。然后根据局部成本Qx编码的特征FFN(Qx)和p的位置嵌入PE(p)构造成本查询Qx，通过交叉注意聚合源像素x的成本记忆Tx信息。

循环光流预测。我们的成本解码器迭代回归流量残差∆f (x)，将每个源像素x的流量细化为f (x)←f (x) +∆f (x)。我们采用ConvGRU模块，并遵循与GMARAFT[25]类似的设计进行流量细化。

然而，我们的循环模块的关键区别是使用成本查询来自适应地聚合来自成本内存的信息，以实现更精确的流量估计。具体来说，在每次迭代中，ConvGRU单元将检索到的成本特征和成本映射补丁Concat(cx, qx)、来自上下文网络的源图像上下文特征tx和当前估计的流量f的拼接作为输入，并输出预测的流量残差如下

实验结果

 

图4 Sintel[3]和KITTI[14]数据集的实验。

' A '表示自动流数据集。“C + T”表示只在FlyingChairs和FlyingThings数据集上进行训练。' + S + K + H '表示对Sintel、KITTI和HD1K训练集的组合进行微调。*表示这些方法使用warm-start策略[46]，该策略依赖于视频中先前的图像帧。†是通过补充阐述的瓦技术估计的。我们的FlowFormer实现了最佳的泛化性能(C+T)，并在Sintel基准测试(C+T+S+K+H)中排名第一。



图5 在steint检验集上的定性比较。

FlowFormer大大减少了对象边界(红色箭头所指)周围的流动泄漏，并使细节(蓝色箭头所指)更加清晰。


图5 在steint检验集上的定性比较。

FlowFormer大大减少了对象边界(红色箭头所指)周围的流动泄漏，并使细节(蓝色箭头所指)更加清晰。




审核编辑：刘清