aikenml/data_mining

1 2[![PyPI](https://img.shields.io/pypi/v/spatial-correlation-sampler.svg)](https://pypi.org/project/spatial-correlation-sampler/)3 4 5# Pytorch Correlation module6 7this is a custom C++/Cuda implementation of Correlation module, used e.g. in [FlowNetC](https://arxiv.org/abs/1504.06852)8 9This [tutorial](http://pytorch.org/tutorials/advanced/cpp_extension.html) was used as a basis for implementation, as well as10[NVIDIA's cuda code](https://github.com/NVIDIA/flownet2-pytorch/tree/master/networks/correlation_package)11 12- Build and Install C++ and CUDA extensions by executing `python setup.py install`,13- Benchmark C++ vs. CUDA by running `python benchmark.py {cpu, cuda}`,14- Run gradient checks on the code by running `python grad_check.py --backend {cpu, cuda}`.15 16# Requirements17 18This module is expected to compile for Pytorch `2.1.0`.19 20Before installation please check compatibility of your GPU and CUDA (_Compute Capability_) [nvidia docs](https://developer.nvidia.com/cuda-gpus). 21e.g RTX 6000 is using CC=8.9 so we are setting the environment variable to22 23`export TORCH_CUDA_ARCH_LIST="8.9+PTX"`24 25# Installation26 27be reminded this module requires `python3-dev` to compile C++ code, e.g. on Ubuntu run:28 29`apt install python3-dev`30 31this module is available on pip32 33`pip install spatial-correlation-sampler`34 35For a cpu-only version, you can install from source with36 37`python setup_cpu.py install`38 39# Known Problems40 41This module needs compatible gcc version and CUDA to be compiled.42Namely, CUDA 9.1 and below will need gcc5, while CUDA 9.2 and 10.0 will need gcc743See [this issue](https://github.com/ClementPinard/Pytorch-Correlation-extension/issues/1) for more information44 45# Usage46 47API has a few difference with NVIDIA's module48 * output is now a 5D tensor, which reflects the shifts horizontal and vertical.49 ```50input (B x C x H x W) -> output (B x PatchH x PatchW x oH x oW)51 ```52 * Output sizes `oH` and `oW` are no longer dependant of patch size, but only of kernel size and padding53 * Patch size `patch_size` is now the whole patch, and not only the radii.54 * `stride1` is now `stride` and`stride2` is `dilation_patch`, which behave like dilated convolutions55 * equivalent `max_displacement` is then `dilation_patch * (patch_size - 1) / 2`.56 * `dilation` is a new parameter, it acts the same way as dilated convolution regarding the correlation kernel57 * to get the right parameters for FlowNetC, you would have58 ```59kernel_size=160patch_size=21,61stride=1,62padding=0,63dilation=164dilation_patch=265 ```66 67 68## Example69```python70import torch71from spatial_correlation_sampler import SpatialCorrelationSampler, 72 73device = "cuda"74batch_size = 175channel = 176H = 1077W = 1078dtype = torch.float3279 80input1 = torch.randint(1, 4, (batch_size, channel, H, W), dtype=dtype, device=device, requires_grad=True)81input2 = torch.randint_like(input1, 1, 4).requires_grad_(True)82 83#You can either use the function or the module. Note that the module doesn't contain any parameter tensor.84 85#function86 87out = spatial_correlation_sample(input1,88	                         input2,89                                 kernel_size=3,90                                 patch_size=1,91                                 stride=2,92                                 padding=0,93                                 dilation=2,94                                 dilation_patch=1)95 96#module97 98correlation_sampler = SpatialCorrelationSampler(99    kernel_size=3,100    patch_size=1,101    stride=2,102    padding=0,103    dilation=2,104    dilation_patch=1)105out = correlation_sampler(input1, input2)106 107```108 109# Benchmark110 111 * 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).112 * Feel free to file an issue to add entries to this with your hardware !113 114## CUDA Benchmark115 116 * 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.117 * Benchmark are launched with environment variable `CUDA_LAUNCH_BLOCKING` set to `1`.118 * Only `float32` is benchmarked.119 * FlowNetC correlation parameters where launched with the following command:120 121 ```bash122 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 float123 124 CUDA_LAUNCH_BLOCKING=1 python NV_correlation_benchmark.py --scale ms -k1 --patch 21 -s1 -p0 --patch_dilation 2 -b4 --height 48 --width 64 -c256125 ```126 127 | implementation | Correlation parameters |  device |     pass |      min time |      avg time |128 | -------------- | ---------------------- | ------- | -------- | ------------: | ------------: |129 |           ours |                default | 980 GTX |  forward |  **5.745 ms** |  **5.851 ms** |130 |           ours |                default | 980 GTX | backward |     77.694 ms |     77.957 ms |131 |         NVIDIA |                default | 980 GTX |  forward |     13.779 ms |     13.853 ms |132 |         NVIDIA |                default | 980 GTX | backward | **73.383 ms** | **73.708 ms** |133 |                |                        |         |          |               |               |134 |           ours |               FlowNetC | 980 GTX |  forward |  **26.102 ms** |  **26.179 ms** |135 |           ours |               FlowNetC | 980 GTX | backward | **208.091 ms** | **208.510 ms** |136 |         NVIDIA |               FlowNetC | 980 GTX |  forward |      35.363 ms |      35.550 ms |137 |         NVIDIA |               FlowNetC | 980 GTX | backward |     283.748 ms |     284.346 ms |138 139### Notes140 * The overhead of our implementation regarding `kernel_size` > 1 during backward needs some investigation, feel free to141 dive in the code to improve it !142 * The backward pass of NVIDIA is not entirely correct when stride1 > 1 and kernel_size > 1, because not everything143 is computed, see [here](https://github.com/NVIDIA/flownet2-pytorch/blob/master/networks/correlation_package/src/correlation_cuda_kernel.cu#L120).144 145## CPU Benchmark146 147  * No other implementation is avalaible on CPU.148  * It is obviously not recommended to run it on CPU if you have a GPU.149 150 | Correlation parameters |               device |     pass |    min time |    avg time |151 | ---------------------- | -------------------- | -------- | ----------: | ----------: |152 |                default | E5-2630 v3 @ 2.40GHz |  forward |  159.616 ms |  188.727 ms |153 |                default | E5-2630 v3 @ 2.40GHz | backward |  282.641 ms |  294.194 ms |154 |               FlowNetC | E5-2630 v3 @ 2.40GHz |  forward |  2.138 s |  2.144 s |155 |               FlowNetC | E5-2630 v3 @ 2.40GHz | backward | 7.006 s | 7.075 s |156