How to convert XGBClassifier with dart booster to ONNX?

81 views Asked by At

I need to convert a XGBClassifier with booster='dart' into an onnx model. But the converting failed with an error on last line:

.venv\Lib\site-packages\onnxmltools\convert\xgboost\common.py", line 40, in get_xgb_params
gbp = config["learner"]["gradient_booster"]["gbtree_model_param"]
KeyError: 'gbtree_model_param'

My demo code:

from onnxmltools.convert import convert_xgboost
from onnxmltools.convert.common import data_types
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
import xgboost as xgb

digits = load_digits()
X, y = digits.data, digits.target  # Our train data shape is (x, 64) where x is total samples
X_train, X_test, y_train, y_test = train_test_split(X, y)

booster = xgb.XGBClassifier(booster='dart', n_estimators=50, max_depth=10, learning_rate=0.1, random_state=42)

booster.fit(X_train, y_train)


initial_type = [('float_input', data_types.FloatTensorType([1, 64]))]
booster_onnx = convert_xgboost(booster, initial_types=initial_type)

How could I solve this?

1

There are 1 answers

1
Krzysztof Kramarz On

It seems to be an issue related to conversion of any xgb model introduced in the latest version of onnxmltools. I recommend using onnxmltools==1.11.0, it worked for me.