【光流】——RAFT: 稠密的光流场估计(代码浅析)

转载自:RAFT: 稠密的光流场估计
代码：https://github.com/princeton-vl/RAFT
尊重原创，请读原文

一、摘要

本文提出了Recurrent All-Pairs Field Transforms(RAFT), 一个光流估计的深度神经网络. RAFT 提取像素级的特征, 为所有像素建立多尺度 4D 关联信息, 通过查找4D关联信息, 循环迭代的更新光流场. 本文算法在KITTI、Sintel数据集上取得了state-of-the-art的表现. 同时, RAFT在多个数据集上有很强的泛化能力, 并且在训练速度、参数数量、推理时间上都有很高的效率.

二、网络结构

RAFT网络由三部分组成:

(1) 特征编码层: 逐像素的提取特征;

(2) 特征关联层: 为所有像素产生WxHxWxH的4D关联信息。视觉相似性(Visual Similarity):视觉相似度计算为所有特征图对的内积。因此，我们将得到一个称为相关体积的四维张量，它提供了关于大小像素位移的关键信息。

(3) 循环更新算子: 基于GRU的循环更新算子查找4D关联信息, 然后迭代的更新光流场, 光流场初始值为0;

三、特征采集模块

对2张图片采用同一个网络进行特征提取。

image1 = 2 * (image1 / 255.0) - 1.0
image2 = 2 * (image2 / 255.0) - 1.0

image1 = image1.contiguous()
image2 = image2.contiguous()

hdim = self.hidden_dim
cdim = self.context_dim

# run the feature network
with autocast(enabled=self.args.mixed_precision):
    fmap1, fmap2 = self.fnet([image1, image2])   

4. 视觉相似模块

上一步得到两幅图像的特征​张量 $g$ ， [公式] ​, correlation volume ​ [公式] 的计算方法:

这样就得到了image1和image2的关系特征矩阵。同时这里采用了类似金字塔的思想。这样一个相关信息张量$C$​是非常大的, 因此在$C$的最后两个维度上采用maxpooling进行降维, 产生4个相关信息​$\left\{C^1,C^2,C^3,C^4,\right\}$，每个​的维度为$H\ast W\ast H/2^k\ast W/2^k$ ​. 保持前两个维度不变, 这种相关信息张量可以保证同时捕捉到较大和较小的像素位移.

不同尺度的关系矩阵

self.corr_pyramid.append(corr)
for i in range(self.num_levels-1):
    corr = F.avg_pool2d(corr, 2, stride=2)
    self.corr_pyramid.append(corr)

5. 迭代更新

迭代更新类似是一个门控循环单元(GRU)序列，但是因为图像特征是2维的，所以这里是分别在水平和竖直两个方向做的GRU
公式

代码：

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

        # vertical
        hx = torch.cat([h, x], dim=1)
        z = torch.sigmoid(self.convz2(hx))
        r = torch.sigmoid(self.convr2(hx))
        q = torch.tanh(self.convq2(torch.cat([r*h, x], dim=1)))       
        h = (1-z) * h + z * q

        return h

6. 上采样模块

转载自：https://blog.csdn.net/weixin_43229348/article/details/120500861
GRU单元输出的光流分辨率为初始图像的1/8，因此作者提出了两种不同的上采样方法来匹配真值分辨率。第一种是光流的双线性插值。它是一种简单快速的方法，但这种方法的质量不如一个可学习名叫Convex Upsampling的上采样模块:

凸上采样(Convex Upsampling)方法表明，全分辨率光流是GRU单元预测的3x3加权网格的凸组合。8 倍图像上采样意味着必须将 1 个像素扩展为 64(8x8) 个像素。凸上采样模块由两个卷积层和末端的softmax激活来预测上采样光流预测中每个新像素的H / 8 × W / 8 × ( 8 × 8 × 9 ) H/8 掩码。现在，上采样图像上的每个像素都是之前粗糙分辨率像素的凸组合，由预测掩码加权，系数为w1 , w2, w3, … w9

code

def upsample_flow(self, flow, mask):
        """ Upsample flow field [H/8, W/8, 2] -> [H, W, 2] using convex combination """
        N, _, H, W = flow.shape
        mask = mask.view(N, 1, 9, 8, 8, H, W)
        mask = torch.softmax(mask, dim=2)

        up_flow = F.unfold(8 * flow, [3,3], padding=1)
        up_flow = up_flow.view(N, 2, 9, 1, 1, H, W)

        up_flow = torch.sum(mask * up_flow, dim=2)
        up_flow = up_flow.permute(0, 1, 4, 2, 5, 3)
        return up_flow.reshape(N, 2, 8*H, 8*W)

7. 整体流程

• 提取image1，image2的特征
  fmap1, fmap2 = self.fnet([image1, image2])
• 对fmap1，fmap2计算4D相关性矩阵$C$
  corr_fn = CorrBlock(fmap1, fmap2, radius=self.args.corr_radius)
• 对image1在提取特征，并在通道上上分解为：GRU的初始状态与feature特征

with autocast(enabled=self.args.mixed_precision):
      cnet = self.cnet(image1)
      net, inp = torch.split(cnet, [hdim, cdim], dim=1)
      net = torch.tanh(net)
      inp = torch.relu(inp)

• 对image1的特征（H/8 * W/8）中的每个点，从4D相关性矩阵$C$中采样81个点。4层C，共有4 * 81=324个关联特征，
• 关联特征和初始的flow，concat之后经过几个conv后，得到motion feature。
• motion feature在和image1的图像特征concat之后送入二维GRU模块

motion_features = self.encoder(flow, corr)
inp = torch.cat([inp, motion_features], dim=1)

• 用GRU的隐藏状态+conv做输出

net = self.gru(net, inp)
delta_flow = self.flow_head(net)

• 用GRU的隐藏状态+conv输出Convex Upsampling上采样的权重
• 最后，对GRU进行一定$ITER$的迭代。

总结

总体，创新点：

• 引入了motion feature，而motion feature的计算通过金字塔4D关系矩阵均匀采样得来。
• 引入了GRU概念进行迭代优化

缺点

• 推理流程多，存在迭代过程

reference