心脏病预测Scikit-learn Pipeline转ONNX后加载推理失败求助

问题描述 #

搭建了用于心脏病预测的Scikit-learn Pipeline，已成功转换为ONNX格式，但加载模型推理时失败，报错如下：

Fail: [ONNXRuntimeError] : 1 : FAIL : Load model from ./pipeline_xgboost.onnx failed:Node (LabelEncoder2) Op (LabelEncoder) [ShapeInferenceError] Input type is not int64 tensor but keys_int64s is set.

当前定义的输入类型：

input_types = [
 ('age', FloatTensorType(shape=[None, 1])), #could not parse age as int 
 ('sex', Int64TensorType(shape=[None, 1])),
 ('cp', Int64TensorType(shape=[None, 1])),
 ('trestbps', Int64TensorType(shape=[None, 1])),
 ('chol', Int64TensorType(shape=[None, 1])),
 ('fbs', Int64TensorType(shape=[None, 1])),
 ('restecg', Int64TensorType(shape=[None, 1])),
 ('thalach', Int64TensorType(shape=[None, 1])),
 ('exang', Int64TensorType(shape=[None, 1])),
 ('oldpeak', FloatTensorType(shape=[None, 1])),
 ('slope', Int64TensorType(shape=[None, 1])),
 ('ca', Int64TensorType(shape=[None, 1])),
 ('thal', StringTensorType(shape=[None, 1]))]

原因分析 #

报错核心是：LabelEncoder2节点配置了keys_int64s（即用整数作为编码映射的键），但该节点对应的输入张量类型不是int64，类型不匹配导致模型加载失败。

解决步骤 #

1. 定位LabelEncoder2对应的特征
   查看你的Scikit-learn Pipeline结构，找到标记为LabelEncoder2的预处理步骤，确认它对应的是哪个输入特征（比如ca、slope或其他特征）。

2. 修正输入类型与LabelEncoder的匹配性

   • 若该LabelEncoder是用来编码整数类特征的（训练时输入为整数），则对应特征的输入类型必须设为Int64TensorType，检查当前输入类型定义中该特征的类型是否正确，错误则修正。
   • 若该LabelEncoder实际处理字符串特征，需确认转换ONNX时是否错误配置了keys_int64s参数，此时应将对应特征的输入类型设为StringTensorType，重新转换模型时指定正确的编码键类型。

3. 重新转换并验证模型
   修正输入类型定义后，使用skl2onnx重新将Pipeline转换为ONNX格式，再尝试加载模型推理，确认问题是否解决。

内容的提问来源于stack exchange，提问作者TheHumanSpider