ShaharA
ShaharA

Reputation: 903

Limiting the number of threads used by XGBoost

I want to use XGBoost for online production purposes (Python 2.7 XGBoost API). In order to be able to do that I want to control and limit the number of threads used by XGBoost at the predict operation.

I'm using the sklearn compatible regressor offered by XGBoost (xgboost.XGBRegressor), and been trying to use the param nthread in the constructor of the regressor to limit the max threads used to 1.

Unfortunately, XGBoost keeps using multiple threads regardless of the value set in nthread.

Is there another way to limit XGBoost and force it to perform the predict operation using n=1 threads?

Upvotes: 7

Views: 6736

Answers (4)

Maxim Egorushkin
Maxim Egorushkin

Reputation: 136435

There is no option to set the number of threads used for prediction using XGBoost-1.5.2 API, unfortunately, setting DMatrix nthreads parameter has no effect.

XGBoost uses OpenMP for parallelization. Set environment variable OMP_THREAD_LIMIT to the maximum number of threads OpenMP can use. E.g. export OMP_THREAD_LIMIT=1 causes the following messages:

OMP: Warning #96: Cannot form a team with 32 threads, using 1 instead.
OMP: Hint Consider unsetting KMP_DEVICE_THREAD_LIMIT (KMP_ALL_THREADS), KMP_TEAMS_THREAD_LIMIT, and OMP_THREAD_LIMIT (if any are set).

In Python add these lines before other imports:

import os
os.environ["OMP_THREAD_LIMIT"] = "1"

Upvotes: 0

mirekphd
mirekphd

Reputation: 6831

Currently n_jobs can be used to limit threads at predict time:

model._Booster.set_param('n_jobs', 2)

More info

Formerly (but now deprecated):

model._Booster.set_param('nthread', 2)

Upvotes: 1

Korin Du
Korin Du

Reputation: 119

I ran into the same issue and figured it out. The correct answer is to set system environment variables. For python scripts:

import os
os.environ['OMP_NUM_THREADS'] = "1"

would work.

Make sure to put these two lines before you import any other package, otherwise it might not work.

Upvotes: 11

ShaharA
ShaharA

Reputation: 903

Answer: Standard set_params with nthread fails, but when using regr._Booster.set_param('nthread', 1) I was able to limit XGBoost to using a single thread.

As mentioned above the env variable OMP_NUM_THREADS=1 works as well.

Upvotes: 2

Related Questions