Why Ray Train?
Just a few of the many reasons why machine learning engineers are choosing to scale their deep learning workloads with Ray Train.
Native multi-GPU support
Easily scale from single threaded to multi-GPU training in under 10 lines of code.
With a best-in-class API for gradient descent, migrate to production or scale to a large cluster without rewrite code.
Seamlessly works with best-in-class deep learning frameworks including PyTorch, Tensorflow, Horovod, and many more
Try It Yourself
Install Ray Train (and PyTorch) with
pip install ray torch and give this example a try.
import ray.train as train from ray.train import Trainer import torch def train_func(): # Setup model. model = torch.nn.Linear(1, 1) model = train.torch.prepare_model(model) loss_fn = torch.nn.MSELoss() optimizer = torch.optim.SGD(model.parameters(), lr=1e-2) # Setup data. input = torch.randn(1000, 1) labels = input * 2 dataset = torch.utils.data.TensorDataset(input, labels) dataloader = torch.utils.data.DataLoader(dataset, batch_size=32) dataloader = train.torch.prepare_data_loader(dataloader) # Train. for _ in range(5): for X, y in dataloader: pred = model(X) loss = loss_fn(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() return model.state_dict() trainer = Trainer(backend="torch", num_workers=4) trainer.start() results = trainer.run(train_func) trainer.shutdown() print(results)
Scale more workloads with Ray
Expand your Ray journey beyond deep learning and bring fast and easy distributed execution to other use cases.
Scale general Python apps.
Scale hyperparameter search.
Scale data loading and processing.
O'Reilly Learning Ray Book
Get your free copy of early release chapters of Learning Ray, the first and only comprehensive book on Ray and its ecosystem, authored by members on the Ray engineering team