You’ve spent countless hours researching, tweaking, and training the perfect XGBoost model.
Your model is performing exceptionally well and you’re ready to celebrate.
But wait, now you need to deploy it, and suddenly, you’re faced with a problem.
How do you save your model for future use?
Don’t worry, there’s a simple solution to this!
In this article, I will walk you through how to save and load your XGBoost models.
This means you can train your model once, save it, and then reload it whenever you need to make predictions.
No more retraining from scratch every time!
Let’s dive in and learn how to save our hard work for future use.
Simple XGBoost Model for Demonstration
Before we dive into saving and loading models, let’s first build a simple XGBoost model we can use for demonstration.
We’ll use the Wine Quality dataset from the UCI Machine Learning Repository.
First, we need to import the necessary libraries and load the dataset.
import pandas as pd
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
# Load dataset
url = "https://archive.ics.uci.edu/ml/machine-learning-databases/wine-quality/winequality-red.csv"
data = pd.read_csv(url, delimiter=";")
Now that we have our data, let’s split it into training and testing sets. We’ll use the ‘quality’ column as our target variable.
# Split data into X and y
X = data.drop('quality', axis=1)
y = data['quality']
# Split data into train and test sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=7)
With our data prepared, we can now build our XGBoost model.
We’ll use the XGBClassifier and fit it to our training data.
# Define the model
model = XGBClassifier()
# Fit the model
model.fit(X_train, y_train)
Great! We now have a simple XGBoost trained on the Wine Quality dataset.
We’ll use this model in the next sections to demonstrate how to save and load models.
Saving XGBoost Models
Now that we have our trained XGBoost model, let’s save it for future use.
XGBoost provides a function called save_model that allows us to do this.
model.save_model('xgb_model.json')
In the code above, we’re saving our model to a file named ‘xgb_model.json’.
The save_model function supports different file formats including JSON and binary format.
It’s important to note that there’s a difference between the sklearn API and the booster method (native API) when it comes to saving models.
When using the sklearn API, you use the model.save_model() function as shown above.
However, when using the Native API, you need to use booster.save_model().
# Save model using booster
booster = model.get_booster()
booster.save_model('booster_model.json')
In the code above, we first get the booster from the model using model.get_booster(), and then save the booster using booster.save_model().
If we had trained our model using the Native API, we would have to use booster.save_model() to save the model with the booster object that was fit to the training data.
Remember: if you use the Scikit-learn API to save the model, you must load it using the Scikit-learn API, the same goes for the Native API.
Loading XGBoost Models
Now that we’ve saved our model, let’s load it back into memory.
We can do this using the load_model function provided by XGBoost.
# Load model from file
loaded_model = XGBClassifier()
loaded_model.load_model('xgb_model.json')
In the code above, we first create a new instance of XGBClassifier, so the library knows what type of model we’re loading.
Then, we load the model from the file ‘xgb_model.json’ into this new instance.
Just like with saving models, there’s a difference between the sklearn API and the Native API when it comes to loading models.
When using the sklearn API, you use the model.load_model() function as shown above.
However, when using the booster method, you need to use booster.load_model().
# Load model using booster
import xgboost
booster = xgboost.Booster()
booster.load_model('booster_model.json')
In the code above, we first create a new instance of a Booster object, and then load the model into it using booster.load_model().
Making Predictions With The Loaded Model
Now that we’ve loaded our model, you can make predictions with it just like you would with any other model.
# Make predictions on test data
y_pred = loaded_model.predict(X_test)
If you are using the Native API, remember to convert your input data to DMatrix before making predictions.
Saving and Loading XGBoost Models With Pickle
Pickle is a Python module that is used for serializing (pickling) and deserializing (unpickling) Python object structures.
It’s another way to save and load your XGBoost models, and it can be faster for large models.
Let’s start by saving our XGBoost model with Pickle.
import pickle
# Save the model to a file
pickle.dump(model, open("xgb.pkl", "wb"))
In the code above, we’re using the dump function from the pickle module to save our model to a file named ‘xgb.pkl’.
The ‘wb’ argument indicates that the file is opened for writing in binary mode.
Now, let’s load the model back into memory using Pickle.
# Load model from file
loaded_model = pickle.load(open("xgb.pkl", "rb"))
Here, we’re using the load function from the pickle module to load our model from the file ‘xgb.pkl’.
The ‘rb’ argument indicates that the file is opened for reading in binary mode.
One of the advantages of using Pickle is that it can be faster for large models, and it saves the model object in a binary format so you don’t need to create a new instance of it before loading.
However, it’s important to note that Pickle is not secure against erroneous or maliciously constructed data.
Never unpickle data received from an untrusted or unauthenticated source.
Saving and Loading XGBoost Models With Joblib
Joblib is another Python library used for serializing Python objects, similar to Pickle.
However, it’s more efficient on objects that carry large numpy arrays, making it suitable for saving and loading large machine learning models like XGBoost.
This is my preferred method.
Let’s start by saving our XGBoost model with Joblib.
from joblib import dump
# Save the model to a file
dump(model, "xgb.pkl")
In the code above, we’re using the dump function from the joblib module to save our model to a file named ‘“xgb.pkl”’.
Now, let’s load the model back into memory using Joblib.
from joblib import load
# Load model from file
loaded_model = load("xgb.pkl")
Here, we’re using the load function from the joblib module to load our model from the file ‘“xgb.pkl”’.
One of the advantages of using Joblib over Pickle is that it can be faster and more efficient for large models, especially those with large numpy arrays.
And it has the same advantage that you don’t need to create a new instance of the model before loading it.
However, similar to Pickle, Joblib is not secure against erroneous or maliciously constructed data.
Never load a joblib serialized file received from an untrusted or unauthenticated source.