Explanation of torch_cluster knn

username: mvillanuevaaylagas
cooked: <p>Hello, I’m trying to compute a batched version of KNN. At the moment I’m looping over scipy’s cKDTree.</p>
<p>I saw that PyTorch geometric has a GPU implementation of <a href="https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#torch_geometric.nn.pool.knn" rel="noopener nofollow ugc">KNN</a>. However, I find that the documentation is not very clear the <code>x</code> and <code>y</code> input variables are matrices of points times features. Now, if I run the example code</p>
<pre><code class="lang-auto">x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch_x = torch.tensor([0, 0, 0, 0])
y = torch.Tensor([[-1, 0], [1, 0]])
batch_y = torch.tensor([0, 0])
assign_index = knn(x, y, 2, batch_x, batch_y)
</code></pre>
<p>the output of the example code is a 2x4 tensor. I was expecting a 2x2, e.g. for each point in y, its two nearest neighbors.<br>
Can someone explaint what is this output?</p>
<p>Also, if I want to run this with <code>x</code> and <code>y</code> being batched vectors (BxNxF), would I have to flatten the tensors and create a kind of one-hot-encoding <code>batch_x</code> that indicates the batch?</p>
<p>Thanks for your help</p>

username: ptrblck
cooked: <p>I’m not completely sure as I haven’t used this method before, but based on the output:</p>
<pre data-code-wrap="python"><code class="lang-python">print(assign_index)
# tensor([[0, 0, 1, 1],
#         [0, 1, 2, 3]])
</code></pre>
<p>I would assume that the first row gives the indices of <code>y</code> while the second row gives the indices of the nearest neighbors in <code>x</code> for the <code>y</code> index.</p>
<p>Let’s change the example a bit and see if this is true.</p>
<pre data-code-wrap="python"><code class="lang-python"># make perfect matches
x = torch.Tensor([[-1, 0], [1, 0], [1, -1], [-1, 0]])
batch_x = torch.tensor([0, 0, 0, 0])
y = torch.Tensor([[-1, 0], [1, 0]]) # first sample should match 0 and 3, second should match 1 and 2
batch_y = torch.tensor([0, 0])
assign_index = knn(x, y, 2, batch_x, batch_y)
print(assign_index)
# tensor([[0, 0, 1, 1],
#         [0, 3, 1, 2]])
</code></pre>
<p>Based on this I guess my assumption is correct.</p>
<aside class="quote no-group" data-username="mvillanuevaaylagas" data-post="1" data-topic="148245">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mvillanuevaaylagas/48/48272_2.png" class="avatar"> mvillanuevaaylagas:</div>
<blockquote>
<p>Also, if I want to run this with <code>x</code> and <code>y</code> being batched vectors (BxNxF),</p>
</blockquote>
</aside>
<p>Isn’t <code>N</code> already defining the batch dimension?</p>
<aside class="quote no-group" data-username="mvillanuevaaylagas" data-post="1" data-topic="148245">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/mvillanuevaaylagas/48/48272_2.png" class="avatar"> mvillanuevaaylagas:</div>
<blockquote>
<p>and create a kind of one-hot-encoding <code>batch_x</code> that indicates the batch?</p>
</blockquote>
</aside>
<p>I don’t know how <code>batch_x</code> and <code>batch_y</code> are used and don’t fully understand the docs, but again I’m not familiar enough with <code>pytorch_geometric</code> so let’s wait for some experts.</p>

username: mvillanuevaaylagas
cooked: <p>Thanks for the explanation, your example makes sense.</p>
<p>Let me elaborate on the batch part of the question while we wait for the experts <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"><br>
Imagine you have a cloud of 2D (F) 10 points (N) <code>x</code> and some 5 points <code>y</code> for which you want to find the NN at time <code>t</code>. Now, the points in <code>x</code> can change over time and you need to find the NN of another set of <code>y</code> points.</p>
<p>So, you want to limit the search of <code>y_t</code> to <code>x_t</code>. If you append all batches of <code>x</code> in an NxF matrix, you would be searching over different timestamps.</p>

username: ptrblck
cooked: <p>Ah OK, thanks for the explanation.<br>
If I thus want to combine both examples into a single operation I could just assign my new example to timestep 1?<br>
This output seems to show that it’s working fine:</p>
<pre><code class="lang-python">x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1], [-1, 0], [1, 0], [1, -1], [-1, 0]])
batch_x = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1])
y = torch.Tensor([[-1, 0], [1, 0], [-1, 0], [1, 0]])
batch_y = torch.tensor([0, 0, 1, 1])
assign_index = knn(x, y, 2, batch_x, batch_y)
print(assign_index)
# tensor([[0, 0, 1, 1, 2, 2, 3, 3],
#         [0, 1, 2, 3, 4, 7, 5, 6]])
# equal to previous answer with offset where offsets are
# tensor([[0, 0, 0, 0, 2, 2, 2, 2],
#         [0, 0, 0, 0, 4, 4, 4, 4]])
</code></pre>
<p>The “second half” of the answer is matching the result of my previous code snippet with the offset mentioned below.</p>