| [](https://pypi.org/project/spatial-correlation-sampler/) | |
| # Pytorch Correlation module | |
| this is a custom C++/Cuda implementation of Correlation module, used e.g. in [FlowNetC](https://arxiv.org/abs/1504.06852) | |
| This [tutorial](http://pytorch.org/tutorials/advanced/cpp_extension.html) was used as a basis for implementation, as well as | |
| [NVIDIA's cuda code](https://github.com/NVIDIA/flownet2-pytorch/tree/master/networks/correlation_package) | |
| - 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 `2.1.0`. | |
| Before installation please check compatibility of your GPU and CUDA (_Compute Capability_) [nvidia docs](https://developer.nvidia.com/cuda-gpus). | |
| e.g RTX 6000 is using CC=8.9 so we are setting the environment variable to | |
| `export TORCH_CUDA_ARCH_LIST="8.9+PTX"` | |
| # Installation | |
| be reminded this module requires `python3-dev` to compile C++ code, e.g. on Ubuntu run: | |
| `apt install python3-dev` | |
| 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](https://github.com/ClementPinard/Pytorch-Correlation-extension/issues/1) for more information | |
| # 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` and`stride2` is `dilation_patch`, which behave like dilated convolutions | |
| * equivalent `max_displacement` is then `dilation_patch * (patch_size - 1) / 2`. | |
| * `dilation` is a new parameter, it acts the same way as dilated convolution regarding the correlation kernel | |
| * to get the right parameters for FlowNetC, you would have | |
| ``` | |
| kernel_size=1 | |
| patch_size=21, | |
| stride=1, | |
| padding=0, | |
| dilation=1 | |
| dilation_patch=2 | |
| ``` | |
| ## Example | |
| ```python | |
| import torch | |
| from spatial_correlation_sampler import SpatialCorrelationSampler, | |
| device = "cuda" | |
| batch_size = 1 | |
| channel = 1 | |
| H = 10 | |
| W = 10 | |
| dtype = torch.float32 | |
| input1 = torch.randint(1, 4, (batch_size, channel, H, W), dtype=dtype, device=device, requires_grad=True) | |
| input2 = torch.randint_like(input1, 1, 4).requires_grad_(True) | |
| #You can either use the function or the module. Note that the module doesn't contain any parameter tensor. | |
| #function | |
| out = spatial_correlation_sample(input1, | |
| input2, | |
| kernel_size=3, | |
| patch_size=1, | |
| stride=2, | |
| padding=0, | |
| dilation=2, | |
| dilation_patch=1) | |
| #module | |
| correlation_sampler = SpatialCorrelationSampler( | |
| kernel_size=3, | |
| patch_size=1, | |
| stride=2, | |
| padding=0, | |
| dilation=2, | |
| dilation_patch=1) | |
| out = correlation_sampler(input1, input2) | |
| ``` | |
| # 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](https://arxiv.org/abs/1504.06852), implemented [here](https://github.com/lmb-freiburg/flownet2) and [here](https://github.com/NVIDIA/flownet2-pytorch/blob/master/networks/FlowNetC.py). | |
| * Feel free to file an issue to add entries to this with your hardware ! | |
| ## CUDA Benchmark | |
| * See [here](https://gist.github.com/ClementPinard/270e910147119831014932f67fb1b5ea) for a benchmark script working with [NVIDIA](https://github.com/NVIDIA/flownet2-pytorch/tree/master/networks/correlation_package)'s code, and Pytorch. | |
| * Benchmark are launched with environment variable `CUDA_LAUNCH_BLOCKING` set to `1`. | |
| * Only `float32` is benchmarked. | |
| * FlowNetC correlation parameters where launched with the following command: | |
| ```bash | |
| CUDA_LAUNCH_BLOCKING=1 python benchmark.py --scale ms -k1 --patch 21 -s1 -p0 --patch_dilation 2 -b4 --height 48 --width 64 -c256 cuda -d float | |
| 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 | |
| ``` | |
| | implementation | Correlation parameters | device | pass | min time | avg time | | |
| | -------------- | ---------------------- | ------- | -------- | ------------: | ------------: | | |
| | ours | default | 980 GTX | forward | **5.745 ms** | **5.851 ms** | | |
| | ours | default | 980 GTX | backward | 77.694 ms | 77.957 ms | | |
| | NVIDIA | default | 980 GTX | forward | 13.779 ms | 13.853 ms | | |
| | NVIDIA | default | 980 GTX | backward | **73.383 ms** | **73.708 ms** | | |
| | | | | | | | | |
| | ours | FlowNetC | 980 GTX | forward | **26.102 ms** | **26.179 ms** | | |
| | ours | FlowNetC | 980 GTX | backward | **208.091 ms** | **208.510 ms** | | |
| | NVIDIA | FlowNetC | 980 GTX | forward | 35.363 ms | 35.550 ms | | |
| | NVIDIA | FlowNetC | 980 GTX | backward | 283.748 ms | 284.346 ms | | |
| ### Notes | |
| * The overhead of our implementation regarding `kernel_size` > 1 during backward 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](https://github.com/NVIDIA/flownet2-pytorch/blob/master/networks/correlation_package/src/correlation_cuda_kernel.cu#L120). | |
| ## CPU Benchmark | |
| * No other implementation is avalaible on CPU. | |
| * It is obviously not recommended to run it on CPU if you have a GPU. | |
| | Correlation parameters | device | pass | min time | avg time | | |
| | ---------------------- | -------------------- | -------- | ----------: | ----------: | | |
| | default | E5-2630 v3 @ 2.40GHz | forward | 159.616 ms | 188.727 ms | | |
| | default | E5-2630 v3 @ 2.40GHz | backward | 282.641 ms | 294.194 ms | | |
| | FlowNetC | E5-2630 v3 @ 2.40GHz | forward | 2.138 s | 2.144 s | | |
| | FlowNetC | E5-2630 v3 @ 2.40GHz | backward | 7.006 s | 7.075 s | | |