XGBoosting Home | About | Contact | Examples

Save XGBoost Model with joblib

While Python’s pickle module is a popular choice for saving and loading models, the joblib package offers a more efficient alternative, especially for large numpy arrays often used in XGBoost.

from joblib import dump, load
from sklearn.datasets import make_classification
from xgboost import XGBClassifier

# Generate a random classification dataset
X, y = make_classification(n_samples=1000, n_classes=2, random_state=42)

# 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 joblib
dump(model, 'xgb_model.joblib') 

# Load model from file
loaded_model = load('xgb_model.joblib') 

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

Here’s the breakdown:

  1. We train an XGBoost classifier on the dataset.
  2. We save the trained model to a file named ‘xgb_model.joblib’ using joblib.dump().
  3. Later, we load the model from the file using joblib.load().
  4. We then use the loaded model to make predictions, just as we would with the original model.

See Also