Home

Awesome

PyPI

** Modified to support feature-wise correlation. Only forward is supported. **

Pytorch 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

Requirements

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

Installation

this module is available on pip

pip install spatial-correlation-sampler

For a cpu-only version, you can install from source with

python setup_cpu.py install

Known Problems

This module needs compatible gcc version and CUDA to be compiled. Namely, CUDA 9.1 and below will need gcc5, while CUDA 9.2 and 10.0 will need gcc7 See this issue for more information

Usage

API has a few difference with NVIDIA's module

input (B x C x H x W) -> output (B x PatchH x PatchW x oH x oW)
kernel_size=1
patch_size=21,
stride=1,
padding=0,
dilation_patch=2

Benchmark

CUDA Benchmark

CUDA_LAUNCH_BLOCKING=1 python benchmark.py --scale ms -k1 --patch 21 -s1 -p0 --patch_dilation 2 -b4 --height 48 --width 64 -c256 cuda

CUDA_LAUNCH_BLOCKING=1 python NV_correlation_benchmark.py --scale ms -k1 --patch 21 -s1 -p0 --patch_dilation 2 -b4 --height 48 --width 64 -c256
implementationCorrelation parametersdevicepassmin timeavg time
oursdefault980 GTXforward5.745 ms5.851 ms
oursdefault980 GTXbackward77.694 ms77.957 ms
NVIDIAdefault980 GTXforward13.779 ms13.853 ms
NVIDIAdefault980 GTXbackward73.383 ms73.708 ms
oursFlowNetC980 GTXforward26.102 ms26.179 ms
oursFlowNetC980 GTXbackward208.091 ms208.510 ms
NVIDIAFlowNetC980 GTXforward35.363 ms35.550 ms
NVIDIAFlowNetC980 GTXbackward283.748 ms284.346 ms

Notes

CPU Benchmark

Correlation parametersdevicepassmin timeavg time
defaultE5-2630 v3 @ 2.40GHzforward159.616 ms188.727 ms
defaultE5-2630 v3 @ 2.40GHzbackward282.641 ms294.194 ms
FlowNetCE5-2630 v3 @ 2.40GHzforward2.138 s2.144 s
FlowNetCE5-2630 v3 @ 2.40GHzbackward7.006 s7.075 s