File size: 1,497 Bytes
916755e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | # 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:
```bash
uv sync --extra root-gnn --extra onnx
```
Export a checkpoint with the project CLI:
```bash
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.
|