flownet2-pytorch

PyTorch 实现的 FlowNet 2.0：基于深度网络的光流估计演进。

支持多 GPU 训练，代码提供了在 MPI-Sintel 清洁版和最终版数据集上进行训练或推理的示例。相同的命令也可用于其他数据集的训练或推理。详情请见下文。

还支持使用 fp16（半精度）进行推理。

如需更多帮助，请输入：

python main.py --help

网络架构

以下是提供的不同 FlowNet 神经网络架构。
每种网络也提供了一个带有批归一化层的版本。

• FlowNet2S
• FlowNet2C
• FlowNet2CS
• FlowNet2CSS
• FlowNet2SD
• FlowNet2

自定义层

FlowNet2 或 FlowNet2C* 架构依赖于自定义层 Resample2d 或 Correlation。
这些层的 PyTorch 实现及其 CUDA 内核可在 ./networks 中找到。
注意：目前这些层尚无半精度内核。

数据加载器

datasets.py 中提供了 FlyingChairs、FlyingThings、ChairsSDHom 和 ImagesFromFolder 的数据加载器。

损失函数

losses.py 中提供了支持多尺度的 L1 和 L2 损失。

安装

# 获取 flownet2-pytorch 源码
git clone https://github.com/NVIDIA/flownet2-pytorch.git
cd flownet2-pytorch

# 安装自定义层
bash install.sh

Python 要求

目前代码支持 Python 3：

• numpy
• PyTorch (== 0.4.1，对于 <= 0.4.0，请参阅分支 python36-PyTorch0.4)
• scipy
• scikit-image
• tensorboardX
• colorama、tqdm、setproctitle

转换后的 Caffe 预训练模型

我们包含了 Caffe 预训练模型。如果您使用这些预训练权重，请遵守 许可协议。

• FlowNet2[620MB]
• FlowNet2-C[149MB]
• FlowNet2-CS[297MB]
• FlowNet2-CSS[445MB]
• FlowNet2-CSS-ft-sd[445MB]
• FlowNet2-S[148MB]
• FlowNet2-SD[173MB]

推理

# MPI Sintel 清洁版示例
python main.py --inference --model FlowNet2 --save_flow --inference_dataset MpiSintelClean \
--inference_dataset_root /path/to/mpi-sintel/clean/dataset \
--resume /path/to/checkpoints 

训练与验证

# MPI Sintel 最终版和清洁版示例，使用 L1Loss 损失函数训练 FlowNet2 模型
python main.py --batch_size 8 --model FlowNet2 --loss=L1Loss --optimizer=Adam --optimizer_lr=1e-4 \
--training_dataset MpiSintelFinal --training_dataset_root /path/to/mpi-sintel/final/dataset  \
--validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset

# MPI Sintel 最终版和清洁版示例，使用 MultiScale 损失函数训练 FlowNet2C 模型 
python main.py --batch_size 8 --model FlowNet2C --optimizer=Adam --optimizer_lr=1e-4 --loss=MultiScale --loss_norm=L1 \
--loss_numScales=5 --loss_startScale=4 --optimizer_lr=1e-4 --crop_size 384 512 \
--training_dataset FlyingChairs --training_dataset_root /path/to/flying-chairs/dataset  \
--validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset

MPI-Sintel 上的结果

参考文献

如果您在工作中发现此实现有用，请适当致谢并引用以下论文：

@InProceedings{IMKDB17,
  author       = "E. Ilg and N. Mayer and T. Saikia and M. Keuper and A. Dosovitskiy and T. Brox",
  title        = "FlowNet 2.0：基于深度网络的光流估计演进",
  booktitle    = "IEEE 计算机视觉与模式识别会议 (CVPR)",
  month        = "Jul",
  year         = "2017",
  url          = "http://lmb.informatik.uni-freiburg.de//Publications/2017/IMKDB17"
}

@misc{flownet2-pytorch,
  author = {Fitsum Reda and Robert Pottorff and Jon Barker and Bryan Catanzaro},
  title = {flownet2-pytorch：FlowNet 2.0：基于深度网络的光流估计演进的 PyTorch 实现},
  year = {2017},
  publisher = {GitHub},
  journal = {GitHub 仓库},
  howpublished = {\url{https://github.com/NVIDIA/flownet2-pytorch}}
}

英伟达相关的光流工作

代码（Caffe 和 PyTorch）：PWC-Net
论文：PWC-Net：利用金字塔、变形和代价体积的光流 CNN。

致谢

本代码的部分内容源自 ClementPinard/FlowNetPytorch，如代码中所注明。

FlowNet2-PyTorch 快速上手指南

FlowNet2-PyTorch 是 NVIDIA 开源的 FlowNet 2.0 光流估计模型的 PyTorch 实现，支持多 GPU 训练、半精度（fp16）推理，并提供多种网络架构变体。

环境准备

系统要求

• 操作系统: Linux (推荐)
• Python: Python 3.x
• GPU: 支持 CUDA 的 NVIDIA 显卡（用于自定义层编译和加速）

前置依赖

请确保已安装以下 Python 库：

pip install numpy scipy scikit-image tensorboardX colorama tqdm setproctitle

PyTorch 版本注意：

• 推荐版本：PyTorch == 0.4.1
• 若使用 <= 0.4.0 版本，请切换至分支 python36-PyTorch0.4。

国内加速建议：安装依赖时可使用清华或阿里镜像源加速：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple numpy scipy scikit-image tensorboardX colorama tqdm setproctitle

安装步骤

1. 克隆项目代码

   git clone https://github.com/NVIDIA/flownet2-pytorch.git
   cd flownet2-pytorch
   

   国内加速：若 GitHub 克隆速度慢，可使用国内镜像：

   git clone https://gitee.com/mirrors/flownet2-pytorch.git
   # 或者
   git clone https://github.com.cnpmjs.org/NVIDIA/flownet2-pytorch.git
   

2. 编译并安装自定义层 该项目包含 Resample2d 和 Correlation 等自定义 CUDA 层，必须编译安装：

   bash install.sh
   

   注意：当前自定义层不支持半精度（half precision）内核。

3. 下载预训练模型（可选） 如需直接进行推理，可从 Google Drive 下载转换后的 Caffe 预训练权重（如 FlowNet2, FlowNet2S 等），并将其放入项目目录。 注：下载链接需遵循相应的许可协议。

基本使用

1. 查看帮助信息

了解所有可用参数：

python main.py --help

2. 推理示例 (Inference)

使用预训练模型在 MPI-Sintel Clean 数据集上进行光流预测并保存结果：

python main.py --inference --model FlowNet2 --save_flow --inference_dataset MpiSintelClean \
--inference_dataset_root /path/to/mpi-sintel/clean/dataset \
--resume /path/to/checkpoints

请将 /path/to/... 替换为本地实际的数据集路径和模型权重文件路径。

3. 训练示例 (Training)

使用 FlyingChairs 数据集训练 FlowNet2C 模型，并在 MPI-Sintel Clean 上进行验证：

python main.py --batch_size 8 --model FlowNet2C --optimizer=Adam --optimizer_lr=1e-4 --loss=MultiScale --loss_norm=L1 \
--loss_numScales=5 --loss_startScale=4 --crop_size 384 512 \
--training_dataset FlyingChairs --training_dataset_root /path/to/flying-chairs/dataset  \
--validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset

支持的网络架构

在命令中通过 --model 参数指定以下任意一种架构：

• FlowNet2S
• FlowNet2C
• FlowNet2CS
• FlowNet2CSS
• FlowNet2SD
• FlowNet2