ONNXMLTools中RandomForestClassifier转换问题的分析与修复
项目地址: https://gitcode.com/gh_mirrors/on/onnxmltools 问题背景
在使用ONNXMLTools将PySpark的RandomForestClassifier模型转换为ONNX格式时，发现了一个影响模型预测结果的严重问题。转换后的ONNX模型在进行推理时，所有输出概率（以及最终的预测结果）都变得完全相同，这显然不符合预期行为。
问题根源分析
经过深入排查，发现问题出在决策树节点的类型转换上。在原始Spark ML模型中，决策树使用的是BRANCH_LEQ（小于等于分支）节点，但在转换为ONNX格式后，这些节点全部被错误地转换成了BRANCH_EQ（等于分支）节点。
这种转换错误直接导致了模型决策逻辑的改变，使得所有样本都沿着相同的路径进行预测，最终产生了完全相同的预测结果。
技术细节剖析
问题的根源可以追溯到ONNXMLTools源码中的两个关键函数：
-
sparkml_tree_dataset_to_sklearn函数（位于tree_ensemble_common.py）
在这个函数中，阈值(threshold)的提取逻辑存在问题。当前代码直接使用了数组形式的leftCategoriesOrThreshold：
threshold.append(item["leftCategoriesOrThreshold"])而实际上，应该像处理元组情况那样，只取数组中的第一个元素：
threshold.append(item["leftCategoriesOrThreshold"][0] if len(item["leftCategoriesOrThreshold"]) >= 1 else -1.0) -
rewrite_ids_and_process函数（位于tree_helper.py）
由于上述阈值提取问题，导致这个函数在处理节点时，将所有BRANCH_LEQ节点错误地识别为BRANCH_EQ节点。
解决方案
修复方案相对直接：修改阈值提取逻辑，确保正确处理数组形式的阈值数据。具体修改如下：
threshold.append(item["leftCategoriesOrThreshold"][0] if len(item["leftCategoriesOrThreshold"]) >= 1 else -1.0)
这一修改与处理元组情况的现有逻辑保持一致，确保了不同类型数据的一致性处理。
影响与意义
这个修复解决了以下关键问题：
- 恢复了RandomForestClassifier在ONNX转换后的正确预测能力
- 确保了Spark ML模型与ONNX模型之间的一致性
- 提高了模型转换的可靠性
对于使用PySpark进行机器学习开发并需要将模型部署到ONNX运行时的用户来说，这一修复至关重要。它确保了模型转换后能够保持原有的预测准确性和决策逻辑。
最佳实践建议
对于使用ONNXMLTools进行模型转换的用户，建议：
- 始终验证转换前后模型的预测结果是否一致
- 对于重要项目，考虑实现自动化测试来验证模型转换的正确性
- 关注ONNXMLTools的更新，及时获取此类重要修复
这个问题的发现和修复过程也提醒我们，在模型格式转换过程中，即使是看似微小的实现细节，也可能对最终模型的预测行为产生重大影响。
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
