PyTorch Lightning Optical Flow

Introduction

This is a collection of state-of-the-art deep model for estimating optical flow. The main goal is to provide a unified framework where multiple models can be trained and tested more easily.

The work and code from many others are present here. I tried to make sure everything is properly referenced, but please let me know if I missed something.

This is still under development, so some things may not work as intended. I plan to add more models in the future, as well keep improving the platform.

• What's new
• Available models
• Results
• Getting started
• Licenses
• Contributing
• Citing
• Acknowledgements

What's new

- v0.5.0

Major update focusing on the correlation computation of the models and on deployment.

• Added features:
  • Integration with the localcorr package, which provides fast correlation kernels (Triton, CUDA, and SYCL)
  • New corr_mode argument to select the implementation used to compute the correlation of the models (e.g., allpairs, torch, or the localcorr engines). If the selected mode is not available, a pure PyTorch implementation is used as fallback. See the installation instructions for more details.
  • Each correlation mode can also be chosen by using its model variant directly, e.g., --model raft_localcorr_tritonfor NVIDIA and Intel GPUs. It is installed automatically on Linux.
  • Support for Intel GPUs (XPU) in the infer, validate, and model_benchmark scripts
  • All models now accept the simple_io argument, and many model operations were revised to be exportable to ONNX and TensorRT
• Added new models:
  • MegaFlow https://arxiv.org/abs/2603.25739
  • UFM https://arxiv.org/abs/2506.09278
• Additions / fixes:
  • The old alternate_corr and corr_fn arguments are still accepted, but they are deprecated and converted into corr_mode
  • Speed up metrics calculation and validation

- v0.4.2

• Added new models:
  • Flow-Anything https://arxiv.org/abs/2506.07740
  • FlowSeek https://arxiv.org/abs/2509.05297
  • MEMFOF https://arxiv.org/abs/2506.23151
  • ReCoVEr https://arxiv.org/abs/2510.13317
  • WAFT https://arxiv.org/abs/2506.21526
• Additions / fixes:
  • Fix compatibility with Lightning < 2.7
  • Add support to log to SwanLab

- v0.4.1

• Added new models:
  • DPFlow https://arxiv.org/abs/2503.14880
  • StreamFlow https://arxiv.org/abs/2311.17099
• Additions / fixes:
  • Add 4K read mode for Spring dataset
  • Add missing arguments for Middlebury-ST and VIPER datasets
  • Fix RAFT config files
  • Fix RAPIDFlow config files

- v0.4.0

Major update to support Lightning 2 (finally!). However, it also introduces breaking changes from the previous v0.3 code. See the details below.

• Transitioning from v0.3 to v0.4: check the v0.4 upgrade guide
• Added features:
  • Support for YAML config files. See the config file documentation
  • Table comparing PTLFlow results with the original papers to check the stability of the included models.
• Added new models:
  • NeuFlow v2 https://arxiv.org/abs/2408.10161
• Add support for more datasets:
  • Middlebury-ST [https://vision.middlebury.edu/stereo/data/scenes2014/]{https://vision.middlebury.edu/stereo/data/scenes2014/}
  • VIPER https://playing-for-benchmarks.org/

- v0.3.2

• Added new models:
  • MemFlow https://arxiv.org/abs/2404.04808
  • NeuFlow https://arxiv.org/abs/2403.10425
  • SEA-RAFT https://arxiv.org/abs/2405.14793
  • SplatFlow https://arxiv.org/abs/2306.08887
• Add support for more datasets:
  • TartanAir https://theairlab.org/tartanair-dataset/
  • Kubric https://github.com/google-research/kubric
• Add ONNX and TensorRT conversion to RAPIDFlow
• Fix LR scheduler when accumulating gradients

- v0.3.1

• Added new models:
  • CCMR https://arxiv.org/abs/2311.02661
  • LLA-Flow https://arxiv.org/abs/2304.08101
  • RAPIDFlow https://hmorimitsu.com/publication/2024-icra-rapidflow/
• Enable FP16 in most models.
  • Except the following models, since they have operations that cannot run in FP16: lcv_raft, matchflow, and separableflow
• Add FP16 mode in infer, model_benchmark, and validate scripts
• Create plot_results.py script
• Move resize operations to CUDA (thanks to coca-huang)

- v0.3.0

• Added new models:
  • DIP https://arxiv.org/abs/2204.00330
  • Flow1D https://arxiv.org/abs/2103.04524
  • FlowFormer++ https://arxiv.org/abs/2303.01237
  • GMFlow+, UniMatch https://arxiv.org/abs/2211.05783
  • MatchFlow https://arxiv.org/abs/2303.08384
  • MS-RAFT+ https://arxiv.org/abs/2210.16900
  • RPKNet https://hmorimitsu.com/publication/2024-aaai-rpknet
  • SeparableFlow https://openaccess.thecvf.com/content/ICCV2021/papers/Zhang_Separable_Flow_Learning_Motion_Cost_Volumes_for_Optical_Flow_Estimation_ICCV_2021_paper.pdf
  • SKFlow https://arxiv.org/abs/2205.14623
  • VideoFlow https://arxiv.org/abs/2303.08340
• speed_benchmark.py becomes model_benchmark.py and records more metrics
• Fix compatibility with PyTorch 2.0
• Fix compatibility with PyTorch Lightning 1.9
• Fix resizing augmentation when the valid mask is sparse
• Add support for more datasets:
  • Middlebury https://vision.middlebury.edu/flow/
  • Monkaa https://lmb.informatik.uni-freiburg.de/resources/datasets/SceneFlowDatasets.en.html
  • Spring https://spring-benchmark.org/

Available models

• CCMR https://arxiv.org/abs/2311.02661
• CRAFT https://arxiv.org/abs/2203.16896
• CSFlow https://arxiv.org/abs/2202.00909
• DICL-Flow https://arxiv.org/abs/2010.14851
• DIP https://arxiv.org/abs/2204.00330
• DPFlow https://arxiv.org/abs/2503.14880
• FastFlowNet https://arxiv.org/abs/2103.04524
• Flow1D https://arxiv.org/abs/2103.04524
• Flow-Anything https://arxiv.org/abs/2506.07740
• FlowFormer https://arxiv.org/abs/2203.16194
• FlowFormer++ https://arxiv.org/abs/2303.01237
• FlowNet https://arxiv.org/abs/1504.06852
• FlowNet2 https://arxiv.org/abs/1612.01925
• FlowSeek https://arxiv.org/abs/2509.05297
• GMA https://arxiv.org/abs/2104.02409
• GMFlow https://arxiv.org/abs/2111.13680
• GMFlow+, UniMatch https://arxiv.org/abs/2211.05783
• GMFlowNet https://arxiv.org/abs/2203.11335
• HD3 https://arxiv.org/abs/1812.06264
• IRR https://arxiv.org/abs/1904.05290
• LCV https://arxiv.org/abs/2007.11431
• LiteFlowNet https://arxiv.org/abs/1805.07036
• LiteFlowNet2 https://arxiv.org/abs/1903.07414
• LiteFlowNet3 https://arxiv.org/abs/2007.09319
• LLA-Flow https://arxiv.org/abs/2304.08101
• MaskFlownet https://arxiv.org/abs/2003.10955
• MatchFlow https://arxiv.org/abs/2303.08384
• MegaFlow https://arxiv.org/abs/2603.25739
• MemFlow https://arxiv.org/abs/2404.04808
• MEMFOF https://arxiv.org/abs/2506.23151
• MS-RAFT+ https://arxiv.org/abs/2210.16900
• NeuFlow v1 https://arxiv.org/abs/2403.10425
• NeuFlow v2 https://arxiv.org/abs/2408.10161
• PWCNet https://arxiv.org/abs/1709.02371
• RAFT https://arxiv.org/abs/2003.12039
• RAPIDFlow https://hmorimitsu.com/publication/2024-icra-rapidflow/
• ReCoVEr https://arxiv.org/abs/2510.13317
• RPKNet https://hmorimitsu.com/publication/2024-aaai-rpknet
• ScopeFlow https://arxiv.org/abs/2002.10770
• SCV https://arxiv.org/abs/2104.02166
• SEA-RAFT https://arxiv.org/abs/2405.14793
• SeparableFlow https://openaccess.thecvf.com/content/ICCV2021/papers/Zhang_Separable_Flow_Learning_Motion_Cost_Volumes_for_Optical_Flow_Estimation_ICCV_2021_paper.pdf
• SKFlow https://arxiv.org/abs/2205.14623
• SplatFlow https://arxiv.org/abs/2306.08887
• STaRFlow https://arxiv.org/abs/2007.05481
• StreamFlow https://arxiv.org/abs/2311.17099
• UFM https://arxiv.org/abs/2506.09278
• VCN https://papers.nips.cc/paper/2019/file/bbf94b34eb32268ada57a3be5062fe7d-Paper.pdf
• VideoFlow https://arxiv.org/abs/2303.08340
• WAFT https://arxiv.org/abs/2506.21526

Read more details about the models on https://ptlflow.readthedocs.io/en/latest/models/models_list.html.

Results

You can see a table with main evaluation results of the available models here. More results are also available in the folder docs/source/results.

Disclaimer: These results are the ones obtained by evaluating the available models in this framework in my machine. Your results may be different due to differences in hardware and software. I also do not guarantee that the results of each model will be similar to the ones presented in the respective papers or other original sources. If you need to replicate the original results from a paper, you should use the original implementations.

Getting started

Please take a look at the documentation to learn how to install and use PTLFlow.

You can also check the notebooks below running on Google Colab for some practical examples:

• Inference with a pretrained model.
• Training and using the learned weights for inference.

If you are using the previous v0.3.X code, then check the v0.3.2 documentation and the following example notebooks:

• Inference with a pretrained model (PTLFlow v0.3).
• Training and using the learned weights for inference (PTLFlow v0.3).

Licenses

The original code of this repository is licensed under the Apache 2.0 license.

Each model may be subjected to different licenses. The license of each model is included in their respective folders. It is your responsibility to make sure that your project is in compliance with all the licenses and conditions involved.

The external pretrained weights all have different licenses, which are listed in their respective folders.

The pretrained weights that were trained within this project are available under the CC BY-NC-SA 4.0 license, which I believe that covers the licenses of the datasets used in the training. That being said, I am not a legal expert so if you plan to use them to any purpose other than research, you should check all the involved licenses by yourself. Additionally, the datasets used for the training usually require the user to cite the original papers, so be sure to include their respective references in your work.

Contributing

Contribution are welcome! Please check CONTRIBUTING.md to see how to contribute.

Citing

BibTeX

@misc{morimitsu2021ptlflow,
  author = {Henrique Morimitsu},
  title = {PyTorch Lightning Optical Flow},
  year = {2021},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/hmorimitsu/ptlflow}}
}

Acknowledgements

• This README file is heavily inspired by the one from the timm repository.
• Some parts of the code were inspired by or taken from FlowNetPytorch.
• flownet2-pytorch was also another important source.
• The current main training routine is based on RAFT.