wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-21 12:50:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

[sgd] fp16 (apex) and scheduler support + move examples page (#7061)

Browse Source 
* Init fp16

* fp16 and schedulers

* scheduler linking and fp16

* to fp16

* loss scaling and documentation

* more documentation

* add tests, refactor config

* moredocs

* more docs

* fix logo, add test mode, add fp16 flag

* fix tests

* fix scheduler

* fix apex

* improve safety

* fix tests

* fix tests

* remove pin memory default

* rm

* fix

* Update doc/examples/doc_code/raysgd_torch_signatures.py

* fix

* migrate changes from other PR

* ok thanks

* pass

* signatures

* lint'

* Update python/ray/experimental/sgd/pytorch/utils.py

* Apply suggestions from code review

Co-Authored-By: Edward Oakes <ed.nmi.oakes@gmail.com>

* should address most comments

* comments

* fix this ci

* fix tests'

* testmode

Co-authored-by: Edward Oakes <ed.nmi.oakes@gmail.com>
This commit is contained in:
Richard Liaw
2020-02-16 19:04:08 -08:00
committed by GitHub
co-authored by Edward Oakes
parent f0e62d733f
commit 94e2fcea2e
17 changed files with 874 additions and 364 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/index.rst
-1
View File
@@ -281,7 +281,6 @@ Getting Involved
raysgd/raysgd.rst
raysgd/raysgd_pytorch.rst
raysgd/raysgd_pytorch_examples.rst
raysgd/raysgd_tensorflow.rst
raysgd/raysgd_ref.rst
doc/source/raysgd/raysgd.rst
+3 -5
View File
@@ -1,10 +1,6 @@
RaySGD: Distributed Deep Learning
=================================
.. tip:: Get in touch with us if you're using or considering using `RaySGD <https://forms.gle/26EMwdahdgm7Lscy9>`_!
.. warning:: This is still an experimental API and is subject to change in the near future.
.. image:: raysgdlogo.png
:scale: 20%
:align: center
@@ -17,6 +13,8 @@ The main features are:
- **Composability**: RaySGD is built on top of the Ray Actor API, enabling seamless integration with existing Ray applications such as RLlib, Tune, and Ray.Serve.
- **Scale up and down**: Start on single CPU. Scale up to multi-node, multi-CPU, or multi-GPU clusters by changing 2 lines of code.
.. tip:: We need your feedback! RaySGD is currently early in its development, and we're hoping to get feedback from people using or considering it. We'd love `to get in touch <https://forms.gle/26EMwdahdgm7Lscy9>`_!
Getting Started
---------------
@@ -53,7 +51,7 @@ You can start a ``PyTorchTrainer`` with the following:
model_creator,
data_creator,
optimizer_creator,
loss_creator,
loss_creator=nn.MSELoss,
num_replicas=2,
use_gpu=True,
batch_size=512,
doc/source/raysgd/raysgd_pytorch.rst
+198 -115
View File
@@ -1,135 +1,124 @@
RaySGD Pytorch
==============
.. warning:: This is still an experimental API and is subject to change in the near future.
.. tip:: Get in touch with us if you're using or considering using `RaySGD <https://forms.gle/26EMwdahdgm7Lscy9>`_!
.. image:: raysgd-pytorch.svg
:align: center
The RaySGD ``PyTorchTrainer`` simplifies distributed model training for PyTorch. The ``PyTorchTrainer`` is a wrapper around ``torch.distributed.launch`` with a Python API to easily incorporate distributed training into a larger Python application, as opposed to needing to execute training outside of Python. For end to end examples, see :ref:`raysgd-pytorch-example`.
The RaySGD ``PyTorchTrainer`` simplifies distributed model training for PyTorch. The ``PyTorchTrainer`` is a wrapper around ``torch.distributed.launch`` with a Python API to easily incorporate distributed training into a larger Python application, as opposed to needing to wrap your training code in bash scripts.
Under the hood, ``PytorchTrainer`` will create *replicas* of your model (controlled by ``num_replicas``), each of which is managed by a Ray actor.
.. image:: raysgd-actors.svg
:align: center
For end to end examples leveraging RaySGD PyTorchTrainer, jump to :ref:`raysgd-pytorch-examples`.
Setting up training
-------------------
The ``PyTorchTrainer`` can be constructed with functions that wrap components of the training script. Specifically, it needs constructors for the Model, Data, Optimizer, and Loss to create replicated copies across different devices and machines.
.. tip:: We need your feedback! RaySGD is currently early in its development, and we're hoping to get feedback from people using or considering it. We'd love `to get in touch <https://forms.gle/26EMwdahdgm7Lscy9>`_!
For example:
The ``PyTorchTrainer`` can be constructed with functions that wrap components of the training script. Specifically, it requires constructors for the Model, Data, Optimizer, Loss, and ``lr_scheduler`` to create replicated copies across different devices and machines.
.. code-block:: python
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_trainer_start__
:end-before: __torch_trainer_end__
import numpy as np
import torch
import torch.nn as nn
from torch import distributed
The below section covers the expected signatures of creator functions. Jump to :ref:`starting-pytorch-trainer`.
from ray.experimental.sgd import PyTorchTrainer
from ray.experimental.sgd.examples.train_example import LinearDataset
Model Creator
~~~~~~~~~~~~~
def model_creator(config):
"""Constructor function for the model(s) to be optimized.
This is the signature needed for ``PyTorchTrainer(model_creator=...)``.
Note that if multiple models are returned, the same number of optimizers
must be returned. You will also need to provide a custom training
function to specify the optimization procedure for multiple models.
Args:
config (dict): Configuration dictionary passed into ``PyTorchTrainer``.
Returns:
One or more torch.nn.Module objects.
"""
return nn.Linear(1, 1)
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_model_start__
:end-before: __torch_model_end__
def optimizer_creator(models, config):
"""Constructor of the optimizers.
Optimizer Creator
~~~~~~~~~~~~~~~~~
Args:
models: The return values from ``model_creator``. This can be one
or more torch nn modules.
config (dict): Configuration dictionary passed into ``PyTorchTrainer``.
This is the signature needed for ``PyTorchTrainer(optimizer_creator=...)``.
Returns:
One or more Torch optimizer objects. You must return as many optimizers
as you have models.
"""
return torch.optim.SGD(model.parameters(), lr=config.get("lr", 1e-4))
def data_creator(config):
"""Constructs torch.utils.data.Dataset objects.
Note that even though two Dataset objects can be returned,
only one dataset will be used for training.
Args:
config: Configuration dictionary passed into ``PyTorchTrainer``
Returns:
One or Two Dataset objects. If only one Dataset object is provided,
``trainer.validate()`` will throw a ValueError.
"""
return LinearDataset(2, 5), LinearDataset(2, 5, size=400)
def loss_creator(config):
"""Constructs the Torch Loss object.
Note that optionally, you can pass in a Torch Loss constructor directly
into the PyTorchTrainer (i.e., ``PyTorchTrainer(loss_creator=nn.BCELoss, ...))``).
Returns:
Torch Loss object.
"""
return torch.nn.BCELoss()
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_optimizer_start__
:end-before: __torch_optimizer_end__
Before instantiating the trainer, you'll have to start or connect to a Ray cluster:
Data Creator
~~~~~~~~~~~~
.. code-block:: python
This is the signature needed for ``PyTorchTrainer(data_creator=...)``.
ray.init()
# or ray.init(address="auto") if a cluster has been started.
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_data_start__
:end-before: __torch_data_end__
Loss Creator
~~~~~~~~~~~~
This is the signature needed for ``PyTorchTrainer(loss_creator=...)``.
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_loss_start__
:end-before: __torch_loss_end__
Scheduler Creator
~~~~~~~~~~~~~~~~~
Optionally, you can provide a creator function for the learning rate scheduler. This is the signature needed
for ``PyTorchTrainer(scheduler_creator=...)``.
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_scheduler_start__
:end-before: __torch_scheduler_end__
.. _starting-pytorch-trainer:
Putting things together
~~~~~~~~~~~~~~~~~~~~~~~
Before instantiating the trainer, first start or connect to a Ray cluster:
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_ray_start__
:end-before: __torch_ray_end__
Instantiate the trainer object:
.. code-block:: python
from ray.experimental.sgd import PyTorchTrainer
trainer = PyTorchTrainer(
model_creator,
data_creator,
optimizer_creator,
loss_creator=nn.MSELoss,
config={"lr": 0.001})
.. literalinclude:: ../../examples/doc_code/raysgd_torch_signatures.py
:language: python
:start-after: __torch_trainer_start__
:end-before: __torch_trainer_end__
You can also set the number of workers and whether the workers will use GPUs:
.. code-block:: python
:emphasize-lines: 8,9
trainer = PyTorchTrainer(
model_creator,
data_creator,
optimizer_creator,
loss_creator=nn.MSELoss,
scheduler_creator=scheduler_creator,
config={"lr": 0.001},
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.
See the documentation on the PyTorchTrainer here: :ref:`ref-pytorch-trainer`. We'll look at the training APIs next.
Training APIs
-------------
@@ -154,17 +143,19 @@ You can customize the exact function that is called by using a customized traini
Shutting down training
----------------------
After training, you may want to reappropriate the Ray cluster. To release Ray resources obtained by the trainer:
After training, you may want to reappropriate the Ray cluster. To release Ray resources obtained by the Trainer:
.. code-block:: python
trainer.shutdown()
.. note:: Be sure to call ``save`` or ``get_model`` before shutting down.
.. note:: Be sure to call ``trainer.save()`` or ``trainer.get_model()`` before shutting down.
Initialization Functions
------------------------
.. warning:: This is still an experimental API and is subject to change without warning.
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:
.. code-block:: python
@@ -193,15 +184,7 @@ Save and Load
If you want to save or reload the training procedure, you can use ``trainer.save``
and ``trainer.load``, which wraps the relevant ``torch.save`` and ``torch.load`` calls. This should work across a distributed cluster even without a NFS because it takes advantage of Ray's distributed object store.
.. code-block::
trainer_1 = PyTorchTrainer(
model_creator,
data_creator,
optimizer_creator,
loss_creator=nn.MSELoss,
num_replicas=num_replicas)
trainer_1.train()
.. code-block:: python
checkpoint_path = os.path.join(tempfile.mkdtemp(), "checkpoint")
trainer_1.save(checkpoint_path)
@@ -210,7 +193,7 @@ and ``trainer.load``, which wraps the relevant ``torch.save`` and ``torch.load``
model_creator,
data_creator,
optimizer_creator,
loss_creator=lambda config: nn.MSELoss(),
loss_creator=nn.MSELoss,
num_replicas=num_replicas)
trainer_2.restore(checkpoint_path)
@@ -220,10 +203,51 @@ Exporting a model for inference
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).
.. code-block::
.. code-block:: python
trainer.train()
model = trainer.get_model()
model = trainer.get_model() # Returns multiple models if the model_creator does.
Mixed Precision (FP16) Training
-------------------------------
You can enable mixed precision training for PyTorch with the ``use_fp16`` flag. This automatically converts the model(s) and optimizer(s) to train using mixed-precision. This requires NVIDIA ``Apex``, which can be installed from `the NVIDIA/Apex repository <https://github.com/NVIDIA/apex#quick-start>`_:
.. code-block:: python
:emphasize-lines: 7
trainer = PyTorchTrainer(
model_creator,
data_creator,
optimizer_creator,
loss_creator=nn.MSELoss,
num_replicas=4,
use_fp16=True
)
``Apex`` is a Pytorch extension with NVIDIA-maintained utilities to streamline mixed precision and distributed training. When ``use_fp16=True``,
you should not manually cast your model or data to ``.half()``. The flag informs the Trainer to call ``amp.initialize`` on the created models and optimizers and optimize using the scaled loss: ``amp.scale_loss(loss, optimizer)``.
To specify particular parameters for ``amp.initialize``, you can use the ``apex_args`` field for the PyTorchTrainer constructor. Valid arguments can be found on the `Apex documentation <https://nvidia.github.io/apex/amp.html#apex.amp.initialize>`_:
.. code-block:: python
:emphasize-lines: 7-12
trainer = PyTorchTrainer(
model_creator,
data_creator,
optimizer_creator,
loss_creator=nn.MSELoss,
num_replicas=4,
use_fp16=True,
apex_args={
opt_level="O3",
num_losses=2,
verbosity=0
}
)
Note that if using a custom training function, you will need to manage loss scaling manually.
Distributed Multi-node Training
@@ -236,9 +260,9 @@ You can start a Ray cluster `via the Ray cluster launcher <autoscaling.html>`_ o
.. code-block:: bash
ray up CLUSTER.yaml
python train.py --address="auto"
ray submit train.py --args="--address='auto'"
Then, you'll be able to scale up the number of workers seamlessly across multiple nodes:
Then, within ``train.py`` you can scale up the number of workers seamlessly across multiple nodes:
.. code-block:: python
@@ -246,7 +270,7 @@ Then, you'll be able to scale up the number of workers seamlessly across multipl
model_creator,
data_creator,
optimizer_creator,
loss_creator=lambda config: nn.MSELoss(),
loss_creator=nn.MSELoss,
num_replicas=100)
@@ -255,7 +279,7 @@ Advanced: Fault Tolerance
For distributed deep learning, jobs are often run on infrastructure where nodes can be pre-empted frequently (i.e., spot instances in the cloud). To overcome this, RaySGD provides **fault tolerance** features that enable training to continue regardless of node failures.
.. code-block:: bash
.. code-block:: python
trainer.train(max_retries=N)
@@ -273,7 +297,7 @@ Note that we assume the Trainer itself is not on a pre-emptible node. It is curr
Users can set ``checkpoint="auto"`` to always checkpoint the current model before executing a pass over the training dataset.
.. code-block:: bash
.. code-block:: python
trainer.train(max_retries=N, checkpoint="auto")
@@ -281,6 +305,8 @@ Users can set ``checkpoint="auto"`` to always checkpoint the current model befor
Advanced: Hyperparameter Tuning
-------------------------------
.. warning:: This is still an experimental API and is subject to change without warning.
``PyTorchTrainer`` naturally integrates with Tune via the ``PyTorchTrainable`` interface. The same arguments to ``PyTorchTrainer`` should be passed into the ``tune.run(config=...)`` as shown below.
.. literalinclude:: ../../../python/ray/experimental/sgd/pytorch/examples/tune_example.py
@@ -288,12 +314,12 @@ Advanced: Hyperparameter Tuning
:start-after: __torch_tune_example__
Simultaneous Multi-model training
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`` and the ``optimizer_creator`` to return multiple values.
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 are returned, you will need to provide a custom training function (and custom validation function if you plan to call ``validate``).
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``).
You can see the `DCGAN script <https://github.com/ray-project/ray/blob/master/python/ray/experimental/sgd/pytorch/examples/dcgan.py>`_ for an end-to-end example.
@@ -336,24 +362,36 @@ You can see the `DCGAN script <https://github.com/ray-project/ray/blob/master/py
Custom Training and Validation Functions
----------------------------------------
.. warning:: This is still an experimental API and is subject to change in the near future.
``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.
Note that this is needed if the model creator returns multiple models, optimizers, or schedulers.
.. code-block:: python
def train(models, dataloader, criterion, optimizers, config):
"""A custom training function.
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:
models: Output of the model_creator passed into PyTorchTrainer.
data_loader: A dataloader wrapping the training dataset created by the ``data_creator`` passed into PyTorchTrainer.
criterion: The instantiation of the ``loss_creator``.
optimizers: Output of the optimizer_creator passed into PyTorchTrainer.
config: The configuration dictionary passed into PyTorchTrainer.
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 dictionary of values/metrics.
A dict of metrics from training.
"""
netD, netG = models
@@ -396,11 +434,56 @@ Note that this is needed if the model creator returns multiple models.
}
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
----------------
Have features that you'd really like to see in RaySGD? Feel free to `open an issue <https://github.com/ray-project/ray>`_.
.. _raysgd-pytorch-examples:
PyTorchTrainer Examples
-----------------------
Here are some examples of using RaySGD for training PyTorch models. If you'd like
to contribute an example, feel free to create a `pull request here <https://github.com/ray-project/ray/>`_.
- `PyTorch training example <https://github.com/ray-project/ray/blob/master/python/ray/experimental/sgd/pytorch/examples/train_example.py>`__:
Simple example of using Ray's PyTorchTrainer.
- `CIFAR10 example <https://github.com/ray-project/ray/blob/master/python/ray/experimental/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.
- `DCGAN example <https://github.com/ray-project/ray/blob/master/python/ray/experimental/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.
doc/source/raysgd/raysgd_pytorch_examples.rst
-38
View File
@@ -1,38 +0,0 @@
.. _raysgd-pytorch-example:
RaySGD PyTorch Examples
=======================
Here are some examples of using RaySGD for training PyTorch models. If you'd like
to contribute an example, feel free to create a `pull request here <https://github.com/ray-project/ray/pull/>`_.
Toy Example
-----------
Below is an example of using Ray's PyTorchTrainer.
.. literalinclude:: ../../../python/ray/experimental/sgd/pytorch/examples/train_example.py
:language: python
:start-after: __torch_train_example__
CIFAR10 Example
---------------
Below is an example of training a ResNet18 model on CIFAR10. It uses a custom training
function, a custom validation function, and custom initialization code for each worker.
.. literalinclude:: ../../../python/ray/experimental/sgd/pytorch/examples/cifar_pytorch_example.py
:language: python
DCGAN Example
-------------
Below is an example of training a Deep Convolutional GAN on MNIST. It constructs
two models and two optimizers and uses a custom training and validation function.
.. literalinclude:: ../../../python/ray/experimental/sgd/pytorch/examples/dcgan.py
:language: python
doc/source/raysgd/raysgd_tensorflow.rst
+2 -4
View File
@@ -1,14 +1,12 @@
RaySGD TensorFlow
=================
.. warning:: This is still an experimental API and is subject to change in the near future.
.. tip:: Help us make RaySGD better; take this 1 minute `User Survey <https://forms.gle/26EMwdahdgm7Lscy9>`_!
RaySGD's ``TFTrainer`` simplifies distributed model training for Tensorflow. The ``TFTrainer`` is a wrapper around ``MultiWorkerMirroredStrategy`` with a Python API to easily incorporate distributed training into a larger Python application, as opposed to write custom logic of setting environments and starting separate processes.
.. important:: This API has only been tested with TensorFlow2.0rc and is still highly experimental. Please file bug reports if you run into any - thanks!
.. tip:: We need your feedback! RaySGD is currently early in its development, and we're hoping to get feedback from people using or considering it. We'd love `to get in touch <https://forms.gle/26EMwdahdgm7Lscy9>`_!
----------
**With Ray**:
doc/source/raysgd/raysgdlogo.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 278 KiB

After

Width:  |  Height:  |  Size: 280 KiB