MessagePassing.jittable() remaining features and planning.

This is an issue to track followup on #1256, #1257, #1258, #1259.

We can annotate below with further issue and PR #s, as progress is made. Let’s use this issue as a forum for discussion about what we want out of the jittable() interface.

During review there were a few functionality issues, or requests that came up (box checked if completed, with PR reference):

• - return a jittable copy of the conv layer instead of the new class (https://github.com/rusty1s/pytorch_geometric/commit/dc87faa6a9ae1b71dbbd6d34f834dc7053803fe6)
• - deal with overloads in a more graceful way
• - add a utility that generates model config files for deployment platforms.

Here’s a list of jittable convolutional ops (box is checked if tested and confirmed):

• Basic test in test/nn/conv/test_message_passing.py
• AGNNConv
• APPNP
• ARMAConv
• CGConv
• ChebConv
• DNAConv
• EdgeConv
• DynamicEdgeConv - #1366
• FeaStConv
• GATConv
• GatedGraphConv
• GCNConv
• GINConv
• GINEConv
• GMMConv
• GraphConv
• GravNetConv - #1366
• HypergraphConv - NB: needs overall refactoring
• NNConv
• PointConv
• PPFConv
• RGCNConv - NB: this needs to wait for better handling of multiple propagate types
• SAGEConv
• SGConv
• SignedConv
• SplineConv
• TAGConv
• XConv

[SonyPony]

I think this depends on torch.scatter_add being supported by ONNX or not (which should give ONNX support for all PyG operators with add or mean aggregation). I haven’t done any experiments with that either, so I’m looking forward to your findings.

@rusty1s I think I’m stuck by another point which is mini-batch inference… In my understanding, PyG processes mini-batch as a “giant graph that holds multiple isolated subgraphs”, and supports arbitrary numbers of nodes for each sample. Since my final destination is TensorRT, batches in PyG may be treated as single samples with highly size variation. This may cause weird problems in TensorRT inference.

Nevertheless, I have successfully tried converting some simple model, including an example message passing model, to ONNX model. Seems that the operations are supported by ONNX at least. Furthermore, I think the TensorRT inference is possible if the data is highly structured. For example, all samples have same number of nodes, the connections of nodes are same, and only the features of nodes are different, thus we can organize the data as traditional mini-batch form.

Don’t know if I understand it right.

@dongb5 check it the prediction is correct. I’ve been struggling with exporting the GN model to ONNX as well. Although the torch.onnx.export runs without an error, the prediction is not correct due to the bad translation of torch.scatter_add to ONNX graph (https://github.com/pytorch/pytorch/issues/32960) if the indices are not unique. Despite being currently supported by ONNX in opset 16 (https://github.com/onnx/onnx/pull/3484), PyTorch currently supports only opset 14.

[chriss2401]

Hi @rusty1s and @lgray ,

This new jit functionality is great! Do you guys also plan onnx support ? As far as I have seen there are issues converting a jittable model to onnx (https://github.com/pytorch/pytorch/issues/34002), and ideally I’d like to export a pytorch_geometric model to onnx.