TorchTrainer#

class ray.train.torch.TorchTrainer(train_loop_per_worker: Callable[[], Any] | Callable[[Dict], Any], *, train_loop_config: Dict | None = None, torch_config: TorchConfig | None = None, scaling_config: ScalingConfig | None = None, run_config: RunConfig | None = None, datasets: Dict[str, Dataset | Callable[[], Dataset]] | None = None, dataset_config: DataConfig | None = None, validation_config: ValidationConfig | None = None, metadata: Dict[str, Any] | None = None, resume_from_checkpoint: Checkpoint | None = None)[source]#

Bases: DataParallelTrainer

A Trainer for data parallel PyTorch training.

At a high level, this Trainer does the following:

  1. Launches multiple workers as defined by the scaling_config.

  2. Sets up a distributed PyTorch environment on these workers as defined by the torch_config.

  3. Ingests the input datasets based on the dataset_config.

  4. Runs the input train_loop_per_worker(train_loop_config) on all workers.

For more details, see:

Examples

>>> import os
>>> import tempfile
>>> import torch
>>> from torch import nn
>>> from torch.nn.parallel import DistributedDataParallel
>>> import ray.train
>>> from ray.train.torch import TorchTrainer
>>> # If using GPUs, set this to True.
>>> use_gpu = False
>>> # Number of processes to run training on.
>>> num_workers = 2
>>> # Define your network structure.
>>> class NeuralNetwork(nn.Module):
...     def __init__(self):
...         super(NeuralNetwork, self).__init__()
...         self.layer1 = nn.Linear(1, 32)
...         self.relu = nn.ReLU()
...         self.layer2 = nn.Linear(32, 1)
...
...     def forward(self, input):
...         return self.layer2(self.relu(self.layer1(input)))
>>> # Training loop.
>>> def train_fn_per_worker(config):
...
...     # Read configurations.
...     lr = config["lr"]
...     batch_size = config["batch_size"]
...     num_epochs = config["num_epochs"]
...
...     # Fetch training dataset.
...     train_dataset_shard = ray.train.get_dataset_shard("train")
...
...     # Instantiate and prepare model for training.
...     model = NeuralNetwork()
...     model = ray.train.torch.prepare_model(model)
...
...     # Define loss and optimizer.
...     loss_fn = nn.MSELoss()
...     optimizer = torch.optim.SGD(model.parameters(), lr=lr)
...
...     # Create data loader.
...     dataloader = train_dataset_shard.iter_torch_batches(
...         batch_size=batch_size, dtypes=torch.float
...     )
...
...     # Train multiple epochs.
...     for epoch in range(num_epochs):
...         # Train epoch.
...         for batch in dataloader:
...             output = model(batch["input"])
...             loss = loss_fn(output, batch["label"])
...             optimizer.zero_grad()
...             loss.backward()
...             optimizer.step()
...
...         # Create checkpoint.
...         base_model = (
...             model.module
...             if isinstance(model, DistributedDataParallel)
...             else model
...         )
...
...         with tempfile.TemporaryDirectory() as temp_checkpoint_dir:
...             torch.save(
...                 {"model_state_dict": base_model.state_dict()},
...                 os.path.join(temp_checkpoint_dir, "model.pt"),
...             )
...             checkpoint = ray.train.Checkpoint.from_directory(temp_checkpoint_dir)
...             # Report metrics and checkpoint.
...             ray.train.report({"loss": loss.item()}, checkpoint=checkpoint)
>>> # Define datasets.
>>> train_dataset = ray.data.from_items(
...     [{"input": [x], "label": [2 * x + 1]} for x in range(128)]
... )
>>> # Initialize the Trainer.
>>> trainer = TorchTrainer(
...     train_fn_per_worker,
...     train_loop_config={"num_epochs": 1, "lr": 0.01, "batch_size": 32},
...     scaling_config=ray.train.ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),
...     datasets={"train": train_dataset},
... )
>>> # Train the model.
>>> result = trainer.fit()
>>> # Inspect the results.
>>> final_loss = result.metrics["loss"]
Parameters:

  • train_loop_per_worker (Callable[[], Any] | Callable[[Dict], Any]) – The training function to execute on each worker. This function can either take in zero arguments or a single Dict argument which is set by defining train_loop_config. Within this function you can use any of the Ray Train Loop utilities.

  • train_loop_config (Dict | None) – A configuration Dict to pass in as an argument to train_loop_per_worker. This is typically used for specifying hyperparameters. Passing large datasets via train_loop_config is not recommended and may introduce large overhead and unknown issues with serialization and deserialization.

  • torch_config (TorchConfig | None) – The configuration for setting up the PyTorch Distributed backend. If set to None, a default configuration will be used in which GPU training uses NCCL and CPU training uses Gloo.

  • scaling_config (ScalingConfig | None) – The configuration for how to scale data parallel training. num_workers determines how many Python processes are used for training, and use_gpu determines whether or not each process should use GPUs. See ScalingConfig for more info.

  • run_config (RunConfig | None) – The configuration for the execution of the training run. See RunConfig for more info.

  • datasets (Dict[str, Dataset | Callable[[], Dataset]] | None) – The Ray Datasets to ingest for training. Datasets are keyed by name ({name: dataset}). Each dataset can be accessed from within the train_loop_per_worker by calling ray.train.get_dataset_shard(name). Sharding and additional configuration can be done by passing in a dataset_config.

  • dataset_config (DataConfig | None) – The configuration for ingesting the input datasets. By default, all the Ray Dataset are split equally across workers. See DataConfig for more details.

  • validation_config (ValidationConfig | None) – [Alpha] Configuration for checkpoint validation. If provided and ray.train.report is called with the validation argument, Ray Train will validate the reported checkpoint using the validation function specified in this config.

  • metadata (Dict[str, Any] | None) – [Deprecated]

  • resume_from_checkpoint (Checkpoint | None) – [Deprecated]

Methods

can_restore

[Deprecated] Checks if a Train experiment can be restored from a previously interrupted/failed run.

fit

Launches the Ray Train controller to run training on workers.

restore

[Deprecated] Restores a Train experiment from a previously interrupted/failed run.