This tutorial and the assets can be downloaded as part of the Wallaroo Tutorials repository.
The following tutorial is a brief example of how to convert a XGBoost ML model to the ONNX standard. This allows organizations that have trained XGBoost models to convert them and use them with Wallaroo.
This tutorial assumes that you have a Wallaroo instance and are running this Notebook from the Wallaroo Jupyter Hub service.
This tutorial provides the following:
housing_model_xgb.pkl
: A pretrained model used as part of the Notebooks in Production tutorial. This model has a total of 18 columns.The first step is to import our libraries we will be using.
import onnx
import pickle
from onnxmltools.convert import convert_xgboost
from skl2onnx.common.data_types import FloatTensorType, DoubleTensorType
The following variables are required to be known before the process can be started:
# set the number of columns
ncols = 18
TARGET_OPSET = 15
Next we will load our model that has been saved in the pickle
format and unpickle it.
# load the xgboost model
with open("housing_model_xgb.pkl", "rb") as f:
xgboost_model = pickle.load(f)
The convert_xgboost
method has the following format and requires the following inputs:
convert_xgboost({XGBoost Model},
{XGBoost Model Type},
[
('input',
{Tensor Data Type}([None, {ncols}]))
],
target_opset={TARGET_OPSET})
tree-based classifier
.FloatTensorType
or DoubleTensorType
from the skl2onnx.common.data_types
library.With all of our data in place we can now convert our XBBoost model to ONNX using the convert_xgboost
method.
onnx_model_converted = convert_xgboost(xgboost_model, 'tree-based classifier',
[('input', FloatTensorType([None, ncols]))],
target_opset=TARGET_OPSET)
With the model converted to ONNX, we can now save it and use it in a Wallaroo pipeline.
onnx.save_model(onnx_model_converted, "housing_model_xgb.onnx")