flownet2-pytorch
Pytorch implementation of FlowNet 2.0: Evolution of Optical Flow Estimation with Deep Networks.
FlowNet2 Caffe implementation : flownet2
Multiple GPU training is supported, and the code provides examples for training or inference on MPI-Sintel clean and final datasets. The same commands can be used for training or inference with other datasets. See below for more detail.
Inference using fp16 (half-precision) is also supported.
For more help, type
python main.py --help
Network architectures
Below are the different flownet neural network architectures that are provided.
A batchnorm version for each network is available.
- FlowNet2S
- FlowNet2C
- FlowNet2CS
- FlowNet2CSS
- FlowNet2SD
- FlowNet2
Custom layers
FlowNet2 or FlowNet2C* achitectures rely on custom layers Resample2d or Correlation.
A pytorch implementation of these layers with cuda kernels are available at ./networks.
Note : Currently, half precision kernels are not available for these layers.
Data Loaders
Dataloaders for FlyingChairs, FlyingThings, ChairsSDHom and ImagesFromFolder are available in datasets.py.
Loss Functions
L1 and L2 losses with multi-scale support are available in losses.py.
Installation
# get flownet2-pytorch source
git clone https://github.com/NVIDIA/flownet2-pytorch.git
cd flownet2-pytorch
# install custom layers
bash install.sh
Docker image
Libraries and other dependencies for this project include: Ubuntu 16.04, Python 2.7, Pytorch 0.2, CUDNN 6.0, CUDA 8.0
A Dockerfile with the above dependencies is available
# Build and launch docker image
bash launch_docker.sh
Inference
# 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
Training and validation
# 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 \
--training_dataset MpiSintelFinal --training_dataset_root /path/to/mpi-sintel/final/dataset \
--validation_dataset MpiSintelClean --validation_dataset_root /path/to/mpi-sintel/clean/dataset
Results on MPI-Sintel
Acknowledgments
The Flownet-S model in this project was derived from ClementPinard/FlowNetPytorch.