Feature request: Correlation module

[ClementPinard]

This module was used by FlowNetC, it is described here (in section 3.)
Here is the caffe implementation of the authors
https://github.com/lmb-freiburg/flownet2/blob/master/src/caffe/layers/correlation_layer.cu
Here is one tensorflow implementation
https://github.com/sampepose/flownet2-tf/tree/master/src/ops/correlation

And also NVIDIA recently released a custom pytorch operator here
https://github.com/NVIDIA/flownet2-pytorch/tree/master/networks/correlation_package
here is another attempt in pytorch (but only high level)
https://github.com/onlytailei/pytorch-flownet2/blob/master/nets/flownetC.py#L45

Would it be useful to add such a feature in functional pytorch library ?

[kingsleyluoxin]

I have been wondering if anyone knows how to understand the codes of these self-defined operations.

[ClementPinard]

Basically, for each point in your output map, you compute the correlation of the two inputs, as if the patch of your input2 was a kernel of convolution applied on patch of ̀input1

other than padding there are 4 fundamental arguments to apply :
kernel size, displacement size, stride and dilation (for simplicity, we'll assume stride and dilation of 1 for now)

with a maximum displacement size, you would get this mapping :
Input1 BxCxHxW Input2 BxCxHxW -> Output BxHxWxHxW
which means you get to treat every possible patch of Input2 as normal convolution kernel, applied to the whole Input1
the minimum displacement you would get :
Output BxHxWx1x1 
The operation done is then just a bunch of inner products, for all (n,i,j) you compute the inner product of the two C-dimension vectors <Input1[n,:,i,j], Input2[n,:,i,j]> and then you apply a avgPool2d of your map with a stride of 1 and kernel_size of your correlation kernel size.

As all these operations are very similar to convolution, I have been studying on how we can make these operations simple GEMM calls, whether it is batched/strided or not (see here ), inspired by this presentation : http://www.prime-project.org/wp-content/uploads/sites/206/2018/02/Talk-10-David-Gregg-Parallel-Multi-Channel-Convolution-using-General-Matrix-Multiplication.pdf , but I cam across nothing conclusive, it seems we are bound to make batch GEMM call of very little matrices, and unable to benefit from CuBLAS capabilities.

I think the best options might be :

• Do what the paper did with convolution, but instead of a GEMM primitive, we use a inner product primitive : construct every possible broadcast inner product of C-dimensional pixels (with spatial shift induced by displacement) followed by a avgpooling
• Do what Fitsum Reda did in NVIDIA's repo : transpose both inputs to be BxHxWxC and then perform explicit correlation operation

I have not looked carefully in the backward operation yet, but it seems that the same kind of correlation operations will be done.

It seemed to me that Reda's technique was optimizable, especially with the transposition and all the multiplications that were done multiple times across different correlation operation, but I don't have any proof that this is not the best way to deal with that very peculiar module

[ClementPinard]

Here is a pytorch way of seeing how can things be done, displacement diameter is D, kernel size is K, padding and stride are ignored for clarity
Output is of size BxD^2xHxW but I think giving an output of size BxDxDxHxW would be semantically better if this module is to be used in other application than FlowNetC

1st solution :

d = (D-1)/2
for i in range(-d, d):
    for j in range(-d, d):
        output[:, i+ D*j, :, :] = torch.sum(input1 * input2[:, :, i:, j:], dim=1)  # [B, D^2, H, W]
final_output = F.avgpool2d(output, kernel_size = K, stride=1)  # padding ignored for clarity

2nd solution :

rinput1 = input1.transpose(0,2,3,1)  # [B, H, W, C]
rinput2 = input2.transpose(0,2,3,1)  # [B, H, W, C]

for b in range(B):
    for i in range(H):
        for j in range(W):
            for k in range(-d, d):
                for l in range(-d, d):
                    # rinput1[b, i:i+K, j:j+K] is a patch of size [K, K, C]
                    output[b, k + D*l, i, j] = torch.sum(rinput1[b, i:i+K, j:j+K] * rinput2[b, i+k:i+k+K, j+l:j+l+K])

[kingsleyluoxin]

Thank you so much for your detailed explanation! However, what indeed puzzles me is how we can properly arrange ways of indexing blocks and threads. That is, how we can understand the way author use to arrange threads for feature maps with shape BHWC or BCHW.

[ClementPinard]

As far as I understood, the key here is to operate on the most compact momery buffer possible. Here is a general summary of this idea : https://devblogs.nvidia.com/how-access-global-memory-efficiently-cuda-c-kernels/

So when you are doing the operation torch.sum(rinput1[b, i:i+K, j:j+K] * rinput2[b, i+k:i+k+K, j+l:j+l+K] in a cuda kernel, you need the strides of both patches to be the lowest possible.

When C is the last dimension, strides are W*C, C, 1 . When C is the second dimension, strides are H*W, W, 1
Now, depending on your number of channels, you might want to use one way of storing information or the other, but more generally, we have C < H, hence the transposition.

The weird thing however is that something similar to the CUDA_KERNEL_LOOP Macro is used (see here : https://devblogs.nvidia.com/cuda-pro-tip-write-flexible-kernels-grid-stride-loops/) but only considering the number of channels. Why not do that directly with the index ? Is there a CUDA master able to explain this ?

we might have to ask Fitsum Reda and Alexey Dosovitskiy for this as I believe Dosovitskiy wrote the first correlation kernel in caffe and Reda adapted it to pytorch, but without touching too much the inner mechanisms.

[kingsleyluoxin]

Thanks so much for your patient reply. I shall read and try to understand these codes with your explanation. :) 👍

[ClementPinard]

Here is a finished implementation, using the cpp-extension tools to make it installable.

https://github.com/ClementPinard/Pytorch-Correlation-extension

Note that this is still subject to change, as DeviceTensor is not an official interface. In the near future, it will most probably be replaced with what comes from #8366

[kingsleyluoxin]

Thank you so much for your implementation I shall learn it by heart. 发自我的小米手机 在 Clément Pinard <notifications@github.com>，2018年8月17日 下午6:25写道： Here is a finished implementation, using the cpp-extension tools to make it installable. https://github.com/ClementPinard/Pytorch-Correlation-extension Note that this is still subject to change, as DeviceTensor is not an official interface. In the near future, it will most probably be replaced with what comes from #8366<#8366> — You are receiving this because you commented. Reply to this email directly, view it on GitHub<#4073 (comment)>, or mute the thread<https://github.com/notifications/unsubscribe-auth/AZg-LaBe4U1NFOnwMmNjSLrokoga3A-jks5uRpoUgaJpZM4Q5yzx>.

Here is a pytorch way of seeing how can things be done, displacement diameter is D, kernel size is K, padding and stride are ignored for clarity
Output is of size BxD^2xHxW but I think giving an output of size BxDxDxHxW would be semantically better if this module is to be used in other application than FlowNetC

1st solution :

d = (D-1)/2
for i in range(-d, d):
    for j in range(-d, d):
        output[:, i+ D*j, :, :] = torch.sum(input1 * input2[:, :, i:, j:], dim=1)  # [B, D^2, H, W]
final_output = F.avgpool2d(output, kernel_size = K, stride=1)  # padding ignored for clarity

2nd solution :

rinput1 = input1.transpose(0,2,3,1)  # [B, H, W, C]
rinput2 = input2.transpose(0,2,3,1)  # [B, H, W, C]

for b in range(B):
    for i in range(H):
        for j in range(W):
            for k in range(-d, d):
                for l in range(-d, d):
                    # rinput1[b, i:i+K, j:j+K] is a patch of size [K, K, C]
                    output[b, k + D*l, i, j] = torch.sum(rinput1[b, i:i+K, j:j+K] * rinput2[b, i+k:i+k+K, j+l:j+l+K])

How much slower would this solution be, compared to the cuda code in Pytorch-Correlation-extension?

@ClementPinard

[ClementPinard]

That would be interesting to test it with JIT tracing ! I haven't tested it, but there's a benchmark script in the Pytorch-Correlaton-extension repo that you can try to adapt to this python code.

My guess is that 1st will be slightly slower, but nothing too serious : I tried this solution in c++ with torch cpp and it was indeed a bit slower that the current implementation. However, 2nd solution has several parrallel computing tricks that makes it the fastest and it would be impossible to code it simply directly in pytorch, even with JIT tracing, so i would guess that it would be painfully slower (5 intricated for loops in python is hardly ever a good idea !)

So my pronostic, from fastest to slowest : 2nd solution (cpp/cuda) -> 1st solution (cpp/cuda) -> (by a short margin) 1st solution (python) -> 2nd solution (python)