onnxruntime
de17ab8d - Bound TreeEnsemble subtree comparison (#32043)

Commit
14 days ago
Bound TreeEnsemble subtree comparison (#32043) This pull request improves the robustness and correctness of categorical node folding in the tree ensemble implementation by refactoring the subtree comparison logic to prevent infinite loops caused by cycles, and by adding a test to ensure cycles are properly detected and rejected. **Categorical Node Folding and Cycle Detection Improvements:** * Refactored the `CheckIfSubtreesAreEqual` method in `tree_ensemble_common.h` to use an iterative approach with explicit cycle detection, replacing the previous recursive implementation. This prevents stack overflows and infinite loops when cycles are present in the tree structure. (`onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h`, [[1]](diffhunk://#diff-098f04f4ce424d09171371b1dc5893d29e38c42a0f880639775de9e0fec67241L115-R116) [[2]](diffhunk://#diff-098f04f4ce424d09171371b1dc5893d29e38c42a0f880639775de9e0fec67241L411-R462) * Modified the categorical node folding logic in `AddNodes` to maintain a set of visited categorical nodes and enforce that no node is revisited, raising an error if a cycle is detected. (`onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h`, [onnxruntime/core/providers/cpu/ml/tree_ensemble_common.hR537-R544](diffhunk://#diff-098f04f4ce424d09171371b1dc5893d29e38c42a0f880639775de9e0fec67241R537-R544)) * Added `#include <set>` to support the new cycle detection logic. (`onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h`, [onnxruntime/core/providers/cpu/ml/tree_ensemble_common.hR8](diffhunk://#diff-098f04f4ce424d09171371b1dc5893d29e38c42a0f880639775de9e0fec67241R8)) **Testing:** * Added a new unit test `TreeRegressorRejectsCategoricalCycle` to verify that the implementation correctly detects and rejects cycles in categorical nodes, ensuring the model fails gracefully with an appropriate error message. (`onnxruntime/test/providers/cpu/ml/treeregressor_test.cc`, [onnxruntime/test/providers/cpu/ml/treeregressor_test.ccR768-R789](diffhunk://#diff-08b3495816c68f145657ecff63d7b5f3d56813586ec62f7324c22977e70e336bR768-R789))
Author
Parents
Loading