radius_graph with CUDA

[PabloVD]

Hi, I'm trying to understand what radius_graph is actually doing to obtain the list of edges.
I see that, when the cpu is used, it is just building a KDTree using the scipy implementation. However, from the documentation and the torch_cluster repository, it is not obvious to me which algorithm is chosen when CUDA is employed.
Which is the expected time complexity in that case, how does it scale with the number of nodes considered?
Many thanks in advance!

[rusty1s]

The CUDA implementation will not build a KDTree. It will leverage the "power of parallel processors" to check if every element is within radius of the central node. As such, it's non-parallel runtime complexity will be O(N^2). Given N parallel processors, it runtime is reduced to O(N).

Edit: Here is the link to the code which sits in torch-cluster.

[PabloVD]

Sure, I find that behavior with the following script:

import torch
import time
from torch_cluster import radius_graph


# use GPUs if available
if torch.cuda.is_available():
    print("CUDA Available")
    device = torch.device('cuda')
else:
    print('CUDA Not Available')
    device = torch.device('cpu')

torch.manual_seed(12345)

num_points = 10000

points = torch.randn([num_points,3])
radius = 0.8
loop = False
max_num_neighbors = points.shape[0]-1

points.to(device)

time_start = time.time()

edge_index = radius_graph(points, r=radius, loop=loop, max_num_neighbors=max_num_neighbors)

time_infer = time.time()-time_start
print("Num points", num_points, "time",time_infer)

For different number of points I get the next results:
Num points 1000 time 0.033
Num points 10000 time 0.298
Num points 100000 time 43.408

Changing max_num_neighbors reduce the time, but the relative time scaling is the same.

I also found the same time scaling with other point clouds, not only these random distributions.

I am using torch 1.10.0+cu113 and torch_cluster 1.5.9

[rusty1s]

Thank you. In the meantime, I just released torch-cluster==1.6.0.

[PabloVD]

So great, thanks!

I just ran my script with that new version but I do not find a significant improvement of the time taken to form the graph (only a gain of <~ 1 ms for 1k and 10k particles). Should we expect a faster performance?

[rusty1s]

It looks like your script is somewhat wrong as you are benchmarking CPU performance since you never assigned points.to(device). Also keep in mind to synchronize CUDA when benchmarking. Furthermore, setting max_num_neighbors to the size of the point cloud can heavely slow down your code and should be used with care. Here is the corrected script:

import torch
import time
from torch_cluster import radius_graph


# use GPUs if available
if torch.cuda.is_available():
    print("CUDA Available")
    device = torch.device('cuda')
else:
    print('CUDA Not Available')
    device = torch.device('cpu')

torch.manual_seed(12345)

num_points = 10000

points = torch.randn([num_points,3])
radius = 0.8
loop = False
max_num_neighbors = points.shape[0]-1

points.to(device)

time_start = time.time()

edge_index = radius_graph(points, r=radius, loop=loop, max_num_neighbors=max_num_neighbors)

time_infer = time.time()-time_start
print("Num points", num_points, "time",time_infer)

This gives me:

  1000: 0.0012s
 10000: 0.013s
100000: 0.335s

[PabloVD]

It seems that points.to(device) is already in my snippet above, although I noticed that running points.is_cuda gives me False. I need to write points=points.to(device) to actually having it on cuda. With that, now I get faster results than before with 10000 points, although slower than yours:
Num points 10000 time 0.045
and the same time for 1000. I don't know why this is not so fast, may it has to do with the cuda/torch version?

By the way, using torch.cuda.Event(enable_timing=True) and torch.cuda.synchronize() I got the same results when computing the time.

Thank you again!