flownet2-pytorch

FlowNet 2.0: 使用深度网络进行光流估计的演变的 Pytorch 实现。

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

还支持使用 fp16(半精度)进行推理。

需要更多帮助,请输入

python main.py --help

网络架构

以下是提供的不同的 flownet 神经网络架构。
每个网络都有一个 batchnorm 版本。

自定义层

FlowNet2FlowNet2C* 架构依赖于自定义层 Resample2dCorrelation
这些层的 pytorch 实现与 cuda 内核可在 ./networks 中找到。
注意:目前,这些层没有半精度内核可用。

数据加载器

FlyingChairs、FlyingThings、ChairsSDHom 和 ImagesFromFolder 的数据加载器可在 datasets.py 中找到。

损失函数

支持多尺度的 L1 和 L2 损失函数可在 losses.py 中找到。

安装

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

# install custom layers
bash install.sh

Python 要求

目前,代码支持 python 3

转换后的 Caffe 预训练模型

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

推理

# Example on MPISintel 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 

训练和验证

# Example on MPISintel Final and Clean, with L1Loss on FlowNet2 model
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

# Example on MPISintel Final and Clean, with MultiScale loss on FlowNet2C model 
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 上的结果

Predicted flows on 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: Evolution of Optical Flow Estimation with Deep Networks",
  booktitle    = "IEEE Conference on Computer Vision and Pattern Recognition (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: Pytorch implementation of FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks},
  year = {2017},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/NVIDIA/flownet2-pytorch}}
}

代码(Caffe 和 Pytorch):PWC-Net
论文 : PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume

鸣谢

此代码的部分内容源自 ClementPinard/FlowNetPytorch,如代码中所述。