Could pytorch provide correlation operator?

Jiayuan_Li: <pre><code class="lang-auto">type or paste code here
</code></pre>
<p>Hi, I am working on optical flow projects, and most of the methods utilizes the correlation of two feature maps to get the similarity. However,  some researchers implement the CUDA programming on correlation, like FlowNet and PWC-net; while others calculate all pairs correlation like, (RAFT, <a href="https://arxiv.org/pdf/2003.12039.pdf" rel="nofollow noopener">https://arxiv.org/pdf/2003.12039.pdf</a>). Could Pytorch support a correlation operator for us? Or can we easily implement the method using existing operators.</p>
<p>The calculation is straightforward.<br>
cost_volume = corr(fmap1, fmap2, search_range)<br>
given fmap1[B, H, W, C] and fmap2[B, H, W, C], search_range=3 (a radius of search range)</p>
<p>The cost_volume should be [B, H, W, 49]. which stores the correlation value between each feature in fmap1 and its corresponding n</p>

ptrblck: <p>You could use <code>F.conv2d</code> and just flip the filter kernel to get a cross-correlation.<br>
Would that work or are you looking for another specific implementation?</p>

Jiayuan_Li: <p>hhh…I see. It’s not the basic math calculation, but a way to calculate  the cost volume between feature maps. Conv2D only applies a window on a image but cost volume is from both feature map</p>

ptrblck: <p>How is this cost volume operation defined?<br>
Could you give an example using two random feature maps?</p>

Jiayuan_Li: <p>Here is an example in PWC-Net. <a href="https://youtu.be/LBJ20kxr1a0?t=775" rel="nofollow noopener">https://youtu.be/LBJ20kxr1a0?t=775</a><br>
The main idea is to calculate the correlation as dot product between two features. Larger result indicates more similar.<br>
Meanwhile, the correlation could be between each feature in map1 and all features in map2; and each feature in map1 and its corresponding neighbor in map2.</p>

ptrblck: <p>It’s a bit unclear from the video, what the scalar output should represent.<br>
A correlation between two signals would result in another signal, so are the authors using the peak of this correlation or another metric/coefficient?</p>

Jiayuan_Li: <p>Hi,<br>
Here is the definition in the PWC-net<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/6/8/68c3fa878103101a1f487aaf39a522cc395184be.png" alt="image" data-base62-sha1="eWNxqHIfEdX5OIYwQG3OJyeV6LI" width="618" height="302"></p>
<p>We can think of two signals you mentioned as two feature maps from CNN, let’s say fmap1 [48, 128, 64] and fmap2 [48, 128, 64] where [H, W, C]. Then “result in another signal” should be the cost volume.</p>
<p>The way to calculate the cost volume should be like this:<br>
For each feature in fmap1 [i, j, 128], we find its corresponding neighbor in fmap2: [i-3, j-3, 128] [i-2,j-3, 128]…[i, j, 128]…[i+3, j+3, 128]. Then we do dot product between feature1 [i, j,  128] between each feature2 then we got N correlation values where N is the number of pixels located in neighbor, and we cat those scalars we get [1, 1, N] correlation feature. Finally loop over all pixels in fmap1 then get a cost volume as [48, 128, N] which stores the correlation value and represents similarity between fmap1 and fmap2</p>

ptrblck: <p>Based on your description this would be similar to a convolution (or correlation, if you flip one kernel), which would be possible to apply using <code>F.conv2d</code> and two inputs.<br>
However, I’m a bit confused by the shapes.<br>
If we scale down the problem a bit and assume that both feature maps have a single channel, the correlation output would have the shape <code>[h1 + 2*h2 -2, w1 + 2*w2 -2]</code>, of you perform a full correlation, wouldn’t it? These output values seem to correspond to the <code>N</code> correlation values you mentioned. In the last step you mention that you would repeat this step over all pixels, whihc is unclear to me.<br>
Could you post a pseudo code or a dummy example in PyTorch or numpy?</p>

Jiayuan_Li: <p>This link from PWC-net might be helpful: <a href="https://github.com/deqings/PWC-Net/issues/1#issuecomment-404290729" rel="nofollow noopener">https://github.com/deqings/PWC-Net/issues/1#issuecomment-404290729</a></p>
<p>As we know that the correlation is a flipped convolution in pure math but here it doesn’t function like that style. Conv2D receives an input feature and then convs it by a kernel defined by user. However, you see that cost volume involves two different feature maps. In this case, all research publication in optical flow needs to implement CUDA programming to do such “correlation”. Like: FlowNet,  FlowNet2, PWC-net. If pytorch is able to provide a official Correlation or CostVolume API, it would be great for both research and industry.</p>
<p>Here is the CUDA and python code from PWC-net.<br>
</p><aside class="onebox whitelistedgeneric">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">
      <a href="https://github.com/NVlabs/PWC-Net/tree/master/PyTorch/external_packages/correlation-pytorch-master/correlation-pytorch/correlation_package" target="_blank" rel="nofollow noopener">GitHub</a>
  </header>
  <article class="onebox-body">
    <img src="https://avatars2.githubusercontent.com/u/2695301?s=400&amp;v=4" class="thumbnail onebox-avatar" width="400" height="400">

<h3><a href="https://github.com/NVlabs/PWC-Net/tree/master/PyTorch/external_packages/correlation-pytorch-master/correlation-pytorch/correlation_package" target="_blank" rel="nofollow noopener">NVlabs/PWC-Net</a></h3>

<p>PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume, CVPR 2018 (Oral) - NVlabs/PWC-Net</p>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


Jiayuan_Li: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/5/95fa54353e22c6c8f504715f4943e28e6d4908f0.png" data-download-href="https://discuss.pytorch.org/uploads/default/95fa54353e22c6c8f504715f4943e28e6d4908f0" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/5/95fa54353e22c6c8f504715f4943e28e6d4908f0.png" alt="image" data-base62-sha1="loLta6LZ43J4KI8zyXZQQ5TQmL6" width="690" height="91" data-dominant-color="95C5FF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">833×110 4.87 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
As Deqin Sun mentioned, two different cost volumes exist.</p>
<p>I write a naive pseudo code</p>
<pre><code class="lang-auto">fmap1 = torch.ones(C, H, W) # [C, H, W]
fmap2 = torch.ones(C, H, W)


# Full cost volume (all-pairs correlation):
cost_vol = torch.ones(H, W, 1, H, W)
for i in range(H):
	for j in range(W):
		vec1 = fmap1(:, i, j)
		corr_ij = dot(vec1.T, fmap2.view(:, -1)) # [1, C] *[C, H*W] = [1, H*W]
		cost_vol(i, j, 0, corr_ij.view(1, H, W))

cost_vol = cost_vol.view(H, W, H*W)

# Partial cost volume
r = 3
dx = linespace(-r, r, 2*r+1)
dy = linespace(-r, r, 2*r+1)
grid = stack(meshgrid(dy, dx), axis=-1)

cost_vol = torch.ones(H, W, (2*r+1)**2)
for i in range(H):
	for j in range(W):
		vec1 = fmap1(:, i, j)
		fmap2_neighbor = grid_sample(fmap2, grid+[i, j]) # sample a sqare fmap2_neighbor at (i, j) in fmap2 with size(2*r+1, 2*r+1)
		corr_ij = dot(vec1.T, fmap2_neighbor.view(:, -1)) # [1, C] *[C, (2*r+1)**2] = [1, 49] 
		cost_vol(i, j, :) = corr_ij

# the first and second channel of final cost volume should be equal to the fmap1 shape. 
# We don't need to care about the order of rest of channels but the product result of the number of other channels indicates that how many feature in fmap2 we calculate the correlation with each feature in fmap1. EX: [H, W, H*W] means all-pairs correlation known as each feature in fmap1 is correlated with each feature in fmap2;   [H, W, 49] means each feature in fmap1 is correlated with a [7, 7] in fmap2.       


</code></pre>

yangyi02: <p>Also need the built-in correlation function.</p>
<p>Due to some reason (gcc =9.3.0, cuda=10.1, and I do not have permission to change them), I cannot compile the existing pytorch <a href="https://github.com/vt-vl-lab/pwc-net.pytorch/tree/master/PWC_src/correlation_package" rel="nofollow noopener">correlation_package</a>.</p>
<p>It will be very helpful if there is such a built-in function.</p>
<p>Thank you.</p>

Ruslan_Baynazarov: <p>Would be great to have this layer added to pytorch.</p>
<p>I’m using Clement Pinard’s implementation: <a href="https://github.com/ClementPinard/Pytorch-Correlation-extension" class="inline-onebox" rel="noopener nofollow ugc">GitHub - ClementPinard/Pytorch-Correlation-extension: Custom implementation of Corrleation Module</a><br>
It works with torch 1.7.0 and cuda 11.1.<br>
However it needs to be compiled from source, which lowers the reproducibility of my code for other developers.</p>
<p>Correlation layer is essential in modern flow architectures (FlowNetC, PWCNet, MaskFlowNet).<br>
Adding this layer to pytorch will also speed up adding it to the runtime frameworks (ONNX, openvino) so that we are able to compute optical flow on edge devices.</p>

Etienne_Perot: <p>would this correspond in pure pytorch? it computes inner product between the central pixel to every pixels in the surrounding:</p>
<pre><code class="lang-auto">import torch
from einops import rearrange
unfold_op = torch.nn.Unfold(k1, dilation=1, padding=k1//2, stride=1)

def pixel_cost_volume(im1, im2, unfold_op):
    """im1/2 b c h w

    -&gt; produces b h w k**2 cost volume, each channel represents
    inner product with a different discrete shift (i,j)
    """
    b,c,h,w = im1.shape
    central_pixel = rearrange(im1, 'b c h w -&gt; b h w c')
    neighbors = unfold_op(im2).reshape(b,-1,h,w)
    neighbors = rearrange(neighbors, 'b (c k2) h w -&gt; b h w k2 c', c=c)
    y = torch.einsum('...i c, ...c -&gt; ...i', neighbors, central_pixel)
    return y # (b, h, w, search_range x search_range)
</code></pre>

Diego_Paez: <p>A native pytorch implementation would be great. Given two same sized images, this layer is similar to conv2D layer between image1 and a padded image2. But there are relevant differences: firstly there are no weights, image2 takes the place of the filters, secondly the multiplications are not added together only the channels are averaged, thirdly each multiplication between each of the images is stored in a different output channel. Take a look at:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/limacv/CorrelationLayer">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/limacv/CorrelationLayer" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/b907e53a40d4a91d9b77b47ad727dd02f45a54bfb84fe21512c2ee80cf292149/limacv/CorrelationLayer" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/limacv/CorrelationLayer" target="_blank" rel="noopener nofollow ugc">GitHub - limacv/CorrelationLayer: Pure Pytorch implementation of Correlation...</a></h3>

  <p>Pure Pytorch implementation of Correlation Layer that commonly used in learning based optical flow estimator - GitHub - limacv/CorrelationLayer: Pure Pytorch implementation of Correlation Layer tha...</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>This layer is very useful for optical flow.</p>

jondo: <p>To which extent is this topic a duplicate of <a href="https://discuss.pytorch.org/t/request-for-correlation-layer/10227" class="inline-onebox">Request for correlation layer</a>?</p>

shakxy42: <p>Has this layer been implemented in PyTorch as of today? I see that RAFT is available as a model in <a href="https://pytorch.org/vision/main/auto_examples/plot_optical_flow.html" rel="noopener nofollow ugc">torchvision</a><br>
RAFT uses the correlation layer to compute a 4D cost volume. The optical flow community would really appreciate if this layer is added to PyTorch</p>

Avishka_Perera: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Are there any Pytorch native implementations yet? Or is there any custom implementation Pytorch would recommend? I am currently trying to compile it from this repo: <a href="https://github.com/NVlabs/PWC-Net" class="inline-onebox" rel="noopener nofollow ugc">GitHub - NVlabs/PWC-Net: PWC-Net: CNNs for Optical Flow Using Pyramid, Warping, and Cost Volume, CVPR 2018 (Oral)</a><br>
Unfortunately, it was developed for a very old version and I am not able to use new Cuda versions or Pytorch versions</p>

vadimkantorov: <p>Might be related torch.sampled_addmm (if we explicitly compute the neighborhood indices + don’t know if it’s good perfwise, as maybe a 2d/3d with explicit local neighborhoods allow for better perf than the generic case): <a href="https://github.com/pytorch/pytorch/issues/105319#issuecomment-1709700269" class="inline-onebox" rel="noopener nofollow ugc">Implementation of torch.sparse.sampled_baddmm · Issue #105319 · pytorch/pytorch · GitHub</a></p>

vadimkantorov: <p>Also, please feel free to chime in <a href="https://github.com/getkeops/keops/issues/148" class="inline-onebox" rel="noopener nofollow ugc">[question] Local video correlation (with temporal context length &gt; 1) · Issue #148 · getkeops/keops · GitHub</a> <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> I think it might be possible to accomplish this in keops</p>