XGBoosting Home | About | Contact | Examples

XGBoost Multithreaded Training and Prediction

XGBoost supports multithreaded training and prediction out of the box.

By configuring the n_jobs parameter, you can specify the number of threads to use, allowing you to take full advantage of your machine’s computing power.

Additionally, setting the OMP_NUM_THREADS environment variable can further optimize performance when using OpenMP BLAS libraries.

Here’s a quick example of how to configure XGBoost for multithreaded training and prediction:

import os
# Set the number of threads to use
os.environ["OMP_NUM_THREADS"] = "4"

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier

# Generate a synthetic dataset
X, y = make_classification(n_samples=10000, n_features=20, n_informative=10, n_redundant=5, random_state=42)

# Split the data into training and testing sets
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# Initialize an XGBClassifier with multithreading enabled
model = XGBClassifier(objective='binary:logistic', n_jobs=-1, random_state=42)

# Fit the model to the training data
model.fit(X_train, y_train)

# Make predictions with the fit model
predictions = model.predict(X_test)
print(predictions[:5])

In this example, we first set the OMP_NUM_THREADS environment variable to 4, indicating that we want to use 4 threads for OpenMP BLAS libraries. This step is optional but can provide additional performance benefits.

Next, we generate a synthetic dataset using sklearn.datasets.make_classification and split it into training and testing sets.

We then initialize an XGBClassifier with the objective set to 'binary:logistic' for binary classification and n_jobs set to -1, which automatically uses all available threads for training and prediction. The random_state parameter is set to 42 for reproducibility.

Finally, we fit the model to the training data using model.fit() and make predictions on the test set using model.predict(). The first five predictions are printed to demonstrate the output.

By leveraging multithreading in XGBoost, you can significantly speed up the training and prediction processes, especially when working with large datasets or complex models.



See Also