Requesting ONNX equivalent to torch_cluster::knn
New Operator
Describe the operator
Requesting an ONNX equivalent to torch_cluster::knn for nearest neighbor computations at runtime. This operator is required for exporting PyTorch Geometric models such as GravNetConv, DynamicEdgeConv, and other graph layers that dynamically build edges based on spatial proximity. Currently we cannot export these models to ONNX.
Some context: I’m a physics grad student researcher with the IceCube Neutrino Observatory, and our collaboration has recently decided to use ONNX as a standard format. Lack of KNN blocks export of many geometric deep learning models, which makes it impractical for us to work with them and slows research.
Can this operator be constructed using existing onnx operators?
Not that I could find.
Is this operator used by any model currently? Which one?
Used by multiple PyG models (GravNetConv, DynamicEdgeConv, etc).
Are you willing to contribute it? (Y/N)
I can assist but I don't have bandwidth to fully implement it myself.
Notes
Here is a minimal working example of the export failing:
import torch
import torch.nn as nn
import torch.onnx as onnx
from torch_geometric.nn import GravNetConv, global_mean_pool
from torch_geometric.nn.aggr import MeanAggregation
in_ch = 16
hidden = 32
out_ch = 5
# set up a simple model w/ GravNetConv layers
class GNNet(nn.Module):
def __init__(self):
super().__init__()
P = 16
self.conv = GravNetConv(in_ch, hidden, space_dimensions=4, propagate_dimensions=P, k=8)
self.act = nn.ReLU()
self.head = nn.Linear(hidden, out_ch)
# have to use aggr="mean" or it crashes
self.conv.aggr_module = MeanAggregation()
self.conv.lin_out2 = nn.Linear(P, self.conv.lin_out2.out_channels, bias=True).to(self.conv.lin_out2.weight)
def forward(self, x, batch):
h = self.conv(x, batch)
h = self.act(h)
return self.head(global_mean_pool(h, batch))
total_nodes = 128
batch_size = 4
# DUMMY BATCH
# random node features
x = torch.randn(total_nodes, in_ch)
nodes_per_graph = total_nodes // batch_size
remainder = total_nodes % batch_size
sizes = torch.full((batch_size,), nodes_per_graph, dtype=torch.long)
sizes[:remainder] += 1
# build the batch
graph_ids = torch.arange(batch_size, dtype=torch.long)
batch = torch.repeat_interleave(graph_ids, sizes)
# evaluate
model = GNNet().eval()
# export not supported and fails with
# torch.onnx.errors.UnsupportedOperatorError: ONNX export failed on an operator with unrecognized namespace torch_cluster::knn.
onnx.export(
model, (x, batch), "gravnetconv.onnx",
input_names=["x", "batch"], output_names=["logits"], opset_version=17,
dynamic_axes={"x": {0: "num_nodes"}, "batch": {0: "num_nodes"}, "logits": {0: "num_graphs"}},
)Source: onnx/onnx