Save XGBoost Model to File Using Pickle

When you’ve spent time training an XGBoost model, you’ll often want to save it for future use or sharing with others. Python’s pickle module provides a quick and easy way to serialize your model to a file.

import pickle
from sklearn.datasets import load_iris
from xgboost import XGBClassifier

# Load example dataset
X, y = load_iris(return_X_y=True)

# Initialize and train an XGBoost model
model = XGBClassifier(n_estimators=100, learning_rate=0.1, random_state=42)
model.fit(X, y)

# Save model to file using pickle
with open('xgb_model.pkl', 'wb') as file:
    pickle.dump(model, file)

# Load model from file
with open('xgb_model.pkl', 'rb') as file:
    loaded_model = pickle.load(file)

# Use the loaded model to make predictions
predictions = loaded_model.predict(X)

Here’s what’s happening:

1. We train an XGBoost classifier on the iris dataset (just as an example).
2. We open a file in write binary mode ('wb') and use pickle.dump() to save the model to the file.
3. Later, we open the file in read binary mode ('rb') and use pickle.load() to load the model from the file.
4. We can then use the loaded model just as we would use the original - here we demonstrate making predictions.