Key Concepts of Ray Train

There are four main concepts in the Ray Train library.

1. Trainers execute distributed training.

2. Configuration objects are used to configure training.

3. Checkpoints are returned as the result of training.

4. Predictors can be used for inference and batch prediction.

Trainers

Trainers are responsible for executing (distributed) training runs. The output of a Trainer run is a Result that contains metrics from the training run and the latest saved Checkpoint. Trainers can also be configured with Datasets and Preprocessors for scalable data ingest and preprocessing.

Deep Learning, Tree-Based, and other Trainers

There are three categories of built-in Trainers:

Deep Learning Trainers

Ray Train supports the following deep learning trainers:

• TorchTrainer

• TensorflowTrainer

• HorovodTrainer

• LightningTrainer

• HuggingFaceTrainer

For these trainers, you usually define your own training function that loads the model and executes single-worker training steps. Refer to the following guides for more details:

• Deep learning user guide

• Quick overview of deep-learning trainers in the Ray AIR documentation

Tree-Based Trainers

Tree-based trainers utilize gradient-based decision trees for training. The most popular libraries for this are XGBoost and LightGBM.

• XGBoostTrainer

• LightGBMTrainer

For these trainers, you just pass a dataset and parameters. The training loop is configured automatically.

• XGBoost/LightGBM user guide

• Quick overview of tree-based trainers in the Ray AIR documentation

Other Trainers

Some trainers don’t fit into the other two categories, such as:

• LightningTrainer for PyTorch Lightning

• HuggingFaceTrainer for NLP

• RLTrainer for reinforcement learning

• SklearnTrainer for (non-distributed) training of sklearn models.

• Other trainers in the Ray AIR documentation

Train Configuration

Trainers are configured with configuration objects. There are two main configuration classes, the ScalingConfig and the RunConfig. The latter contains subconfigurations, such as the FailureConfig, SyncConfig and CheckpointConfig.

Check out the Configurations User Guide for an in-depth guide on using these configurations.

Train Checkpoints

Calling Trainer.fit() returns a Result object, which includes information about the run such as the reported metrics and the saved checkpoints.

Checkpoints have the following purposes:

• They can be passed to a Trainer to resume training from the given model state.

• They can be used to create a Predictor / BatchPredictor for scalable batch prediction.

• They can be deployed with Ray Serve.

Train Predictors

Predictors are the counterpart to Trainers. A Trainer trains a model on a dataset, and a predictor uses the resulting model and performs inference on it.

Each Trainer has a respective Predictor implementation that is compatible with its generated checkpoints.

Example: XGBoostPredictor

import numpy as np
import ray

from ray.train.xgboost import XGBoostTrainer, XGBoostPredictor
from ray.air.config import ScalingConfig

train_dataset = ray.data.from_items([{"x": x, "y": x + 1} for x in range(32)])
trainer = XGBoostTrainer(
    label_column="y",
    params={"objective": "reg:squarederror"},
    scaling_config=ScalingConfig(num_workers=3),
    datasets={"train": train_dataset},
)
result = trainer.fit()

predictor = XGBoostPredictor.from_checkpoint(result.checkpoint)
predictions = predictor.predict(np.expand_dims(np.arange(32, 64), 1))

A predictor can be passed into a BatchPredictor is used to scale up prediction over a Ray cluster. It takes a Ray Dataset as input.

Example: Batch prediction with XGBoostPredictor

import pandas as pd
from ray.train.batch_predictor import BatchPredictor

batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, XGBoostPredictor)
predict_dataset = ray.data.from_pandas(pd.DataFrame({"x": np.arange(32)}))
predictions = batch_predictor.predict(
    data=predict_dataset,
    batch_size=8,
    min_scoring_workers=2,
)
predictions.show()

See the Predictors user guide for more information and examples.