onnxruntime
41f9056c - Fix O(n²) model load time for TreeEnsemble with categorical feature chains (#27391)

Commit
152 days ago
Fix O(n²) model load time for TreeEnsemble with categorical feature chains (#27391) ### Description Profiling shows that CheckIfSubtreesAreEqual is invoked recursively for many node pairs for LightGBM models with categorical features. A significant portion of this work consists of self-comparisons (left_id == right_id), leading to effectively O(n²) comparing trees to themselves during model loading. This change adds a fast-path for trivial equality, avoiding unnecessary recursive comparisons. Example results: - model with 7K BRANCH_EQ nodes: 527 ms → 47 ms (~11× faster) - model with 106K BRANCH_EQ nodes: 141 s → 80 ms (~1760× faster) ### Motivation and Context We have some LightGBM exported models that make heavy use of categorical features and exhibit extremely slow load times (minutes for a single 2.5mb model). Heres a diagram to illustrate the issue: <img width="1008" height="1229" alt="image" src="https://github.com/user-attachments/assets/348e16cb-9eec-448f-ac5c-e1edb60e2a3d" /> the 106K model has much longer "member of" chains, with chains that lead into more chains: <details> <summary>"trees"</summary> <img width="1405" height="593" alt="image" src="https://github.com/user-attachments/assets/12f0c43f-5987-4b33-9001-2a2b526e537f" /> </details> Interestingly we did also try using the new onnx.ml opset 5 node that has MEMBER, but it seems even slower as it recreates these branch EQ chains.
Author
Parents
Loading