K-NNGraph: make k variable for batched examples

[RasmusOrsoe]

Hey!

I use torch_geometric.transforms.KNNGraph quite often, and on data where the number of nodes in a graph can vary quite a bit. Having a fixed k for all examples in a batch is suboptimal. In terms of optimization, it could be very nice to specify a k for each example in a batch of graphs. This could be achieved by making KNNGraph accept a list as argument for k.

Mock-example:
Let's say we have batch_size = 5, so 5 graphs in a batch, and that the number of nodes in the graphs are: [10, 4, 30, 50, 8].

Instead of specifying a fixed k for the entire batch, one could pass a list of k's. In this case, perhaps something like k = [3, 2, 8, 10, 3] such that the edge index for the first graph is calculated using k = 3, and k=2 for the second graph, k = 8 for the third, and so on.

I can write a dirty for-loop to achieve this, but it's highly inefficient. Any ideas on how to implement this?

Thanks

[rusty1s]

I think this is a good feature request, thanks for sharing. This would need to be implemented in torch-cluster, both on CPU and GPU side. I cannot make this top prio on my end, but happy to take in any contributions. Let me know if you are interested :)

[RasmusOrsoe]

I had a look at the source code in torch-cluster, and I don't think I can produce an elegant fix for this. This is my best attempt:

def knn_graph_batch(batch: Batch, k: List[int]):
    """Calculates the k-nearest-neighbours with an individual k for each event in batch.

    Args:
        batch (Batch): A torch_geometric.data.Batch of events
        k (List[int]): A list of k's

    Returns:
        Batch: returns the same batch of events, but with updated edges.
    """
    data_list = batch.to_data_list()
    for i in range(len(data_list)):
        data_list[i].edge_index = knn_graph(x=data_list[i].x[:, 0:3], k=k[i])
    return Batch.from_data_list(data_list)

If anyone has ideas on how to do this more elegantly, I'd like to see it!

[rusty1s]

Ideally, we would integrate this into the C++/CUDA implementation of torch-cluster (although the presented solution obviously works as well :)).