Correlation module

this is a custom C++/Cuda implementation of Correlation module, used e.g. in FlowNetC

This tutorial was used as a basis for implementation, as well as NVIDIA's cuda code

• Build and Install C++ and CUDA extensions by executing python setup.py install,
• Benchmark C++ vs. CUDA by running python benchmark.py {cpu, cuda},
• Run gradient checks on the code by running python grad_check.py --backend {cpu, cuda}.

Requirements

This module is expected to compile for Pytorch 0.4.1, on Python > 3.5 and Python 2.7.

Installation

this module is available on pip

pip install spatial-correlation-sampler

Usage

API has a few difference with NVIDIA's module

• output is now a 5D tensor, which reflects the shifts horizontal and vertical.

input (B x C x H x W) -> output (B x PatchH x PatchW x oH x oW)

• Output sizes oH and oW are no longer dependant of patch size, but only of kernel size and padding
• Patch size patch_size is now the whole patch, and not only the radii.
• stride1 is now stride andstride2 is dilation_patch, which behave like dilated convolutions
• equivalent max_displacement is then dilation_patch * (patch_size - 1) / 2.
• to get the right parameters for FlowNetC, you would have

kernel_size=1
patch_size=21,
stride=1,
padding=0,
dilation_patch=2

Benchmark

• default parameters are from benchmark.py, FlowNetC parameters are same as use in FlowNetC with a batch size of 4, described in this paper, implemented here and here.
• Feel free to file an issue to add entries to this with your hardware !

CUDA Benchmark

• See here for a benchmark script working with NVIDIA's code, and Pytorch 0.3.
• Benchmark are launched with environment variable CUDA_LAUNCH_BLOCKING set to 1.
• Only float32 is benchmarked.

Notes

• The large overhead of our implementation regarding kernel_size > 1 needs some investigation, feel free to dive in the code to improve it !
• The backward pass of NVIDIA is not entirely correct when stride1 > 1 and kernel_size > 1, because not everything is computed, see here.

CPU Benchmark

• No other implementation is avalaible on CPU.