microsoft / microsoft/onnxruntime
[Feature Request] Release the retained NodeProtos after session initialization
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the feature request
`Graph::AddNode` copies each `NodeProto`'s attributes into the owning `Node`, so once a session is initialized the `NodeProto` list the `Graph` still holds from the loaded model is a pure duplicate of state the `Node` instances own. Nothing reads it again: `Node::original_node_proto_` is cleared at the end of the first `Graph::Resolve`, and serialization goes through `Graph::ToGraphProto`, which rebuilds the node list from the `Node` instances.
For most models this duplication is negligible, because the bulk of the model is initializers. It is not negligible for models whose payload is node attributes — ai.onnx.ml tree ensembles above all, where the whole model is one node's attributes and the retained copy is a second copy of the entire model, held for the lifetime of the session.
I measured this on a synthetic TreeEnsembleRegressor (8000 trees, 1,016,000 nodes, 37 MB on disk, no initializers) using onnxruntime_perf_test on macOS arm64. Releasing the retained protos after initialization drops the session's live heap from 150.5 MiB to 41.5 MiB — 109 MiB and 1,016,043 allocations freed, the allocation delta matching the model's node count exactly. What remains is essentially the tree kernel's own optimized representation. Peak RSS is unchanged, since both copies necessarily coexist during load. Steady-state RSS is unchanged on macOS because the system allocator retains the freed pages; on glibc the large repeated-attribute arrays are above MMAP_THRESHOLD, so free() munmaps them and most of the saving should appear as an RSS reduction (I verified that allocator behaviour separately on Debian 12 / glibc 2.36, but have not yet run ORT itself on Linux).
Proposed shape: a new session config key `session.release_node_protos_after_init`, off by default. When set to "1", `InferenceSession::Initialize()` calls a new `Graph::ReleaseNodeProtos()` right after `PruneRemovableAttributes()`, which clears the node list for the main graph and recurses into subgraphs, and flags each graph as needing a proto sync so serialization still rebuilds correctly from the `Node` instances. The one thing it gives up is resolving the graph again afterwards, because `Graph::Resolve` validates nodes via `Node::ToProto()` without updating subgraphs.
Alternatives considered:
- Doing it unconditionally when no optimized model is being saved. Avoids the config key, but silently removes the ability to re-resolve a graph and can't be turned off without a rebuild if it breaks someone. An opt-in key seemed like the right way to introduce it; happy to switch if you'd prefer the default flipped once it has mileage.
- Not copying attributes into `Node` in the first place. A much larger change to graph construction, and the `Node` copy is the one inference actually uses.
I have a PR ready with three unit tests (option on, option off, and a control-flow model that runs and then round-trips both subgraph bodies through `ToGraphProto`). Happy to open it, or to adjust the design first if you'd rather it look different.
### Describe scenario use case
Long-lived inference processes that host ai.onnx.ml tree ensemble models. Gradient-boosted tree models converted to ONNX carry their entire payload in node attributes rather than initializers, so every such session currently holds a full second copy of the model for as long as it lives. Where a single process hosts many models, that overhead scales with the number of loaded models and is pure waste — it is never read after initialization.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start at InferenceSession::Initialize() and inspect Graph::PruneRemovableAttributes(), Graph::Resolve(), Graph::ToGraphProto(), and the proposed Graph::ReleaseNodeProtos(). Review the three mentioned unit tests for option-on, option-off, and control-flow round trips. Done means the retained protos are released when configured, subgraphs serialize correctly, and existing behavior remains unchanged by default.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- backend, machine-learning
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100