使用 TensorFlow 进行光流预测

本仓库提供了基于 TensorFlow 的实现，复现了 Deqing Sun 等人发表于 CVPR 2018 的优秀论文《PWC-Net：利用金字塔、扭曲和代价体积的 CNN 光流网络》。

目前已有若干尝试使用 TensorFlow 实现 PWC-Net。然而，这些实现要么沿用了论文中已过时的 CNN 网络架构，要么仅提供 TF 推理（无 TF 训练），要么仅支持 Linux 平台，且不支持多 GPU 训练。

本实现同时提供了基于 TensorFlow 的训练与推理功能。它具有良好的可移植性：由于未使用任何动态加载的 CUDA 版 TensorFlow 用户自定义操作，因此可在 Linux 和 Windows 上运行。此外，它还支持多 GPU 训练（此处展示的笔记本及结果均在 GTX 1080 Ti 配合 Titan X 的配置下完成）。代码还支持混合精度训练。

最后，正如“预训练模型链接”部分所示，我们在极具挑战性的 MPI-Sintel ‘final’ 数据集上取得了优于官方论文报道的结果。

目录

• 背景
• 环境搭建
• 预训练模型链接
• PWC-Net
  • 基本思想
  • 网络结构
  • Jupyter 笔记本
  • 训练
    • 多步学习率调度
    • 循环学习率调度
    • 混合精度训练
  • 评估
  • 推理
    • 在数据集的测试集上运行推理
    • 对图像对进行推理
• 数据集
• 参考文献
• 致谢

背景

光流估计的目的是生成一幅从一帧视频到下一帧视频之间发生的运动的稠密二维实数值（u,v 向量）映射。这一信息在解决计算机视觉问题时非常有用，例如 目标跟踪、动作识别、视频目标分割 等。

下图 [2017a] (a) 展示了来自 Middlebury 光流数据集的训练样本对（黑白帧 0 和 1），以及它们的彩色编码光流真实值。图 (b) 则说明了用于直观显示 (u,v) 流场的色彩编码规则：通常，向量方向由颜色色调表示，而向量长度则由颜色饱和度编码：

评估光流估计质量最常用的指标是 角度误差 (AE) 和 端点误差 (EPE)。两个光流向量 (u₀, v₀) 和 (u₁, v₁) 之间的角度误差定义为 arccos((u₀, v₀) · (u₁, v₁))。端点误差则衡量两个光流向量 (u₀, v₀) 和 (u₁, v₁) 的端点之间的距离，定义为 sqrt((u₀ - u₁)² + (v₀ - v₁)²)。

环境搭建

本仓库中的代码是在 Anaconda3 v.5.2.0 下开发并测试的。如需复现我们的 conda 环境，请参阅以下文件：

在 Ubuntu 上：

• conda list 和 conda env export

在 Windows 上：

• conda list 和 conda env export

预训练模型链接

预训练模型可在此处获取 这里。模型分为两种：“小”模型（sm，拥有 4,705,064 个学习参数）不使用密集连接或残差连接；“大”模型（lg，拥有 14,079,050 个学习参数）则使用。所有模型均采用 6 层金字塔结构，在每个维度上以 2×4 的比例进行上采样以生成最终预测，并在每一层构建一个 81 通道的代价体积，搜索范围（最大位移）为 4。

请注意，我们使用略有不同的数据集和学习率调度来训练这些模型。官方论文 [2018a] 中讨论的多步调度如下：S_long 训练 120 万次迭代，批次大小 8；随后 S_fine 微调 50 万次迭代，批次大小 4。而我们的调度仅为 S_long，训练 120 万次迭代，批次大小 8，训练数据混合了 FlyingChairs 和 FlyingThings3DHalfRes。其中，FlyingThings3DHalfRes 是我们自己版本的 FlyingThings3D，其每一对输入图像及其对应的光流真实值均已在每个维度上 缩小了两倍。此外，我们还采用了 不同的增强技术组合。

模型性能

模型名称	笔记本	FlyingChairs (384x512) AEPE	Sintel clean (436x1024) AEPE	Sintel final (436x1024) AEPE
pwcnet-lg-6-2-multisteps-chairsthingsmix	训练	1.44 (笔记本)	2.60 (笔记本)	3.70 (笔记本)
pwcnet-sm-6-2-multisteps-chairsthingsmix	训练	1.71 (笔记本)	2.96 (笔记本)	3.83 (笔记本)

作为参考，以下是官方报告的结果：

模型推理时间

我们还在几款 GPU 上测量了以下 MPI-Sintel（436×1024）数据集的推理时间：

模型名称	Titan X	GTX 1080	GTX 1080 Ti
pwcnet-lg-6-2-cyclic-chairsthingsmix	90ms	81ms	68ms
pwcnet-sm-6-2-cyclic-chairsthingsmix	68.5ms	64.4ms	53.8ms

关于上述数字的几点说明……

首先，请注意，本实现从设计上就具有可移植性，即它不使用任何用户自定义的 CUDA 内核，而官方的 NVIDIA 实现则会使用。我们的实现可以在任何操作系统和任何硬件配置上运行（甚至在没有 GPU 的系统上也能运行），只要该系统能够支持 TensorFlow。

其次，我们报告的这些时间数据是基于在 FlyingChairs 和 FlyingThings3DHalfRes 数据集上训练的模型的推理时间。如果您希望，可以对这些模型进行更长时间的训练，或者使用额外的数据集对其进行微调。换句话说，这些图尚未经过冻结处理。

在典型的生产环境中，您会在完成最终训练/微调后对模型进行冻结，并利用 TensorFlow XLA 或 TensorRT 将计算图优化至您需要分发到的目标平台。在这种重要场景下，我们在未优化的计算图上报告的推理时间其实并无太大意义。

PWC-Net

基本思想

根据 [2018a]，PWC Net 通过将领域知识融入网络设计，对 FlowNet2 [2016a] 进行了改进。光流估计的基本思想是：尽管像素在相邻帧之间发生了位置变化，但其亮度在时间上仍能保持基本不变（即“亮度恒定”）。我们可以先在第一帧视频中选取一个像素周围的小区域，然后在第二帧视频中寻找另一个小区域，使得这两个区域之间的某种函数值（例如归一化互相关）达到最大。接着将这个小区域在整个第一帧上滑动并寻找峰值，从而生成所谓的“代价体积”（PWC 中的 C）。这种技术相当稳健（对颜色变化具有不变性），但计算成本较高。在某些情况下，为了减少第一帧中的误检数量，可能需要使用较大的搜索区域，这进一步增加了计算复杂度。

为降低生成代价体积的成本，第一个优化措施是采用“金字塔处理”（PWC 中的 P）。通过使用较低分辨率的图像，可以在较小的第二帧版本上以较小的搜索区域进行滑动搜索，从而得到一个较小的运动向量，再将这一信息作为提示，在金字塔的下一更高分辨率层级上进行更有针对性的搜索。这种多尺度运动估计既可以在图像域进行，也可以在特征域进行（即利用卷积神经网络生成的下采样特征图）。实际上，PWC 会用在较低分辨率下估计出的运动流的上采样版本来“扭曲”（PWC 中的 W）第一帧，因为这样可以在金字塔的下一更高分辨率层级上搜索更小的运动增量，从而缩小搜索范围。以下是 Deqing Sun 在一次演讲中展示的使用两层金字塔的示意图：

需要注意的是，这里使用的三项优化措施（P/W/C）并非 PWC-Net 所独有。这些技术同样被用于 SpyNet [2016b] 和 FlowNet2 [2016a]。然而，在 PWC-Net 中，这些优化是在“CNN 特征”上进行的，而不是在图像金字塔上：

作者还承认，为了达到最佳性能，必须进行细致的数据增强（例如添加水平翻转）。此外，为了提高鲁棒性，作者建议在多个数据集上进行训练（例如 Sintel+KITTI+HD1K），并仔细平衡各类样本的分布不均问题。

由于该算法每次仅处理两帧连续图像，因此它与仅使用图像对的方法（而非 n 帧、n>2）具有相同的局限性。也就是说，如果某个物体移出了画面，预测的光流很可能会出现较大的 EPE。正如作者所指出的，使用更多帧的技术可以通过随时间传播运动信息来弥补这一局限性。此外，该模型有时也会对小型、快速移动的物体失效。

网络

以下是 [2018a] 中描述的网络架构图：

Jupyter 笔记本

测试此实现的推荐方法是使用以下 Jupyter 笔记本：

• 光流数据集（准备与检查）：在该笔记本中，我们：

  • 加载光流数据集，并（自动）创建训练模型所需的额外数据（首次加载时执行一次操作）。
  • 展示每个数据集的样本图像/光流。请注意，您必须已经下载并解压了主数据文件。有关各数据集的下载链接，请参阅[数据集]。

• PWC-Net 大模型训练（采用多步学习率调度）：在该笔记本中，我们：

  • 使用 PWC-Net 大模型（带有密集连接和残差连接），6 层金字塔结构，将第 2 层上采样放大 4 倍作为最终的光流预测。
  • 使用 [2016a] 中描述的 S_long 调度，在 FlyingChairs 和 FlyingThings3DHalfRes 数据集的混合数据上训练模型。
  • 在 PWC-Net 小模型训练（采用多步学习率调度） 中，我们训练模型的小版本（无密集或残差连接）。
  • 在 PWC-Net 大模型训练（采用循环学习率调度） 中，我们使用 Cyclic_short 调度训练大模型。
  • 在 PWC-Net 小模型训练（采用循环学习率调度） 中，我们使用 Cyclic_short 调度训练小模型（无密集或残差连接）。

• PWC-Net 大模型评估（在 FlyingChairs 验证集上）：在该笔记本中，我们：

  • 评估在 FlyingChairs 和 FlyingThings3DHalfRes 数据集混合数据上、采用 S_long 调度训练的 PWC-Net 大模型。
  • 在 FlyingChairs 数据集的验证集上进行评估，得到平均 EPE 为 1.44。
  • 进行基本误差分析。

• PWC-Net 大模型评估（在 MPI-Sintel 'clean' 上）：在该笔记本中，我们：

  • 评估在 FlyingChairs 和 FlyingThings3DHalfRes 数据集混合数据上、采用 S_long 调度训练的 PWC-Net 大模型。
  • 在 MPI-Sintel 数据集的**'clean'** 版本上进行评估，得到平均 EPE 为 2.60。
  • 进行基本误差分析。

• PWC-Net 大模型评估（在 MPI-Sintel 'final' 上）：在该笔记本中，我们：

  • 评估在 FlyingChairs 和 FlyingThings3DHalfRes 数据集混合数据上、采用 S_long 调度训练的 PWC-Net 大模型。
  • 在 MPI-Sintel 数据集的**'final'** 版本上进行评估，得到平均 EPE 为 3.70。
  • 进行基本误差分析。

训练

多步学习率调度

与原始论文不同，我们并未按顺序在 FlyingChairs 和 FlyingThings3D 上进行训练（即先在 FlyingChairs 上预训练，再在 FlyingThings3D 上微调）。这是因为 MPI-Sintel 数据集上的平均光流幅值仅为 13.5，而 FlyingChairs 和 FlyingThings3D 的平均光流幅值分别为 11.1 和 38。在我们的实验中，如果在 FlyingThings3D 上进行微调，反而会在 MPI-Sintel 上得到更差的结果。

通过使用 FlyingThings3D 数据集的半分辨率版本，其平均光流幅值为 19，与 FlyingChairs 和 MPI-Sintel 更为接近，我们获得了更为稳定的结果。随后，我们在 FlyingChairs 和 FlyingThings3DHalfRes 数据集的混合数据上进行训练。当然，这一混合数据还可以进一步扩展至其他数据集。

以下是上述列出的 S_long 训练笔记本的训练曲线：

请注意，如果您在运行上述训练笔记本时点击 Tensorboard 中的 IMAGE 选项卡，便可以可视化训练过程中几个验证样本的进展（包括各金字塔层级上的预测光流），如下所示：

循环学习率调度

如果您不想使用长时间的训练调度，但仍想尝试这段代码，不妨试试我们非常短的循环学习率调度（10 万次迭代，批量大小 8）。虽然结果远不如前者，但能够快速进行实验：

模型名称	笔记本	FlyingChairs (384x512) AEPE	Sintel clean (436x1024) AEPE	Sintel final (436x1024) AEPE
pwcnet-lg-6-2-cyclic-chairsthingsmix	训练	2.67 (笔记本)	3.99 (笔记本)	5.08 (笔记本)
pwcnet-sm-6-2-cyclic-chairsthingsmix	训练	2.79 (笔记本)	4.34 (笔记本)	5.3 (笔记本)

以下是 Cyclic_short 训练笔记本的训练曲线：

混合精度训练

您可以通过使用混合精度训练进一步加快训练速度。不过，同样不要指望能达到相同的精度水平：

模型名称	笔记本	FlyingChairs (384x512) AEPE	Sintel clean (436x1024) AEPE	Sintel final (436x1024) AEPE
pwcnet-sm-6-2-cyclic-chairsthingsmix-fp16	训练	2.47 (笔记本)	3.77 (笔记本)	4.90 (笔记本)

评估

如评估笔记本所示，且正如预期，当一帧到下一帧的平均光流幅值较高时，PWC-Net 模型进行准确光流预测的难度会增加：

尤其对于这一模型——以及任何其他基于两帧的运动估计器——当图像元素简单地从画面中消失或突然飞入时，生成准确预测尤为困难：

尽管如此，当平均运动适中时，小型和大型模型都能产生令人瞩目的结果：

推理

您可以通过两种方式调用此处提供的代码，为自己的数据集生成光流预测：

• 使用 ModelPWCNet 对象的 predict_from_img_pairs() 方法，传入一组图像对；
• 将 OpticalFlowDataset 对象传递给 ModelPWCNet 对象，并调用其 predict() 方法。

在图像对上运行推理

如果您想在自己的图像集上使用预训练的 PWC-Net 模型，可以按照如下示例，通过 ModelPWCNet 对象的 predict_from_img_pairs() 方法传入一组图像对：

from __future__ import absolute_import, division, print_function
from copy import deepcopy
from skimage.io import imread
from model_pwcnet import ModelPWCNet, _DEFAULT_PWCNET_TEST_OPTIONS
from visualize import display_img_pairs_w_flows

# 构建待处理的图像对列表
img_pairs = []
for pair in range(1, 4):
    image_path1 = f'./samples/mpisintel_test_clean_ambush_1_frame_00{pair:02d}.png'
    image_path2 = f'./samples/mpisintel_test_clean_ambush_1_frame_00{pair+1:02d}.png'
    image1, image2 = imread(image_path1), imread(image_path2)
    img_pairs.append((image1, image2))

# TODO：设置用于推理的设备
# 这里我们使用 GPU（若要使用 CPU 进行推理，请使用 '/device:CPU:0'）
gpu_devices = ['/device:GPU:0']  
controller = '/device:GPU:0'

# TODO：设置训练好的模型路径（请确保已从 http://bit.ly/tfoptflow 下载该模型）
ckpt_path = './models/pwcnet-lg-6-2-multisteps-chairsthingsmix/pwcnet.ckpt-595000'

# 配置模型以进行推理，首先采用默认选项
nn_opts = deepcopy(_DEFAULT_PWCNET_TEST_OPTIONS)
nn_opts['verbose'] = True
nn_opts['ckpt_path'] = ckpt_path
nn_opts['batch_size'] = 1
nn_opts['gpu_devices'] = gpu_devices
nn_opts['controller'] = controller

# 我们以四分之一分辨率模式运行 PWC-Net-large 模型
# 即采用 6 层金字塔结构，并在最终光流预测中将第 2 层以 4 倍放大
nn_opts['use_dense_cx'] = True
nn_opts['use_res_cx'] = True
nn_opts['pyr_lvls'] = 6
nn_opts['flow_pred_lvl'] = 2

# 本数据集中的图像尺寸并非 64 的倍数，而模型生成的光流则被填充至 64 的倍数。
# 因此，我们需要将预测的光流裁剪回原始尺寸
nn_opts['adapt_info'] = (1, 436, 1024, 2)

# 实例化模型以进行推理，并显示模型配置
nn = ModelPWCNet(mode='test', options=nn_opts)
nn.print_config()

# 生成预测并显示结果
pred_labels = nn.predict_from_img_pairs(img_pairs, batch_size=1, verbose=False)
display_img_pairs_w_flows(img_pairs, pred_labels)

上述代码可在 pwcnet_predict_from_img_pairs.ipynb 笔记本及 pwcnet_predict_from_img_pairs.py 脚本中找到。

在数据集的测试集上运行推理

如果您想从头训练一个 PWC-Net 模型，或使用自己的数据集对预训练的 PWC-Net 模型进行微调，则需要实现一个数据集处理器，该处理器需继承自 dataset_base.py 中的 OpticalFlowDataset 基类。

我们为多个知名数据集提供了若干数据集处理器，例如 MPI-Sintel（dataset_mpisintel.py）、FlyingChairs（dataset_flyingchairs.py）、FlyingThings3D（dataset_flyingthings3d.py）以及 KITTI（dataset_kitti.py）。这些处理器均可作为您实现自定义数据集的起点。请注意，这并不复杂；派生类只需告知基类哪些文件列表用于训练、验证和测试，其余繁重工作则由基类完成。

一旦您拥有数据处理器，便可将其传递给 ModelPWCNet 对象，并调用其 predict() 方法，为其测试集生成光流预测，具体操作可参见 pwcnet_predict.ipynb 笔记本及 pwcnet_predict.py 脚本。

数据集

用于光流估计的常见数据集包括：

• FlyingThings3D [图像对 + 光流 + all_unused_files.txt]
• FlyingChairs [图像对 + 光流 + FlyingChairs 训练与验证分割]
• MPI Sintel [zip]
• KITTI Flow 2012 [zip] 和/或 KITTI Flow 2015 [zip]

其他光流数据集（未在此处使用）：

• Middlebury 光流 [网页]
• 海德堡 HD1K 光流 [网页]

根据 [2018a]，KITTI 和 Sintel 目前是光流领域最具挑战性且应用最广泛的基准。KITTI 基准面向自动驾驶应用，其半稠密的真实光流数据通过激光雷达采集。2012 年的数据集仅包含静态场景。而 2015 年的数据集则通过人工标注扩展至动态场景，由于大运动、严重光照变化及遮挡，对现有方法提出了更高要求。

Sintel 基准则基于开源动画电影《Sintel》制作，分为清洁版和最终版两个阶段。其中，最终版包含强烈的大气效果、运动模糊及相机噪声，这些因素均会给现有方法带来严峻挑战。

参考文献

2018年

• [2018a] 孙等人，2018年。PWC-Net：基于金字塔、扭曲和代价体积的卷积神经网络光流估计。[arXiv] [网页] [PyTorch（官方）] [PyTorch] [PyTorch] [Caffe（官方）] [TensorFlow] [TensorFlow] [视频] [视频]

2017年

• [2017a] Baghaie等人，2017年。用于光流估计的稠密描述子：一项比较研究。[网页]

2016年

• [2016a] Ilg等人，2016年。FlowNet 2.0：深度网络在光流估计中的演进。[arXiv] [PyTorch（官方）] [TensorFlow]
• [2016b] Ranjan等人，2016年。SpyNet：基于空间金字塔网络的光流估计。[arXiv] [Torch（官方）] [PyTorch]

2015年

• [2015a] Fischer等人，2015年。FlowNet：利用卷积神经网络学习光流。[arXiv] [Tensorflow（FlowNet-S）]

致谢

我们还感谢以下其他TensorFlow实现：

• daigo0927的https://github.com/daigo0927/PWC-Net_tf
• djl11的https://github.com/djl11/PWC_Net_TensorFlow
• PatWie的https://github.com/tensorpack/tensorpack/tree/master/examples/OpticalFlow

@InProceedings{Sun2018PWC-Net,
  author    = {Deqing Sun and Xiaodong Yang and Ming-Yu Liu and Jan Kautz},
  title     = {{PWC-Net}: {CNNs} for Optical Flow Using Pyramid, Warping, and Cost Volume},
  booktitle = CVPR,
  year      = {2018},
}

@InProceedings\{DFIB15,
  author       = "A. Dosovitskiy and P. Fischer and E. Ilg and P. H{\"a}usser and C. Hazirbas and V. Golkov and P. v.d. Smagt and D. Cremers and T. Brox",
  title        = "FlowNet: Learning Optical Flow with Convolutional Networks",
  booktitle    = "IEEE International Conference on Computer Vision (ICCV)",
  month        = "Dec",
  year         = "2015",
  url          = "http://lmb.informatik.uni-freiburg.de//Publications/2015/DFIB15"
}

联系方式

如果您对本工作有任何疑问，请随时通过以下方式联系我们：

tfoptflow 快速上手指南

环境准备

• 系统：Linux / Windows 64 位
• Python：3.6+（推荐 Anaconda3 5.2.0）
• GPU：可选，支持单卡/多卡（CUDA 10.0+ 驱动即可）
• 依赖：TensorFlow ≥1.12（CPU/GPU 均可）

安装步骤

1. 克隆仓库

   git clone https://github.com/philferriere/tfoptflow.git
   cd tfoptflow
   

2. 创建并激活 Conda 环境

   • Linux

   conda env create -f setup/dlubu36.yml
   conda activate tfoptflow
   

   • Windows

   conda env create -f setup/dlwin36.yml
   activate tfoptflow
   

3. 安装额外依赖（如已用 yml 文件可跳过）

   pip install opencv-python tqdm
   

基本使用

1. 下载预训练模型

# 下载地址（浏览器打开即可）
http://bit.ly/tfoptflow
# 选择 pwcnet-lg-6-2-multisteps-chairsthingsmix.zip 或 pwcnet-sm-6-2-multisteps-chairsthingsmix.zip
unzip pwcnet-lg-6-2-multisteps-chairsthingsmix.zip -d models/

2. 单对图片推理

from model_pwcnet import ModelPWCNet
from visualize import display_img_pairs_w_flows

# 初始化模型
nn = ModelPWCNet(mode='test', ckpt_path='models/pwcnet-lg-6-2-multisteps-chairsthingsmix')

# 推理
img_pairs = [('data/frame_0010.png', 'data/frame_0011.png')]
pred_labels = nn.predict_from_img_pairs(img_pairs, batch_size=1, verbose=True)

# 可视化
display_img_pairs_w_flows(img_pairs, pred_labels)

3. 批量推理（数据集）

# 准备数据集目录结构
# data/
#   sintel/
#     test/
#       final/
#         alley_1/
#           frame_0001.png
#           frame_0002.png
#           ...

# 运行
nn.predict_from_dataset(
    dataset_root='data/sintel/test/final',
    save_dir='output',
    batch_size=4
)

4. 快速体验 Jupyter

jupyter notebook
# 打开 pwcnet_predict_from_img_pairs.ipynb 直接体验交互式推理

至此，您已完成 tfoptflow 的快速上手。如需训练或微调，请打开 pwcnet_train_lg-6-2-multisteps-chairsthingsmix.ipynb 参考完整流程。