Are there any examples on running a GNN model through ONNX?

[Wantcha]

Ask a Question

Question

I'm using Pytorch-Geometric to create a GNN and I'd like to be able to run it inside an application that supports ONNX inference, but I can't seem to find any information on how or whether this is possible. How do I export a PyG GNN model and how can I run inference on it?

[M-Polonczyk]

Managed to save the model to ONNX format. Threw a lot of errors, but saved. Don't know if it works properly as I only tested it with a single straightforward example.

model, best_val_acc = train(
    train_loader=train_loader,
    val_loader=val_loader,
    model_config=model_params["gcn_standard"],
    num_epochs=training_config["epochs"],
    learning_rate=training_config["learning_rate"],
    device=device,
)
print(f"Best validation accuracy: {best_val_acc:.4f}")

# Save the trained model
torch.save(model.state_dict(), "classifier_model.pth")

input_dim = next(iter(train_loader)).x.shape[1]
num_nodes = 10
x = torch.randn(num_nodes, input_dim).to(device)
batch = torch.zeros(num_nodes, dtype=torch.long).to(device)
row = torch.arange(num_nodes).repeat(num_nodes)
col = torch.arange(num_nodes).unsqueeze(1).repeat(1, num_nodes).flatten()
edge_index = torch.stack([row, col], dim=0).to(device)

# For some reason it doesn't work on GPU
model.to("cpu")
x = x.to("cpu")
edge_index = edge_index.to("cpu")
batch = batch.to("cpu")

onnx_program = torch.onnx.export(
    model,
    (x, edge_index, batch),
    dynamo=True,
    input_names=["x", "edge_index", "batch"],
    output_names=["output"],
    export_params=True,
)
onnx_program.save("classifier_model.onnx")

Model:

import torch
import torch.nn.functional as F
from torch import nn
from torch_geometric.nn import (
    BatchNorm,
    GCNConv,
    global_add_pool,
    global_max_pool,
    global_mean_pool,
)

from .base import BaseGNN


class GCN(BaseGNN):
    def __init__(
        self,
        input_dim,
        hidden_dim=128,
        num_classes=2,
        num_layers=3,
        dropout_rate=0.3,
        use_batch_norm=True,
        pool_type="mean",
    ) -> None:
        super().__init__(
            input_dim,
            hidden_dim,
            num_classes,
            num_layers,
            dropout_rate,
        )

        self.hidden_dim = hidden_dim
        self.dropout_rate = dropout_rate
        self.use_batch_norm = use_batch_norm
        self.pool_type = pool_type

        # GCN layers
        self.convs = nn.ModuleList()
        if use_batch_norm:
            self.batch_norms = nn.ModuleList()

        # First layer
        self.convs.append(GCNConv(input_dim, hidden_dim))
        if use_batch_norm:
            self.batch_norms.append(BatchNorm(hidden_dim))

        # Hidden layers
        for _ in range(num_layers - 2):
            self.convs.append(GCNConv(hidden_dim, hidden_dim))
            if use_batch_norm:
                self.batch_norms.append(BatchNorm(hidden_dim))

        # Final conv layer
        if num_layers > 1:
            self.convs.append(GCNConv(hidden_dim, hidden_dim))
            if use_batch_norm:
                self.batch_norms.append(BatchNorm(hidden_dim))

        # Classifier head
        if pool_type == "combined":
            classifier_input_dim = hidden_dim * 2
        else:
            classifier_input_dim = hidden_dim

        self.classifier = nn.Sequential(
            nn.Linear(classifier_input_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(dropout_rate),
            nn.Linear(hidden_dim // 2, num_classes),
        )

    def forward(self, x, edge_index, batch):

        if self.use_batch_norm and len(self.batch_norms) > 0:
            for conv, bn in zip(self.convs, self.batch_norms, strict=False):
                x = conv(x=x, edge_index=edge_index)
                x = bn(x)
                x = F.relu(x)
                x = F.dropout(x, p=self.dropout_rate, training=self.training)
            if len(self.convs) > len(self.batch_norms):
                for i, cconv in enumerate(self.convs):
                    if i >= len(self.batch_norms):
                        x = cconv(x=x, edge_index=edge_index)
                        x = F.relu(x)
                        x = F.dropout(x, p=self.dropout_rate, training=self.training)
        else:
            for conv in self.convs:
                x = conv(x=x, edge_index=edge_index)
                x = F.relu(x)
                x = F.dropout(x, p=self.dropout_rate, training=self.training)
                
        # Graph-level pooling
        if self.pool_type == "mean":
            x = global_mean_pool(x, batch)
        elif self.pool_type == "max":
            x = global_max_pool(x, batch)
        elif self.pool_type == "add":
            x = global_add_pool(x, batch)
        else:  # combined pooling
            x_mean = global_mean_pool(x, batch)
            x_max = global_max_pool(x, batch)
            x = torch.cat([x_mean, x_max], dim=1)

        return self.classifier(x)

Here is how I ran onnx model from Go code:

package main

import (
    "fmt"
    ort "github.com/yalue/onnxruntime_go"
)

func main() {
    // This line _may_ be optional; by default the library will try to load
    // "onnxruntime.dll" on Windows, and "onnxruntime.so" on any other system.
    // For stability, it is probably a good idea to always set this explicitly.
    ort.SetSharedLibraryPath("lib/libonnxruntime.so")

    err := ort.InitializeEnvironment()
    if err != nil {
        panic(err)
    }
    defer ort.DestroyEnvironment()

    // For a slight performance boost and convenience when re-using existing
    // tensors, this library expects the user to create all input and output
    // tensors prior to creating the session. If this isn't ideal for your use
    // case, see the DynamicAdvancedSession type in the documnentation, which
    // allows input and output tensors to be specified when calling Run()
    // rather than when initializing a session.
    inputData := []float32{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9}
    inputShape := ort.NewShape(2, 5)
    inputTensor, err := ort.NewTensor(inputShape, inputData)
    defer inputTensor.Destroy()
    // This hypothetical network maps a 2x5 input -> 2x3x4 output.
    outputShape := ort.NewShape(2, 3, 4)
    outputTensor, err := ort.NewEmptyTensor[float32](outputShape)
    defer outputTensor.Destroy()

    session, err := ort.NewAdvancedSession("classifier_model.onnx",
        []string{"Input 1 Name"}, []string{"Output 1 Name"},
        []ort.Value{inputTensor}, []ort.Value{outputTensor}, nil)
    defer session.Destroy()

    // Calling Run() will run the network, reading the current contents of the
    // input tensors and modifying the contents of the output tensors.
    err = session.Run()

    // Get a slice view of the output tensor's data.
    outputData := outputTensor.GetData()
    fmt.Println("Model loaded!", outputData)
}

However it seems to have many flaws, especially for pytorch geometric library. Resources what I found:

• https://indico.cern.ch/event/1071461/contributions/4505362/attachments/2323544/3957178/IRIS-HEP_Submit.pdf
• SAGEConv ONNX export pytorch/pytorch#64769
• GNN (SAGE Convolutional) ONNX export pytorch/pytorch#65138
• torch.onnx support pyg-team/pytorch_geometric#728