Training on GPU

CatBoost supports training on GPUs.

Choose the implementation for more details on the parameters that are required to start training on GPU.

The parameters that enable and customize training on GPU are set in the constructors of the following classes:

ClassParametersDescription
task_type

The processing unit type to use for training.

Possible values:
  • CPU
  • GPU
devices

IDs of the GPU devices to use for training (indices are zero-based).

Format

  • <unit ID> for one device (for example, 3)
  • <unit ID1>:<unit ID2>:..:<unit IDN> for multiple devices (for example, devices='0:1:3')
  • <unit ID1>-<unit IDN> for a range of devices (for example, devices='0-3')
Note. Other training parameters are also available. Some of them are CPU-specific or GPU-specific. See the Training parameters section for more details.
For example, use the following code to train a classification model on GPU:
from catboost import CatBoostClassifier

train_data = [[0,3],
              [4,1],
              [8,1],
              [9,1]]
train_labels = [0,0,1,1]

model = CatBoostClassifier(iterations=1000, task_type = "GPU")
model.fit(train_data, train_labels, verbose = False)