Save XGBoost Model in ONNX Format

ONNX (Open Neural Network Exchange) is an open format built to represent machine learning models.

Saving your XGBoost models in ONNX format enables interoperability with various frameworks and environments, such as PyTorch, Caffe2, and ML.NET, making your models more portable and accessible.

First, we must install the onnxmltools library using our preferred package manager, such as pip:

pip install onnxmltools

Next, we can save our XGBoost model in ONNX format.

from sklearn.datasets import make_classification
from xgboost import XGBClassifier
import onnxmltools
from onnxconverter_common.data_types import FloatTensorType

# Generate a synthetic dataset
X, y = make_classification(n_samples=100, n_features=20, random_state=42)

# Fit an XGBoost model
model = XGBClassifier(eval_metric='logloss')
model.fit(X, y)

# Convert the XGBoost model to ONNX format
onnx_model = onnxmltools.convert_xgboost(model, initial_types=[('input', FloatTensorType([None, X.shape[1]]))])

# Save the ONNX model
with open("xgboost_model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

Here’s what we’re doing:

1. We fit the XGBoost model on the dataset.
2. We convert the trained model to ONNX format using convert_xgboost() from onnxmltools. This function takes the model and a name for the converted model as parameters.
3. We save the ONNX model to a file named xgboost_model.onnx using SerializeToString().

Once saved in ONNX format, the model can be loaded and used for inference in different frameworks that support ONNX.