mirror of
https://github.com/wassname/ray.git
synced 2026-07-18 12:40:56 +08:00
[raysgd] Custom training operator (#7211)
This commit is contained in:
Richard Liaw
Binary file not shown.
|
After Width: | Height: | Size: 30 KiB |
@@ -10,6 +10,8 @@ Under the hood, ``PytorchTrainer`` will create *replicas* of your model (control
|
||||
|
||||
For end to end examples leveraging RaySGD PyTorchTrainer, jump to :ref:`raysgd-pytorch-examples`.
|
||||
|
||||
.. contents:: :local:
|
||||
|
||||
Setting up training
|
||||
-------------------
|
||||
|
||||
@@ -81,6 +83,7 @@ for ``PyTorchTrainer(scheduler_creator=...)``.
|
||||
:start-after: __torch_scheduler_start__
|
||||
:end-before: __torch_scheduler_end__
|
||||
|
||||
|
||||
.. _starting-pytorch-trainer:
|
||||
|
||||
Putting things together
|
||||
@@ -115,30 +118,17 @@ You can also set the number of workers and whether the workers will use GPUs:
|
||||
num_replicas=100,
|
||||
use_gpu=True)
|
||||
|
||||
See the documentation on the PyTorchTrainer here: :ref:`ref-pytorch-trainer`. We'll look at the training APIs next.
|
||||
|
||||
Training APIs
|
||||
-------------
|
||||
|
||||
Now that the trainer is constructed, you'll naturally want to train the model.
|
||||
Now that the trainer is constructed, here's how to train the model.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
trainer.train()
|
||||
|
||||
This takes one pass over the training data.
|
||||
|
||||
To run the model on the validation data passed in by the ``data_creator``, you can simply call:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
trainer.validate()
|
||||
|
||||
You can customize the exact function that is called by using a customized training function (see :ref:`raysgd-custom-training`).
|
||||
for i in range(10):
|
||||
metrics = trainer.train()
|
||||
val_metrics = trainer.validate()
|
||||
|
||||
|
||||
Shutting down training
|
||||
----------------------
|
||||
Each ``train`` call makes one pass over the training data, and each ``validate`` call runs the model on the validation data passed in by the ``data_creator``. Provide a custom training operator (:ref:`raysgd-custom-training`) to calculate custom metrics or customize the training/validation process.
|
||||
|
||||
After training, you may want to reappropriate the Ray cluster. To release Ray resources obtained by the Trainer:
|
||||
|
||||
@@ -148,15 +138,131 @@ After training, you may want to reappropriate the Ray cluster. To release Ray re
|
||||
|
||||
.. note:: Be sure to call ``trainer.save()`` or ``trainer.get_model()`` before shutting down.
|
||||
|
||||
Initialization Functions
|
||||
------------------------
|
||||
See the documentation on the PyTorchTrainer here: :ref:`ref-pytorch-trainer`.
|
||||
|
||||
You may want to run some initializers on each worker when they are started. This may be something like setting an environment variable or downloading some data. You can do this via the ``initialization_hook`` parameter:
|
||||
|
||||
.. _raysgd-custom-training:
|
||||
|
||||
Custom Training and Validation (Operators)
|
||||
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
|
||||
|
||||
``PyTorchTrainer`` allows you to run a custom training and validation loops in parallel on each worker, providing a flexible interface similar to using PyTorch natively.
|
||||
This is done via the :ref:`ref-pytorch-operator` interface.
|
||||
|
||||
For both training and validation, there are two granularities that you can provide customization - per epoch and per batch. These correspond to ``train_batch``,
|
||||
``train_epoch``, ``validate``, and ``validate_batch``. Other useful methods to override include ``setup``, ``save`` and ``restore``. You can use these
|
||||
to manage state (like a classifier neural network for calculating inception score, or a heavy tokenizer).
|
||||
|
||||
Providing a custom operator is necessary if creator functions return multiple models, optimizers, or schedulers.
|
||||
|
||||
Below is a partial example of a custom ``TrainingOperator`` that provides a ``train_batch`` implementation for a Deep Convolutional GAN.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
from ray.util.sgd.pytorch import TrainingOperator
|
||||
|
||||
def initialization_hook(runner):
|
||||
class GANOperator(TrainingOperator):
|
||||
def setup(self, config):
|
||||
"""Custom setup for this operator.
|
||||
|
||||
Args:
|
||||
config (dict): Custom configuration value to be passed to
|
||||
all creator and operator constructors. Same as ``self.config``.
|
||||
"""
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
def train_batch(self, batch, batch_info):
|
||||
"""Trains on one batch of data from the data creator.
|
||||
|
||||
Example taken from:
|
||||
https://github.com/eriklindernoren/PyTorch-GAN/blob/
|
||||
a163b82beff3d01688d8315a3fd39080400e7c01/implementations/dcgan/dcgan.py
|
||||
|
||||
Args:
|
||||
batch: One item of the validation iterator.
|
||||
batch_info (dict): Information dict passed in from ``train_epoch``.
|
||||
|
||||
Returns:
|
||||
A dict of metrics. Defaults to "loss" and "num_samples",
|
||||
corresponding to the total number of datapoints in the batch.
|
||||
"""
|
||||
Tensor = torch.cuda.FloatTensor if cuda else torch.FloatTensor
|
||||
discriminator, generator = self.models
|
||||
optimizer_D, optimizer_G = self.optimizers
|
||||
|
||||
# Adversarial ground truths
|
||||
valid = Variable(Tensor(batch.shape[0], 1).fill_(1.0), requires_grad=False)
|
||||
fake = Variable(Tensor(batch.shape[0], 1).fill_(0.0), requires_grad=False)
|
||||
|
||||
# Configure input
|
||||
real_imgs = Variable(batch.type(Tensor))
|
||||
|
||||
# -----------------
|
||||
# Train Generator
|
||||
# -----------------
|
||||
|
||||
optimizer_G.zero_grad()
|
||||
|
||||
# Sample noise as generator input
|
||||
z = Variable(Tensor(np.random.normal(0, 1, (
|
||||
batch.shape[0], self.config["latent_dim"]))))
|
||||
|
||||
# Generate a batch of images
|
||||
gen_imgs = generator(z)
|
||||
|
||||
# Loss measures generator's ability to fool the discriminator
|
||||
g_loss = adversarial_loss(discriminator(gen_imgs), valid)
|
||||
|
||||
g_loss.backward()
|
||||
optimizer_G.step()
|
||||
|
||||
# ---------------------
|
||||
# Train Discriminator
|
||||
# ---------------------
|
||||
|
||||
optimizer_D.zero_grad()
|
||||
|
||||
# Measure discriminator's ability to classify real from generated samples
|
||||
real_loss = adversarial_loss(discriminator(real_imgs), valid)
|
||||
fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake)
|
||||
d_loss = (real_loss + fake_loss) / 2
|
||||
|
||||
d_loss.backward()
|
||||
optimizer_D.step()
|
||||
|
||||
return {
|
||||
"loss_g": g_loss.item(),
|
||||
"loss_d": d_loss.item(),
|
||||
"num_samples": imgs.shape[0]
|
||||
}
|
||||
|
||||
trainer = PyTorchTrainer(
|
||||
model_creator,
|
||||
data_creator,
|
||||
optimizer_creator,
|
||||
nn.BCELoss,
|
||||
training_operator_cls=GANOperator,
|
||||
num_replicas=num_replicas,
|
||||
config=config,
|
||||
use_gpu=True,
|
||||
batch_size=128)
|
||||
|
||||
for i in range(5):
|
||||
stats = trainer.train()
|
||||
print(stats)
|
||||
|
||||
See the `DCGAN example <https://github.com/ray-project/ray/blob/master/python/ray/util/sgd/pytorch/examples/dcgan.py>`__ for an end to end example. It constructs two models and two optimizers and uses a custom training operator to provide a non-standard training loop.
|
||||
|
||||
|
||||
Initialization Functions
|
||||
------------------------
|
||||
|
||||
Use the ``initialization_hook`` parameter to initialize state on each worker process when they are started. This is useful when setting an environment variable:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def initialization_hook():
|
||||
print("NCCL DEBUG SET")
|
||||
# Need this for avoiding a connection restart issue
|
||||
os.environ["NCCL_SOCKET_IFNAME"] = "^docker0,lo"
|
||||
@@ -193,8 +299,8 @@ and ``trainer.load``, which wraps the relevant ``torch.save`` and ``torch.load``
|
||||
trainer_2.restore(checkpoint_path)
|
||||
|
||||
|
||||
Exporting a model for inference
|
||||
-------------------------------
|
||||
Retrieving the model
|
||||
--------------------
|
||||
|
||||
The trained torch model can be extracted for use within the same Python program with ``trainer.get_model()``. This will load the state dictionary of the model(s).
|
||||
|
||||
@@ -242,22 +348,23 @@ To specify particular parameters for ``amp.initialize``, you can use the ``apex_
|
||||
}
|
||||
)
|
||||
|
||||
Note that if using a custom training function, you will need to manage loss scaling manually.
|
||||
Note that if using a custom training operator (:ref:`raysgd-custom-training`), you will need to manage loss scaling manually.
|
||||
|
||||
|
||||
Distributed Multi-node Training
|
||||
-------------------------------
|
||||
|
||||
You can scale out your training onto multiple nodes without making any modifications to your training code. To train across a cluster, simply make sure that the Ray cluster is started.
|
||||
You can scale your training to multiple nodes without making any modifications to your training code.
|
||||
|
||||
You can start a Ray cluster `via the Ray cluster launcher <autoscaling.html>`_ or `manually <using-ray-on-a-cluster.html>`_.
|
||||
To train across a cluster, first make sure that the Ray cluster is started. You can start a Ray cluster `via the Ray cluster launcher <autoscaling.html>`_ or `manually <using-ray-on-a-cluster.html>`_.
|
||||
|
||||
.. code-block:: bash
|
||||
Then, in your program, you'll need to connect to this cluster via ``ray.init``:
|
||||
|
||||
ray up CLUSTER.yaml
|
||||
ray submit train.py --args="--address='auto'"
|
||||
.. code-block:: python
|
||||
|
||||
Then, within ``train.py`` you can scale up the number of workers seamlessly across multiple nodes:
|
||||
ray.init(address="auto") # or a specific redis address of the form "ip-address:port"
|
||||
|
||||
After connecting, you can scale up the number of workers seamlessly across multiple nodes:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@@ -266,7 +373,10 @@ Then, within ``train.py`` you can scale up the number of workers seamlessly acro
|
||||
data_creator,
|
||||
optimizer_creator,
|
||||
loss_creator=nn.MSELoss,
|
||||
num_replicas=100)
|
||||
num_replicas=100
|
||||
)
|
||||
trainer.train()
|
||||
model = trainer.get_model()
|
||||
|
||||
|
||||
Advanced: Fault Tolerance
|
||||
@@ -310,22 +420,37 @@ Advanced: Hyperparameter Tuning
|
||||
Simultaneous Multi-model Training
|
||||
---------------------------------
|
||||
|
||||
In certain scenarios such as training GANs, you may want to use multiple models in the training loop. You can do this in the ``PyTorchTrainer`` by allowing the ``model_creator``, ``optimizer_creator``, and ``scheduler_creator`` to return multiple values.
|
||||
|
||||
If multiple models, optimizers, or schedulers are returned, you will need to provide a custom training function (and custom validation function if you plan to call ``validate``).
|
||||
In certain scenarios, such as training GANs, you may want to use multiple models in the training loop. You can do this in the ``PyTorchTrainer`` by allowing the ``model_creator``, ``optimizer_creator``, and ``scheduler_creator`` to return multiple values. Provide a custom TrainingOperator (:ref:`raysgd-custom-training`) to train across multiple models.
|
||||
|
||||
You can see the `DCGAN script <https://github.com/ray-project/ray/blob/master/python/ray/util/sgd/pytorch/examples/dcgan.py>`_ for an end-to-end example.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from ray.util.sgd.pytorch import PyTorchTrainer, TrainingOperator
|
||||
|
||||
def train(*, model=None, criterion=None, optimizer=None, dataloader=None):
|
||||
model.train()
|
||||
train_loss = 0
|
||||
correct = 0
|
||||
total = 0
|
||||
for batch_idx, (inputs, targets) in enumerate(dataloader):
|
||||
optimizer.zero_grad()
|
||||
outputs = model(inputs)
|
||||
loss = criterion(outputs, targets)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
train_loss += loss.item()
|
||||
_, predicted = outputs.max(1)
|
||||
total += targets.size(0)
|
||||
correct += predicted.eq(targets).sum().item()
|
||||
return {
|
||||
"accuracy": correct / total,
|
||||
"train_loss": train_loss / (batch_idx + 1)
|
||||
}
|
||||
|
||||
def model_creator(config):
|
||||
netD = Discriminator()
|
||||
netD.apply(weights_init)
|
||||
|
||||
netG = Generator()
|
||||
netG.apply(weights_init)
|
||||
return netD, netG
|
||||
|
||||
return Discriminator(), Generator()
|
||||
|
||||
def optimizer_creator(models, config):
|
||||
net_d, net_g = models
|
||||
@@ -335,125 +460,27 @@ You can see the `DCGAN script <https://github.com/ray-project/ray/blob/master/py
|
||||
net_g.parameters(), lr=config.get("lr", 0.01), betas=(0.5, 0.999))
|
||||
return discriminator_opt, generator_opt
|
||||
|
||||
|
||||
def custom_train(models, dataloader, criterion, optimizers, config):
|
||||
result = {}
|
||||
for i, (model, optimizer) in enumerate(zip(models, optimizers)):
|
||||
result["model_{}".format(i)] = train(model, dataloader, criterion,
|
||||
optimizer, config)
|
||||
return result
|
||||
class CustomOperator(TrainingOperator):
|
||||
def train_epoch(self, dataloader, info):
|
||||
result = {}
|
||||
for i, (model, optimizer) in enumerate(
|
||||
zip(self.models, self.optimizers)):
|
||||
result["model_{}".format(i)] = train(
|
||||
model=model,
|
||||
criterion=self.criterion,
|
||||
optimizer=optimizer,
|
||||
dataloader=dataloader)
|
||||
return result
|
||||
|
||||
trainer = PyTorchTrainer(
|
||||
model_creator,
|
||||
data_creator,
|
||||
optimizer_creator,
|
||||
loss_creator=nn.BCELoss,
|
||||
train_function=custom_train)
|
||||
training_operator_cls=CustomOperator)
|
||||
|
||||
.. _raysgd-custom-training:
|
||||
trainer.train()
|
||||
|
||||
Custom Training and Validation Functions
|
||||
----------------------------------------
|
||||
|
||||
``PyTorchTrainer`` allows you to run a custom training and validation step in parallel on each worker, providing a flexibility similar to using PyTorch natively. This is done via the ``train_function`` and ``validation_function`` parameters.
|
||||
|
||||
Note that this is needed if the model creator returns multiple models, optimizers, or schedulers.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
def train(config, model, train_iterator, criterion, optimizer, scheduler=None):
|
||||
"""Runs one standard training pass over the train_iterator.
|
||||
|
||||
Raises:
|
||||
ValueError if multiple models/optimizers/schedulers are provided. You
|
||||
are expected to have a custom training function if you wish
|
||||
to use multiple models/optimizers/schedulers.
|
||||
|
||||
Args:
|
||||
config: (dict): A user configuration provided into the Trainer
|
||||
constructor.
|
||||
model: The model(s) as created by the model_creator.
|
||||
train_iterator: An iterator created from the DataLoader which
|
||||
wraps the provided Dataset.
|
||||
criterion: The loss object created by the loss_creator.
|
||||
optimizer: The torch.optim.Optimizer(s) object
|
||||
as created by the optimizer_creator.
|
||||
scheduler (optional): The torch.optim.lr_scheduler(s) object
|
||||
as created by the scheduler_creator.
|
||||
|
||||
Returns:
|
||||
A dict of metrics from training.
|
||||
"""
|
||||
|
||||
netD, netG = models
|
||||
optimD, optimG = optimizers
|
||||
real_label = 1
|
||||
fake_label = 0
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
for i, data in enumerate(dataloader, 0):
|
||||
netD.zero_grad()
|
||||
real_cpu = data[0].to(device)
|
||||
b_size = real_cpu.size(0)
|
||||
label = torch.full((b_size, ), real_label, device=device)
|
||||
output = netD(real_cpu).view(-1)
|
||||
errD_real = criterion(output, label)
|
||||
errD_real.backward()
|
||||
|
||||
noise = torch.randn(b_size, latent_vector_size, 1, 1, device=device)
|
||||
fake = netG(noise)
|
||||
label.fill_(fake_label)
|
||||
output = netD(fake.detach()).view(-1)
|
||||
errD_fake = criterion(output, label)
|
||||
errD_fake.backward()
|
||||
errD = errD_real + errD_fake
|
||||
optimD.step()
|
||||
|
||||
netG.zero_grad()
|
||||
label.fill_(real_label)
|
||||
output = netD(fake).view(-1)
|
||||
errG = criterion(output, label)
|
||||
errG.backward()
|
||||
optimG.step()
|
||||
|
||||
is_score, is_std = inception_score(fake)
|
||||
|
||||
return {
|
||||
"loss_g": errG.item(),
|
||||
"loss_d": errD.item(),
|
||||
"inception": is_score
|
||||
}
|
||||
|
||||
|
||||
def custom_validate(config, model, val_iterator, criterion, scheduler=None):
|
||||
"""Runs one standard validation pass over the val_iterator.
|
||||
|
||||
Args:
|
||||
config: (dict): A user configuration provided into the Trainer
|
||||
constructor.
|
||||
model: The model(s) as created by the model_creator.
|
||||
train_iterator: An iterator created from the DataLoader which
|
||||
wraps the provided Dataset.
|
||||
criterion: The loss object created by the loss_creator.
|
||||
scheduler (optional): The torch.optim.lr_scheduler object(s)
|
||||
as created by the scheduler_creator.
|
||||
|
||||
Returns:
|
||||
A dict of metrics from the evaluation.
|
||||
"""
|
||||
...
|
||||
return {"validation_accuracy": 0.5}
|
||||
|
||||
|
||||
trainer = PyTorchTrainer(
|
||||
model_creator,
|
||||
data_creator,
|
||||
optimizer_creator,
|
||||
nn.BCELoss,
|
||||
train_function=train,
|
||||
validation_function=custom_validate,
|
||||
...
|
||||
)
|
||||
|
||||
Feature Requests
|
||||
----------------
|
||||
@@ -472,9 +499,7 @@ to contribute an example, feel free to create a `pull request here <https://gith
|
||||
Simple example of using Ray's PyTorchTrainer.
|
||||
|
||||
- `CIFAR10 example <https://github.com/ray-project/ray/blob/master/python/ray/util/sgd/pytorch/examples/cifar_pytorch_example.py>`__:
|
||||
Training a ResNet18 model on CIFAR10. It uses a custom training
|
||||
function, a custom validation function, and custom initialization code for each worker.
|
||||
Training a ResNet18 model on CIFAR10.
|
||||
|
||||
- `DCGAN example <https://github.com/ray-project/ray/blob/master/python/ray/util/sgd/pytorch/examples/dcgan.py>`__:
|
||||
Training a Deep Convolutional GAN on MNIST. It constructs
|
||||
two models and two optimizers and uses a custom training and validation function.
|
||||
Training a Deep Convolutional GAN on MNIST. It constructs two models and two optimizers and uses a custom training operator.
|
||||
|
||||
@@ -11,6 +11,14 @@ PyTorchTrainer
|
||||
|
||||
.. automethod:: __init__
|
||||
|
||||
.. _ref-pytorch-operator:
|
||||
|
||||
PyTorch TrainingOperator
|
||||
------------------------
|
||||
|
||||
.. autoclass:: ray.util.sgd.pytorch.TrainingOperator
|
||||
:members:
|
||||
|
||||
|
||||
PyTorchTrainable
|
||||
----------------
|
||||
|
||||
Reference in New Issue
Block a user