Speed issues

[de-gozaru]

Hi @rusty1s ,

I was comparing the speed of torch_cluster.knn_graph on the GPU with the function bellow:

def knn(x, k):
    inner = -2 * torch.matmul(x.transpose(2, 1), x)
    xx = torch.sum(x**2, dim=1, keepdim=True)
    pairwise_distance = -xx - inner - xx.transpose(2, 1)

    idx = pairwise_distance.topk(k=k, dim=-1)[1]   # (batch_size, num_points, k)
    return idx

and found that torch_cluster.knn_graph is way slower, around x8 slower, is this something expected?

Thank you in advance!

[rusty1s]

knn_graph parallelizes over the batch dimension. What's the batch_size you are using? Also note that you are trading memory with runtime in your implementation. knn_graph computes distances on the fly and therefore does not need to compute the intermediate O(N^2) matrix - but it is also slower in return.

[de-gozaru]

Hi @rusty1s
I'm using small batches actually!
Also, the fact that trading memory with speed answers my question!

Thank you!