onnx / onnxmltools

ONNXMLTools enables conversion of models to ONNX
https://onnx.ai
Apache License 2.0
992 stars 181 forks source link

XGBoostRegression Convert Sklearn #524

Open drvshavva opened 2 years ago

drvshavva commented 2 years ago

Hi, I am trying to save my xgboost regression model as an onnx but I got the error as below:

Unable to create node 'TreeEnsembleRegressor' with name='TreeEnsembleRegressor'.

My versions:

onnx 1.8.0 onnxconverter-common 1.6.1 onnxmltools 1.6.1 onnxruntime 1.3.0 skl2onnx 1.6.1 xgboost 1.0.2

Code:

from xgboost import XGBRegressor

from skl2onnx.common.shape_calculator import calculate_linear_regressor_output_shapes
from skl2onnx import update_registered_converter
from onnxmltools.convert.xgboost.operator_converters.XGBoost import convert_xgboost
from sklearn.datasets import load_iris
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

update_registered_converter(XGBRegressor, 'XgBoostRegression', calculate_linear_regressor_output_shapes,
convert_xgboost)

data = load_iris()
X = data.data[:, :2]
y = data.target

ind = numpy.arange(X.shape[0])
numpy.random.shuffle(ind)
X = X[ind, :].copy()
y = y[ind].copy()

pipe = Pipeline([('scaler', StandardScaler()),
('XGBRegressor', XGBRegressor(n_estimators=3))])
pipe.fit(X, y)

convert_sklearn(pipe, 'pipeline_xgboost',
[('input', FloatTensorType([None, 2]))])

Can you help with this ? Thanks.

xadupre commented 2 years ago

This code works with the latest version of the packages:

import numpy
from xgboost import XGBRegressor

from skl2onnx.common.shape_calculator import calculate_linear_regressor_output_shapes
from skl2onnx import update_registered_converter
from onnxmltools.convert.xgboost.operator_converters.XGBoost import convert_xgboost
from sklearn.datasets import load_iris
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

update_registered_converter(XGBRegressor, 'XgBoostRegression', calculate_linear_regressor_output_shapes,
convert_xgboost)

data = load_iris()
X = data.data[:, :2]
y = data.target

ind = numpy.arange(X.shape[0])
numpy.random.shuffle(ind)
X = X[ind, :].copy()
y = y[ind].copy()

pipe = Pipeline([
    ('scaler', StandardScaler()),
    ('XGBRegressor', XGBRegressor(n_estimators=3))])
pipe.fit(X, y)

onx = convert_sklearn(
    pipe, 'pipeline_xgboost',
    [('input', FloatTensorType([None, 2]))],
    target_opset={'': 15, 'ai.onnx.ml': 2})
print(onx)