docs/export.md

ROOT-GNN ONNX export

ONNX export is an inference-only adapter for prepared ROOT-GNN graph batches. It does not read ROOT files or move feature construction into the deployment model. Install the optional dependencies with:

uv sync --extra root-gnn --extra onnx

Export a checkpoint with the project CLI:

uv run gnn4colliders export \
  export.checkpoint=/path/to/checkpoint.pt \
  export.output=model.onnx \
  data.cache.path=/path/to/graph-cache.pt

The exported model accepts six tensor inputs: node_features, edge_features, edge_src, edge_dst, node_batch, and global_features. node_batch identifies the graph for each node; edge membership is derived from node_batch[edge_src]. Graph, node, edge, and batch dimensions are dynamic. The model returns raw logits; task postprocessing and metadata stay outside ONNX.

The exporter uses ONNX opset 17, validates the structure with onnx.checker, and writes compact provenance/schema metadata to model.onnx.json. It supports EdgeNetwork multiclass models and FineTunedEdgeNetwork binary models. Empty graphs are invalid; single-node graphs (zero edges) are represented and pooled with a zero edge contribution.

Direct DGL export was not retained: the active model uses DGL graph mutation and reductions (apply_edges, update_all, and graph pooling) that are not a portable ONNX contract. The adapter expresses those operations with standard tensor indexing, index_add, and per-graph means.