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.