wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-25 13:30:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Examples updated with actors. (#358)

Browse Source 
* Updated examples with actors

* Small changes, and convert documentation from MD to RST.
This commit is contained in:
Wapaul1
2017-03-11 15:30:31 -08:00
committed by Robert Nishihara
parent 3b7788bf88
commit b1cb48159a
11 changed files with 631 additions and 684 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/example-lbfgs.md
-150
View File
@@ -1,150 +0,0 @@
# Batch L-BFGS
This document provides a walkthrough of the L-BFGS example. To run the
application, first install these dependencies.
- SciPy
- [TensorFlow](https://www.tensorflow.org/)
Then from the directory `ray/examples/lbfgs/` run the following.
```
python driver.py
```
Optimization is at the heart of many machine learning algorithms. Much of
machine learning involves specifying a loss function and finding the parameters
that minimize the loss. If we can compute the gradient of the loss function,
then we can apply a variety of gradient-based optimization algorithms. L-BFGS is
one such algorithm. It is a quasi-Newton method that uses gradient information
to approximate the inverse Hessian of the loss function in a computationally
efficient manner.
## The serial version
First we load the data in batches. Here, each element in `batches` is a tuple
whose first component is a batch of `100` images and whose second component is a
batch of the `100` corresponding labels. For simplicity, we use TensorFlow's
built in methods for loading the data.
```python
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
batch_size = 100
num_batches = mnist.train.num_examples // batch_size
batches = [mnist.train.next_batch(batch_size) for _ in range(num_batches)]
```
Now, suppose we have defined a function which takes a set of model parameters
`theta` and a batch of data (both images and labels) and computes the loss for
that choice of model parameters on that batch of data. Similarly, suppose we've
also defined a function that takes the same arguments and computes the gradient
of the loss for that choice of model parameters.
```python
def loss(theta, xs, ys):
# compute the loss on a batch of data
return loss
def grad(theta, xs, ys):
# compute the gradient on a batch of data
return grad
def full_loss(theta):
# compute the loss on the full data set
return sum([loss(theta, xs, ys) for (xs, ys) in batches])
def full_grad(theta):
# compute the gradient on the full data set
return sum([grad(theta, xs, ys) for (xs, ys) in batches])
```
Since we are working with a small dataset, we don't actually need to separate
these methods into the part that operates on a batch and the part that operates
on the full dataset, but doing so will make the distributed version clearer.
Now, if we wish to optimize the loss function using L-BFGS, we simply plug these
functions, along with an initial choice of model parameters, into
`scipy.optimize.fmin_l_bfgs_b`.
```python
theta_init = 1e-2 * np.random.normal(size=dim)
result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, fprime=full_grad)
```
## The distributed version
In this example, the computation of the gradient itself can be done in parallel
on a number of workers or machines.
First, let's turn the data into a collection of remote objects.
```python
batch_ids = [(ray.put(xs), ray.put(ys)) for (xs, ys) in batches]
```
We can load the data on the driver and distribute it this way because MNIST
easily fits on a single machine. However, for larger data sets, we will need to
use remote functions to distribute the loading of the data.
Now, lets turn `loss` and `grad` into remote functions.
```python
@ray.remote
def loss(theta, xs, ys):
# compute the loss
return loss
@ray.remote
def grad(theta, xs, ys):
# compute the gradient
return grad
```
The only difference is that we added the `@ray.remote` decorator.
Now, it is easy to speed up the computation of the full loss and the full
gradient.
```python
def full_loss(theta):
theta_id = ray.put(theta)
loss_ids = [loss.remote(theta_id, xs_id, ys_id) for (xs_id, ys_id) in batch_ids]
return sum(ray.get(loss_ids))
def full_grad(theta):
theta_id = ray.put(theta)
grad_ids = [grad.remote(theta_id, xs_id, ys_id) for (xs_id, ys_id) in batch_ids]
return sum(ray.get(grad_ids)).astype("float64") # This conversion is necessary for use with fmin_l_bfgs_b.
```
Note that we turn `theta` into a remote object with the line `theta_id =
ray.put(theta)` before passing it into the remote functions. If we had written
```python
[loss.remote(theta, xs_id, ys_id) for (xs_id, ys_id) in batch_ids]
```
instead of
```python
theta_id = ray.put(theta)
[loss.remote(theta_id, xs_id, ys_id) for (xs_id, ys_id) in batch_ids]
```
then each task that got sent to the scheduler (one for every element of
`batch_ids`) would have had a copy of `theta` serialized inside of it. Since
`theta` here consists of the parameters of a potentially large model, this is
inefficient. *Large objects should be passed by object ID to remote functions
and not by value*.
We use remote functions and remote objects internally in the implementation of
`full_loss` and `full_grad`, but the user-facing behavior of these methods is
identical to the behavior in the serial version.
We can now optimize the objective with the same function call as before.
```python
theta_init = 1e-2 * np.random.normal(size=dim)
result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, fprime=full_grad)
```
doc/source/example-lbfgs.rst
+156
View File
@@ -0,0 +1,156 @@
Batch L-BFGS
============
This document provides a walkthrough of the L-BFGS example. To run the
application, first install these dependencies.
.. code-block:: bash
pip install tensorflow
pip install scipy
Then you can run the example as follows.
.. code-block:: bash
python ray/examples/lbfgs/driver.py
Optimization is at the heart of many machine learning algorithms. Much of
machine learning involves specifying a loss function and finding the parameters
that minimize the loss. If we can compute the gradient of the loss function,
then we can apply a variety of gradient-based optimization algorithms. L-BFGS is
one such algorithm. It is a quasi-Newton method that uses gradient information
to approximate the inverse Hessian of the loss function in a computationally
efficient manner.
The serial version
------------------
First we load the data in batches. Here, each element in ``batches`` is a tuple
whose first component is a batch of ``100`` images and whose second component is a
batch of the ``100`` corresponding labels. For simplicity, we use TensorFlow's
built in methods for loading the data.
.. code-block:: python
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
batch_size = 100
num_batches = mnist.train.num_examples // batch_size
batches = [mnist.train.next_batch(batch_size) for _ in range(num_batches)]
Now, suppose we have defined a function which takes a set of model parameters
``theta`` and a batch of data (both images and labels) and computes the loss for
that choice of model parameters on that batch of data. Similarly, suppose we've
also defined a function that takes the same arguments and computes the gradient
of the loss for that choice of model parameters.
.. code-block:: python
def loss(theta, xs, ys):
# compute the loss on a batch of data
return loss
def grad(theta, xs, ys):
# compute the gradient on a batch of data
return grad
def full_loss(theta):
# compute the loss on the full data set
return sum([loss(theta, xs, ys) for (xs, ys) in batches])
def full_grad(theta):
# compute the gradient on the full data set
return sum([grad(theta, xs, ys) for (xs, ys) in batches])
Since we are working with a small dataset, we don't actually need to separate
these methods into the part that operates on a batch and the part that operates
on the full dataset, but doing so will make the distributed version clearer.
Now, if we wish to optimize the loss function using L-BFGS, we simply plug these
functions, along with an initial choice of model parameters, into
``scipy.optimize.fmin_l_bfgs_b``.
.. code-block:: python
theta_init = 1e-2 * np.random.normal(size=dim)
result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, fprime=full_grad)
The distributed version
-----------------------
In this example, the computation of the gradient itself can be done in parallel
on a number of workers or machines.
First, let's turn the data into a collection of remote objects.
.. code-block:: python
batch_ids = [(ray.put(xs), ray.put(ys)) for (xs, ys) in batches]
We can load the data on the driver and distribute it this way because MNIST
easily fits on a single machine. However, for larger data sets, we will need to
use remote functions to distribute the loading of the data.
Now, lets turn ``loss`` and ``grad`` into methods of an actor that will contain our network.
.. code-block:: python
class Network(object):
def __init__():
# Initialize network.
def loss(theta, xs, ys):
# compute the loss
return loss
def grad(theta, xs, ys):
# compute the gradient
return grad
Now, it is easy to speed up the computation of the full loss and the full
gradient.
.. code-block:: python
def full_loss(theta):
theta_id = ray.put(theta)
loss_ids = [actor.loss(theta_id) for actor in actors]
return sum(ray.get(loss_ids))
def full_grad(theta):
theta_id = ray.put(theta)
grad_ids = [actor.grad(theta_id) for actor in actors]
return sum(ray.get(grad_ids)).astype("float64") # This conversion is necessary for use with fmin_l_bfgs_b.
Note that we turn ``theta`` into a remote object with the line ``theta_id =
ray.put(theta)`` before passing it into the remote functions. If we had written
.. code-block:: python
[actor.loss(theta_id) for actor in actors]
instead of
.. code-block:: python
theta_id = ray.put(theta)
[actor.loss(theta_id) for actor in actors]
then each task that got sent to the scheduler (one for every element of
``batch_ids``) would have had a copy of ``theta`` serialized inside of it. Since
``theta`` here consists of the parameters of a potentially large model, this is
inefficient. *Large objects should be passed by object ID to remote functions
and not by value*.
We use remote actors and remote objects internally in the implementation of
``full_loss`` and ``full_grad``, but the user-facing behavior of these methods is
identical to the behavior in the serial version.
We can now optimize the objective with the same function call as before.
.. code-block:: python
theta_init = 1e-2 * np.random.normal(size=dim)
result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, fprime=full_grad)
doc/source/example-rl-pong.md
-113
View File
@@ -1,113 +0,0 @@
# Learning to Play Pong
In this example, we'll be training a neural network to play Pong using the
OpenAI Gym. This application is adapted, with minimal modifications, from Andrej
Karpathy's
[code](https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5) (see
the accompanying [blog post](http://karpathy.github.io/2016/05/31/rl/)). To run
the application, first install this dependency.
- [Gym](https://gym.openai.com/)
Then from the directory `ray/examples/rl_pong/` run the following.
```
python driver.py
```
## The distributed version
At the core of [Andrej's
code](https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5), a
neural network is used to define a "policy" for playing Pong (that is, a
function that chooses an action given a state). In the loop, the network
repeatedly plays games of Pong and records a gradient from each game. Every ten
games, the gradients are combined together and used to update the network.
This example is easy to parallelize because the network can play ten games in
parallel and no information needs to be shared between the games. We define a
remote function `compute_gradient`, which plays a game of pong and returns an
estimate of the gradient. Below is a simplified pseudocode version of this
function.
```python
@ray.remote(num_return_vals=2)
def compute_gradient(model):
# Retrieve the game environment.
env = ray.env.env
# Reset the game.
observation = env.reset()
while not done:
# Choose an action using policy_forward.
# Take the action and observe the new state of the world.
# Compute a gradient using policy_backward. Return the gradient and reward.
return gradient, reward_sum
```
Calling this remote function inside of a for loop, we launch multiple tasks to
perform rollouts and compute gradients. If we have at least ten worker
processes, then these tasks will all be executed in parallel.
```python
model_id = ray.put(model)
grads, reward_sums = [], []
# Launch tasks to compute gradients from multiple rollouts in parallel.
for i in range(10):
grad_id, reward_sum_id = compute_gradient.remote(model_id)
grads.append(grad_id)
reward_sums.append(reward_sum_id)
```
### Reusing the Gym environment
Workers are long-running Python processes, and though we'd like to think of
workers as being stateless, sometimes it's important to have a variable that
gets shared between different tasks on the same worker (perhaps because it is
expensive to initialize the variable).
In this example, we'd like each worker to have access to a Pong environment. The
Pong environment has state that gets mutated by the task, and this state is
shared between tasks that run on the same worker, so there is some danger that
the output of the overall program will depend on which tasks are scheduled on
which workers. This can be avoided if the state of the Pong environment is reset
between tasks.
To accomplish this, the user must mark the Pong environment as an environment
variable. This is done by providing a method for initializing the gym, and
storing it in `ray.env`.
```python
# Function for initializing the gym environment.
def env_initializer():
return gym.make("Pong-v0")
# Create an environment variable for the gym environment.
ray.env.env = ray.EnvironmentVariable(env_initializer)
```
A remote task can then call `ray.env.env` to retrieve the variable.
By default, whenever a task uses the `ray.env.env` variable, the worker
that the task was scheduled on will rerun the initialization code
`env_initializer` after the task has finished so that state will not leak
between the tasks.
However, sometimes the initialization code is expensive, and there may be a
faster way to reinitialize the variable (or maybe no reinitialization is needed
at all). In these cases, the user can provide a custom **reinitializer**, which
gets run after any task that uses the variable.
```python
# Function for initializing the gym environment.
def env_initializer():
return gym.make("Pong-v0")
# Function for reinitializing the gym environment in order to guarantee that
# the state of the game is reset after each remote task.
def env_reinitializer(env):
env.reset()
return env
# Create an environment variable for the gym environment.
ray.env.env = ray.EnvironmentVariable(env_initializer, env_reinitializer)
```
doc/source/example-rl-pong.rst
+70
View File
@@ -0,0 +1,70 @@
Learning to Play Pong
=====================
In this example, we'll be training a neural network to play Pong using the
OpenAI Gym. This application is adapted, with minimal modifications, from Andrej
Karpathy's
[code](https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5) (see
the accompanying [blog post](http://karpathy.github.io/2016/05/31/rl/)). To run
the application, first install some dependencies.
.. code-block:: bash
pip install gym[atari]
Then you can run the example as follows.
.. code-block:: bash
python ray/examples/rl_pong/driver.py
The distributed version
-----------------------
At the core of [Andrej's
code](https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5), a
neural network is used to define a "policy" for playing Pong (that is, a
function that chooses an action given a state). In the loop, the network
repeatedly plays games of Pong and records a gradient from each game. Every ten
games, the gradients are combined together and used to update the network.
This example is easy to parallelize because the network can play ten games in
parallel and no information needs to be shared between the games.
We define an **actor** for the Pong environment, which includes a method for
performing a rollout and computing a gradient update. Below is pseudocode for
the actor.
.. code-block:: python
@ray.actor
class PongEnv(object):
def __init__(self):
self.env = gym.make("Pong-v0")
def compute_gradient(self, model):
# Reset the game.
observation = self.env.reset()
while not done:
# Choose an action using policy_forward.
# Take the action and observe the new state of the world.
# Compute a gradient using policy_backward. Return the gradient and reward.
return [gradient, reward_sum]
We then create a number of actors, so that we can perform rollouts in parallel.
.. code-block:: python
actors = [PongEnv() for _ in range(batch_size)]
Calling this remote function inside of a for loop, we launch multiple tasks to
perform rollouts and compute gradients in parallel.
.. code-block:: python
model_id = ray.put(model)
actions = []
# Launch tasks to compute gradients from multiple rollouts in parallel.
for i in range(batch_size):
action_id = actors[i].compute_gradient(model_id)
actions.append(action_id)
doc/source/index.rst
+3 -3
View File
@@ -30,9 +30,9 @@ learning and reinforcement learning applications.*
example-policy-gradient.rst
example-resnet.rst
example-a3c.rst
example-lbfgs.md
example-rl-pong.md
using-ray-with-tensorflow.md
example-lbfgs.rst
example-rl-pong.rst
using-ray-with-tensorflow.rst
.. toctree::
:maxdepth: 1
doc/source/using-ray-with-tensorflow.md
-320
View File
@@ -1,320 +0,0 @@
# Using Ray with TensorFlow
This document describes best practices for using Ray with TensorFlow. If you are
training a deep network in the distributed setting, you may need to ship your
deep network between processes (or machines). For example, you may update your
model on one machine and then use that model to compute a gradient on another
machine. However, shipping the model is not always straightforward.
For example, a straightforward attempt to pickle a TensorFlow graph gives mixed
results. Some examples fail, and some succeed (but produce very large strings).
The results are similar with other pickling libraries as well.
Furthermore, creating a TensorFlow graph can take tens of seconds, and so
serializing a graph and recreating it in another process will be inefficient.
The better solution is to create the same TensorFlow graph on each worker once
at the beginning and then to ship only the weights between the workers.
Suppose we have a simple network definition (this one is modified from the
TensorFlow documentation).
```python
import tensorflow as tf
import numpy as np
x_data = tf.placeholder(tf.float32, shape=[100])
y_data = tf.placeholder(tf.float32, shape=[100])
w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
b = tf.Variable(tf.zeros([1]))
y = w * x_data + b
loss = tf.reduce_mean(tf.square(y - y_data))
optimizer = tf.train.GradientDescentOptimizer(0.5)
grads = optimizer.compute_gradients(loss)
train = optimizer.apply_gradients(grads)
init = tf.global_variables_initializer()
sess = tf.Session()
```
To extract the weights and set the weights, you can use the following helper
method.
```python
import ray
variables = ray.experimental.TensorFlowVariables(loss, sess)
```
The `TensorFlowVariables` object provides methods for getting and setting the
weights as well as collecting all of the variables in the model.
Now we can use these methods to extract the weights, and place them back in the
network as follows.
```python
# First initialize the weights.
sess.run(init)
# Get the weights
weights = variables.get_weights() # Returns a dictionary of numpy arrays
# Set the weights
variables.set_weights(weights)
```
**Note:** If we were to set the weights using the `assign` method like below,
each call to `assign` would add a node to the graph, and the graph would grow
unmanageably large over time.
```python
w.assign(np.zeros(1)) # This adds a node to the graph every time you call it.
b.assign(np.zeros(1)) # This adds a node to the graph every time you call it.
```
## Complete Example
Putting this all together, we would first create the graph on each worker using
environment variables. Within the environment variables, we would use the
`get_weights` and `set_weights` methods of the `TensorFlowVariables` class. We
would then use those methods to ship the weights (as a dictionary of variable
names mapping to tensorflow tensors) between the processes without shipping the
actual TensorFlow graphs, which are much more complex Python objects. Note that
to avoid namespace collision with already created variables on the workers, we
use a separate graph for each network.
```python
import tensorflow as tf
import numpy as np
import ray
ray.init(num_workers=5)
BATCH_SIZE = 100
NUM_BATCHES = 1
NUM_ITERS = 201
def net_vars_initializer():
# Use a separate graph for each network.
with tf.Graph().as_default():
# Seed TensorFlow to make the script deterministic.
tf.set_random_seed(0)
# Define the inputs.
x_data = tf.placeholder(tf.float32, shape=[BATCH_SIZE])
y_data = tf.placeholder(tf.float32, shape=[BATCH_SIZE])
# Define the weights and computation.
w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
b = tf.Variable(tf.zeros([1]))
y = w * x_data + b
# Define the loss.
loss = tf.reduce_mean(tf.square(y - y_data))
optimizer = tf.train.GradientDescentOptimizer(0.5)
grads = optimizer.compute_gradients(loss)
train = optimizer.apply_gradients(grads)
# Define the weight initializer and session.
init = tf.global_variables_initializer()
sess = tf.Session()
# Additional code for setting and getting the weights
variables = ray.experimental.TensorFlowVariables(loss, sess)
# Return all of the data needed to use the network.
return variables, sess, grads, train, loss, x_data, y_data, init
def net_vars_reinitializer(net_vars):
return net_vars
# Define an environment variable for the network variables.
ray.env.net_vars = ray.EnvironmentVariable(net_vars_initializer, net_vars_reinitializer)
# Define a remote function that trains the network for one step and returns the
# new weights.
@ray.remote
def step(weights, x, y):
variables, sess, _, train, _, x_data, y_data, _ = ray.env.net_vars
# Set the weights in the network.
variables.set_weights(weights)
# Do one step of training.
sess.run(train, feed_dict={x_data: x, y_data: y})
# Return the new weights.
return variables.get_weights()
variables, sess, _, train, loss, x_data, y_data, init = ray.env.net_vars
# Initialize the network weights.
sess.run(init)
# Get the weights as a dictionary of numpy arrays.
weights = variables.get_weights()
# Define a remote function for generating fake data.
@ray.remote(num_return_vals=2)
def generate_fake_x_y_data(num_data, seed=0):
# Seed numpy to make the script deterministic.
np.random.seed(seed)
x = np.random.rand(num_data)
y = x * 0.1 + 0.3
return x, y
# Generate some training data.
batch_ids = [generate_fake_x_y_data.remote(BATCH_SIZE, seed=i) for i in range(NUM_BATCHES)]
x_ids = [x_id for x_id, y_id in batch_ids]
y_ids = [y_id for x_id, y_id in batch_ids]
# Generate some test data.
x_test, y_test = ray.get(generate_fake_x_y_data.remote(BATCH_SIZE, seed=NUM_BATCHES))
# Do some steps of training.
for iteration in range(NUM_ITERS):
# Put the weights in the object store. This is optional. We could instead pass
# the variable weights directly into step.remote, in which case it would be
# placed in the object store under the hood. However, in that case multiple
# copies of the weights would be put in the object store, so this approach is
# more efficient.
weights_id = ray.put(weights)
# Call the remote function multiple times in parallel.
new_weights_ids = [step.remote(weights_id, x_ids[i], y_ids[i]) for i in range(NUM_BATCHES)]
# Get all of the weights.
new_weights_list = ray.get(new_weights_ids)
# Add up all the different weights. Each element of new_weights_list is a dict
# of weights, and we want to add up these dicts component wise using the keys
# of the first dict.
weights = {variable: sum(weight_dict[variable] for weight_dict in new_weights_list) / NUM_BATCHES for variable in new_weights_list[0]}
# Print the current weights. They should converge to roughly to the values 0.1
# and 0.3 used in generate_fake_x_y_data.
if iteration % 20 == 0:
print("Iteration {}: weights are {}".format(iteration, weights))
```
## How to Train in Parallel using Ray
In some cases, you may want to do data-parallel training on your network. We use the network
above to illustrate how to do this in Ray. The only differences are in the remote function
`step` and the driver code.
In the function `step`, we run the grad operation rather than the train operation to get the gradients.
Since Tensorflow pairs the gradients with the variables in a tuple, we extract the gradients to avoid
needless computation.
### Extracting numerical gradients
Code like the following can be used in a remote function to compute numerical gradients.
```python
x_values = [1] * 100
y_values = [2] * 100
numerical_grads = sess.run([grad[0] for grad in grads], feed_dict={x_data: x_values, y_data: y_values})
```
### Using the returned gradients to train the network
By pairing the symbolic gradients with the numerical gradients in a feed_dict, we can update the network.
```python
# We can feed the gradient values in using the associated symbolic gradient
# operation defined in tensorflow.
feed_dict = {grad[0]: numerical_grad for (grad, numerical_grad) in zip(grads, numerical_grads)}
sess.run(train, feed_dict=feed_dict)
```
You can then run `variables.get_weights()` to see the updated weights of the network.
For reference, the full code is below:
```python
import tensorflow as tf
import numpy as np
import ray
ray.init(num_workers=5)
BATCH_SIZE = 100
NUM_BATCHES = 1
NUM_ITERS = 201
def net_vars_initializer():
# Use a separate graph for each network.
with tf.Graph().as_default():
# Seed TensorFlow to make the script deterministic.
tf.set_random_seed(0)
# Define the inputs.
x_data = tf.placeholder(tf.float32, shape=[BATCH_SIZE])
y_data = tf.placeholder(tf.float32, shape=[BATCH_SIZE])
# Define the weights and computation.
w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
b = tf.Variable(tf.zeros([1]))
y = w * x_data + b
# Define the loss.
loss = tf.reduce_mean(tf.square(y - y_data))
optimizer = tf.train.GradientDescentOptimizer(0.5)
grads = optimizer.compute_gradients(loss)
train = optimizer.apply_gradients(grads)
# Define the weight initializer and session.
init = tf.global_variables_initializer()
sess = tf.Session()
# Additional code for setting and getting the weights
variables = ray.experimental.TensorFlowVariables(loss, sess)
# Return all of the data needed to use the network.
return variables, sess, grads, train, loss, x_data, y_data, init
def net_vars_reinitializer(net_vars):
return net_vars
# Define an environment variable for the network variables.
ray.env.net_vars = ray.EnvironmentVariable(net_vars_initializer, net_vars_reinitializer)
# Define a remote function that trains the network for one step and returns the
# new weights.
@ray.remote
def step(weights, x, y):
variables, sess, grads, _, _, x_data, y_data, _ = ray.env.net_vars
# Set the weights in the network.
variables.set_weights(weights)
# Do one step of training. We only need the actual gradients so we filter over the list.
actual_grads = sess.run([grad[0] for grad in grads], feed_dict={x_data: x, y_data: y})
return actual_grads
variables, sess, grads, train, loss, x_data, y_data, init = ray.env.net_vars
# Initialize the network weights.
sess.run(init)
# Get the weights as a dictionary of numpy arrays.
weights = variables.get_weights()
# Define a remote function for generating fake data.
@ray.remote(num_return_vals=2)
def generate_fake_x_y_data(num_data, seed=0):
# Seed numpy to make the script deterministic.
np.random.seed(seed)
x = np.random.rand(num_data)
y = x * 0.1 + 0.3
return x, y
# Generate some training data.
batch_ids = [generate_fake_x_y_data.remote(BATCH_SIZE, seed=i) for i in range(NUM_BATCHES)]
x_ids = [x_id for x_id, y_id in batch_ids]
y_ids = [y_id for x_id, y_id in batch_ids]
# Generate some test data.
x_test, y_test = ray.get(generate_fake_x_y_data.remote(BATCH_SIZE, seed=NUM_BATCHES))
# Do some steps of training.
for iteration in range(NUM_ITERS):
# Put the weights in the object store. This is optional. We could instead pass
# the variable weights directly into step.remote, in which case it would be
# placed in the object store under the hood. However, in that case multiple
# copies of the weights would be put in the object store, so this approach is
# more efficient.
weights_id = ray.put(weights)
# Call the remote function multiple times in parallel.
gradients_ids = [step.remote(weights_id, x_ids[i], y_ids[i]) for i in range(NUM_BATCHES)]
# Get all of the weights.
gradients_list = ray.get(gradients_ids)
# Take the mean of the different gradients. Each element of gradients_list is a list
# of gradients, and we want to take the mean of each one.
mean_grads = [sum([gradients[i] for gradients in gradients_list]) / len(gradients_list) for i in range(len(gradients_list[0]))]
feed_dict = {grad[0]: mean_grad for (grad, mean_grad) in zip(grads, mean_grads)}
sess.run(train, feed_dict=feed_dict)
weights = variables.get_weights()
# Print the current weights. They should converge to roughly to the values 0.1
# and 0.3 used in generate_fake_x_y_data.
if iteration % 20 == 0:
print("Iteration {}: weights are {}".format(iteration, weights))
```
doc/source/using-ray-with-tensorflow.rst
+322
View File
@@ -0,0 +1,322 @@
Using Ray with TensorFlow
=========================
This document describes best practices for using Ray with TensorFlow.
To see more involved examples using TensorFlow, take a look at `hyperparameter optimization`_,
`A3C`_, `ResNet`_, `Policy Gradients`_, and `LBFGS`_.
.. _`hyperparameter optimization`: http://ray.readthedocs.io/en/latest/example-hyperopt.html
.. _`A3C`: http://ray.readthedocs.io/en/latest/example-a3c.html
.. _`ResNet`: http://ray.readthedocs.io/en/latest/example-resnet.html
.. _`Policy Gradients`: http://ray.readthedocs.io/en/latest/example-policy-gradient.html
.. _`LBFGS`: http://ray.readthedocs.io/en/latest/example-lbfgs.html
If you are training a deep network in the distributed setting, you may need to
ship your deep network between processes (or machines). For example, you may
update your model on one machine and then use that model to compute a gradient
on another machine. However, shipping the model is not always straightforward.
For example, a straightforward attempt to pickle a TensorFlow graph gives mixed
results. Some examples fail, and some succeed (but produce very large strings).
The results are similar with other pickling libraries as well.
Furthermore, creating a TensorFlow graph can take tens of seconds, and so
serializing a graph and recreating it in another process will be inefficient.
The better solution is to create the same TensorFlow graph on each worker once
at the beginning and then to ship only the weights between the workers.
Suppose we have a simple network definition (this one is modified from the
TensorFlow documentation).
.. code-block:: python
import tensorflow as tf
import numpy as np
x_data = tf.placeholder(tf.float32, shape=[100])
y_data = tf.placeholder(tf.float32, shape=[100])
w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
b = tf.Variable(tf.zeros([1]))
y = w * x_data + b
loss = tf.reduce_mean(tf.square(y - y_data))
optimizer = tf.train.GradientDescentOptimizer(0.5)
grads = optimizer.compute_gradients(loss)
train = optimizer.apply_gradients(grads)
init = tf.global_variables_initializer()
sess = tf.Session()
To extract the weights and set the weights, you can use the following helper
method.
.. code-block:: python
import ray
variables = ray.experimental.TensorFlowVariables(loss, sess)
The ``TensorFlowVariables`` object provides methods for getting and setting the
weights as well as collecting all of the variables in the model.
Now we can use these methods to extract the weights, and place them back in the
network as follows.
.. code-block:: python
# First initialize the weights.
sess.run(init)
# Get the weights
weights = variables.get_weights() # Returns a dictionary of numpy arrays
# Set the weights
variables.set_weights(weights)
**Note:** If we were to set the weights using the ``assign`` method like below,
each call to ``assign`` would add a node to the graph, and the graph would grow
unmanageably large over time.
.. code-block:: python
w.assign(np.zeros(1)) # This adds a node to the graph every time you call it.
b.assign(np.zeros(1)) # This adds a node to the graph every time you call it.
Complete Example
----------------
Putting this all together, we would first embed the graph in an actor. Within
the actor, we would use the ``get_weights`` and ``set_weights`` methods of the
``TensorFlowVariables`` class. We would then use those methods to ship the weights
(as a dictionary of variable names mapping to numpy arrays) between the
processes without shipping the actual TensorFlow graphs, which are much more
complex Python objects.
.. code-block:: python
import tensorflow as tf
import numpy as np
import ray
ray.init()
BATCH_SIZE = 100
NUM_BATCHES = 1
NUM_ITERS = 201
class Network(object):
def __init__(self, x, y):
# Seed TensorFlow to make the script deterministic.
tf.set_random_seed(0)
# Define the inputs.
self.x_data = tf.constant(x, dtype=tf.float32)
self.y_data = tf.constant(y, dtype=tf.float32)
# Define the weights and computation.
w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
b = tf.Variable(tf.zeros([1]))
y = w * self.x_data + b
# Define the loss.
self.loss = tf.reduce_mean(tf.square(y - self.y_data))
optimizer = tf.train.GradientDescentOptimizer(0.5)
self.grads = optimizer.compute_gradients(self.loss)
self.train = optimizer.apply_gradients(self.grads)
# Define the weight initializer and session.
init = tf.global_variables_initializer()
self.sess = tf.Session()
# Additional code for setting and getting the weights
self.variables = ray.experimental.TensorFlowVariables(self.loss, self.sess)
# Return all of the data needed to use the network.
self.sess.run(init)
# Define a remote function that trains the network for one step and returns the
# new weights.
def step(self, weights):
# Set the weights in the network.
self.variables.set_weights(weights)
# Do one step of training.
self.sess.run(self.train)
# Return the new weights.
return self.variables.get_weights()
def get_weights(self):
return self.variables.get_weights()
# Define a remote function for generating fake data.
@ray.remote(num_return_vals=2)
def generate_fake_x_y_data(num_data, seed=0):
# Seed numpy to make the script deterministic.
np.random.seed(seed)
x = np.random.rand(num_data)
y = x * 0.1 + 0.3
return x, y
# Generate some training data.
batch_ids = [generate_fake_x_y_data.remote(BATCH_SIZE, seed=i) for i in range(NUM_BATCHES)]
x_ids = [x_id for x_id, y_id in batch_ids]
y_ids = [y_id for x_id, y_id in batch_ids]
# Generate some test data.
x_test, y_test = ray.get(generate_fake_x_y_data.remote(BATCH_SIZE, seed=NUM_BATCHES))
# Create actors to store the networks.
remote_network = ray.actor(Network)
actor_list = [remote_network(x_ids[i], y_ids[i]) for i in range(NUM_BATCHES)]
# Get initial weights of some actor.
weights = ray.get(actor_list[0].get_weights())
# Do some steps of training.
for iteration in range(NUM_ITERS):
# Put the weights in the object store. This is optional. We could instead pass
# the variable weights directly into step.remote, in which case it would be
# placed in the object store under the hood. However, in that case multiple
# copies of the weights would be put in the object store, so this approach is
# more efficient.
weights_id = ray.put(weights)
# Call the remote function multiple times in parallel.
new_weights_ids = [actor.step(weights_id) for actor in actor_list]
# Get all of the weights.
new_weights_list = ray.get(new_weights_ids)
# Add up all the different weights. Each element of new_weights_list is a dict
# of weights, and we want to add up these dicts component wise using the keys
# of the first dict.
weights = {variable: sum(weight_dict[variable] for weight_dict in new_weights_list) / NUM_BATCHES for variable in new_weights_list[0]}
# Print the current weights. They should converge to roughly to the values 0.1
# and 0.3 used in generate_fake_x_y_data.
if iteration % 20 == 0:
print("Iteration {}: weights are {}".format(iteration, weights))
How to Train in Parallel using Ray
----------------------------------
In some cases, you may want to do data-parallel training on your network. We use the network
above to illustrate how to do this in Ray. The only differences are in the remote function
``step`` and the driver code.
In the function ``step``, we run the grad operation rather than the train operation to get the gradients.
Since Tensorflow pairs the gradients with the variables in a tuple, we extract the gradients to avoid
needless computation.
Extracting numerical gradients
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Code like the following can be used in a remote function to compute numerical gradients.
.. code-block:: python
x_values = [1] * 100
y_values = [2] * 100
numerical_grads = sess.run([grad[0] for grad in grads], feed_dict={x_data: x_values, y_data: y_values})
Using the returned gradients to train the network
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
By pairing the symbolic gradients with the numerical gradients in a feed_dict, we can update the network.
.. code-block:: python
# We can feed the gradient values in using the associated symbolic gradient
# operation defined in tensorflow.
feed_dict = {grad[0]: numerical_grad for (grad, numerical_grad) in zip(grads, numerical_grads)}
sess.run(train, feed_dict=feed_dict)
You can then run ``variables.get_weights()`` to see the updated weights of the network.
For reference, the full code is below:
.. code-block:: python
import tensorflow as tf
import numpy as np
import ray
ray.init()
BATCH_SIZE = 100
NUM_BATCHES = 1
NUM_ITERS = 201
class Network(object):
def __init__(self, x, y):
# Seed TensorFlow to make the script deterministic.
tf.set_random_seed(0)
# Define the inputs.
x_data = tf.constant(x, dtype=tf.float32)
y_data = tf.constant(y, dtype=tf.float32)
# Define the weights and computation.
w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
b = tf.Variable(tf.zeros([1]))
y = w * x_data + b
# Define the loss.
self.loss = tf.reduce_mean(tf.square(y - y_data))
optimizer = tf.train.GradientDescentOptimizer(0.5)
self.grads = optimizer.compute_gradients(self.loss)
self.train = optimizer.apply_gradients(self.grads)
# Define the weight initializer and session.
init = tf.global_variables_initializer()
self.sess = tf.Session()
# Additional code for setting and getting the weights
self.variables = ray.experimental.TensorFlowVariables(self.loss, self.sess)
# Return all of the data needed to use the network.
self.sess.run(init)
# Define a remote function that trains the network for one step and returns the
# new weights.
def step(self, weights):
# Set the weights in the network.
self.variables.set_weights(weights)
# Do one step of training. We only need the actual gradients so we filter over the list.
actual_grads = self.sess.run([grad[0] for grad in self.grads])
return actual_grads
def get_weights(self):
return self.variables.get_weights()
# Define a remote function for generating fake data.
@ray.remote(num_return_vals=2)
def generate_fake_x_y_data(num_data, seed=0):
# Seed numpy to make the script deterministic.
np.random.seed(seed)
x = np.random.rand(num_data)
y = x * 0.1 + 0.3
return x, y
# Generate some training data.
batch_ids = [generate_fake_x_y_data.remote(BATCH_SIZE, seed=i) for i in range(NUM_BATCHES)]
x_ids = [x_id for x_id, y_id in batch_ids]
y_ids = [y_id for x_id, y_id in batch_ids]
# Generate some test data.
x_test, y_test = ray.get(generate_fake_x_y_data.remote(BATCH_SIZE, seed=NUM_BATCHES))
# Create actors to store the networks.
remote_network = ray.actor(Network)
actor_list = [remote_network(x_ids[i], y_ids[i]) for i in range(NUM_BATCHES)]
local_network = Network(x_test, y_test)
# Get initial weights of local network.
weights = local_network.get_weights()
# Do some steps of training.
for iteration in range(NUM_ITERS):
# Put the weights in the object store. This is optional. We could instead pass
# the variable weights directly into step.remote, in which case it would be
# placed in the object store under the hood. However, in that case multiple
# copies of the weights would be put in the object store, so this approach is
# more efficient.
weights_id = ray.put(weights)
# Call the remote function multiple times in parallel.
gradients_ids = [actor.step(weights_id) for actor in actor_list]
# Get all of the weights.
gradients_list = ray.get(gradients_ids)
# Take the mean of the different gradients. Each element of gradients_list is a list
# of gradients, and we want to take the mean of each one.
mean_grads = [sum([gradients[i] for gradients in gradients_list]) / len(gradients_list) for i in range(len(gradients_list[0]))]
feed_dict = {grad[0]: mean_grad for (grad, mean_grad) in zip(local_network.grads, mean_grads)}
local_network.sess.run(local_network.train, feed_dict=feed_dict)
weights = local_network.get_weights()
# Print the current weights. They should converge to roughly to the values 0.1
# and 0.3 used in generate_fake_x_y_data.
if iteration % 20 == 0:
print("Iteration {}: weights are {}".format(iteration, weights))
examples/a3c/driver.py
+3 -3
View File
@@ -76,6 +76,6 @@ if __name__ == '__main__':
if gym.__version__[:3] == '0.8':
raise Exception("This example currently does not work with gym==0.8.0. "
"Please downgrade to gym==0.7.4.");
NW = int(sys.argv[1])
ray.init(num_workers=NW, num_cpus=NW)
train(NW)
num_workers = int(sys.argv[1])
ray.init(num_cpus=num_workers)
train(num_workers)
examples/lbfgs/driver.py
+31 -40
View File
@@ -3,10 +3,10 @@ from __future__ import division
from __future__ import print_function
import ray
import numpy as np
import scipy.optimize
import tensorflow as tf
import os
from tensorflow.examples.tutorials.mnist import input_data
@@ -17,8 +17,8 @@ class LinearModel(object):
are set via self.variables.set_weights.
Example:
net = LinearModel([10,10])
weights = [np.random.normal(size=[10,10]), np.random.normal(size=[10])]
net = LinearModel([10, 10])
weights = [np.random.normal(size=[10, 10]), np.random.normal(size=[10])]
variable_names = [v.name for v in net.variables]
net.variables.set_weights(dict(zip(variable_names, weights)))
@@ -60,53 +60,45 @@ class LinearModel(object):
"""Computes the gradients of the network."""
return self.sess.run(self.cross_entropy_grads, feed_dict={self.x: xs, self.y_: ys})
def net_initialization():
with tf.Graph().as_default():
return LinearModel([784,10])
@ray.actor
class NetActor(object):
def __init__(self, xs, ys):
os.environ["CUDA_VISIBLE_DEVICES"] = ""
with tf.device("/cpu:0"):
self.net = LinearModel([784, 10])
self.xs = xs
self.ys = ys
# By default, when an environment variable is used by a remote function, the
# initialization code will be rerun at the end of the remote task to ensure
# that the state of the variable is not changed by the remote task. However,
# the initialization code may be expensive. This case is one example, because
# a TensorFlow network is constructed. In this case, we pass in a special
# reinitialization function which gets run instead of the original
# initialization code. As users, if we pass in custom reinitialization code,
# we must ensure that no state is leaked between tasks.
def net_reinitialization(net):
return net
# Compute the loss on a batch of data.
def loss(self, theta):
net = self.net
net.variables.set_flat(theta)
return net.loss(self.xs, self.ys)
# Register the network with Ray and create an environment variable for it.
ray.env.net = ray.EnvironmentVariable(net_initialization, net_reinitialization)
# Compute the gradient of the loss on a batch of data.
def grad(self, theta):
net = self.net
net.variables.set_flat(theta)
gradients = net.grad(self.xs, self.ys)
return np.concatenate([g.flatten() for g in gradients])
# Compute the loss on a batch of data.
@ray.remote
def loss(theta, xs, ys):
net = ray.env.net
net.variables.set_flat(theta)
return net.loss(xs,ys)
# Compute the gradient of the loss on a batch of data.
@ray.remote
def grad(theta, xs, ys):
net = ray.env.net
net.variables.set_flat(theta)
gradients = net.grad(xs, ys)
return np.concatenate([g.flatten() for g in gradients])
def get_flat_size(self):
return self.net.variables.get_flat_size()
# Compute the loss on the entire dataset.
def full_loss(theta):
theta_id = ray.put(theta)
loss_ids = [loss.remote(theta_id, xs_id, ys_id) for (xs_id, ys_id) in batch_ids]
loss_ids = [actor.loss(theta_id) for actor in actors]
return sum(ray.get(loss_ids))
# Compute the gradient of the loss on the entire dataset.
def full_grad(theta):
theta_id = ray.put(theta)
grad_ids = [grad.remote(theta_id, xs_id, ys_id) for (xs_id, ys_id) in batch_ids]
grad_ids = [actor.grad(theta_id) for actor in actors]
return sum(ray.get(grad_ids)).astype("float64") # This conversion is necessary for use with fmin_l_bfgs_b.
if __name__ == "__main__":
ray.init(num_workers=10)
ray.init(redirect_output=True)
# From the perspective of scipy.optimize.fmin_l_bfgs_b, full_loss is simply a
# function which takes some parameters theta, and computes a loss. Similarly,
@@ -121,14 +113,13 @@ if __name__ == "__main__":
# Load the mnist data and turn the data into remote objects.
print("Downloading the MNIST dataset. This may take a minute.")
mnist = input_data.read_data_sets("MNIST_data", one_hot=True)
batch_size = 100
num_batches = mnist.train.num_examples // batch_size
num_batches = 10
batch_size = mnist.train.num_examples // num_batches
batches = [mnist.train.next_batch(batch_size) for _ in range(num_batches)]
print("Putting MNIST in the object store.")
batch_ids = [(ray.put(xs), ray.put(ys)) for (xs, ys) in batches]
actors = [NetActor(xs, ys) for (xs, ys) in batches]
# Initialize the weights for the network to the vector of all zeros.
dim = ray.env.net.variables.get_flat_size()
dim = ray.get(actors[0].get_flat_size())
theta_init = 1e-2 * np.random.normal(size=dim)
# Use L-BFGS to minimize the loss function.
examples/resnet/resnet_main.py
+2 -1
View File
@@ -21,6 +21,7 @@ tf.app.flags.DEFINE_string('eval_data_path', '',
'Filepattern for eval data')
tf.app.flags.DEFINE_string('num_gpus', 0, 'Number of gpus to run with')
use_gpu = 1 if int(FLAGS.num_gpus) > 0 else 0
@ray.remote(num_return_vals=4)
def get_data(path, size):
os.environ['CUDA_VISIBLE_DEVICES'] = ''
@@ -136,7 +137,7 @@ class ResNetTestActor(object):
def train():
"""Training loop."""
num_gpus = int(FLAGS.num_gpus)
ray.init(num_workers=2, num_gpus=num_gpus)
ray.init(num_gpus=num_gpus)
train_data = get_data.remote(FLAGS.train_data_path, 50000)
test_data = get_data.remote(FLAGS.eval_data_path, 10000)
if num_gpus > 0:
examples/rl_pong/driver.py
+44 -54
View File
@@ -19,19 +19,6 @@ decay_rate = 0.99 # decay factor for RMSProp leaky sum of grad^2
D = 80 * 80 # input dimensionality: 80x80 grid
# Function for initializing the gym environment.
def env_initializer():
return gym.make("Pong-v0")
# Function for reinitializing the gym environment in order to guarantee that
# the state of the game is reset after each remote task.
def env_reinitializer(env):
env.reset()
return env
# Create an environment variable for the gym environment.
ray.env.env = ray.EnvironmentVariable(env_initializer, env_reinitializer)
def sigmoid(x):
return 1.0 / (1.0 + np.exp(-x)) # sigmoid "squashing" function to interval [0,1]
@@ -69,48 +56,52 @@ def policy_backward(eph, epx, epdlogp, model):
dW1 = np.dot(dh.T, epx)
return {"W1": dW1, "W2": dW2}
@ray.remote(num_return_vals=2)
def compute_gradient(model):
env = ray.env.env
observation = env.reset()
prev_x = None # used in computing the difference frame
xs, hs, dlogps, drs = [], [], [], []
reward_sum = 0
done = False
while not done:
cur_x = preprocess(observation)
x = cur_x - prev_x if prev_x is not None else np.zeros(D)
prev_x = cur_x
@ray.actor
class PongEnv(object):
def __init__(self):
self.env = gym.make("Pong-v0")
aprob, h = policy_forward(x, model)
action = 2 if np.random.uniform() < aprob else 3 # roll the dice!
def compute_gradient(self, model):
# Reset the game.
observation = self.env.reset()
prev_x = None # used in computing the difference frame
xs, hs, dlogps, drs = [], [], [], []
reward_sum = 0
done = False
while not done:
cur_x = preprocess(observation)
x = cur_x - prev_x if prev_x is not None else np.zeros(D)
prev_x = cur_x
xs.append(x) # observation
hs.append(h) # hidden state
y = 1 if action == 2 else 0 # a "fake label"
dlogps.append(y - aprob) # grad that encourages the action that was taken to be taken (see http://cs231n.github.io/neural-networks-2/#losses if confused)
aprob, h = policy_forward(x, model)
action = 2 if np.random.uniform() < aprob else 3 # roll the dice!
observation, reward, done, info = env.step(action)
reward_sum += reward
xs.append(x) # observation
hs.append(h) # hidden state
y = 1 if action == 2 else 0 # a "fake label"
dlogps.append(y - aprob) # grad that encourages the action that was taken to be taken (see http://cs231n.github.io/neural-networks-2/#losses if confused)
drs.append(reward) # record reward (has to be done after we call step() to get reward for previous action)
observation, reward, done, info = self.env.step(action)
reward_sum += reward
epx = np.vstack(xs)
eph = np.vstack(hs)
epdlogp = np.vstack(dlogps)
epr = np.vstack(drs)
xs, hs, dlogps, drs = [], [], [], [] # reset array memory
drs.append(reward) # record reward (has to be done after we call step() to get reward for previous action)
# compute the discounted reward backwards through time
discounted_epr = discount_rewards(epr)
# standardize the rewards to be unit normal (helps control the gradient estimator variance)
discounted_epr -= np.mean(discounted_epr)
discounted_epr /= np.std(discounted_epr)
epdlogp *= discounted_epr # modulate the gradient with advantage (PG magic happens right here.)
return policy_backward(eph, epx, epdlogp, model), reward_sum
epx = np.vstack(xs)
eph = np.vstack(hs)
epdlogp = np.vstack(dlogps)
epr = np.vstack(drs)
xs, hs, dlogps, drs = [], [], [], [] # reset array memory
# compute the discounted reward backwards through time
discounted_epr = discount_rewards(epr)
# standardize the rewards to be unit normal (helps control the gradient estimator variance)
discounted_epr -= np.mean(discounted_epr)
discounted_epr /= np.std(discounted_epr)
epdlogp *= discounted_epr # modulate the gradient with advantage (PG magic happens right here.)
return policy_backward(eph, epx, epdlogp, model), reward_sum
if __name__ == "__main__":
ray.init(num_workers=10)
ray.init()
# Run the reinforcement learning
running_reward = None
@@ -120,18 +111,17 @@ if __name__ == "__main__":
model["W2"] = np.random.randn(H) / np.sqrt(H)
grad_buffer = {k: np.zeros_like(v) for k, v in model.items()} # update buffers that add up gradients over a batch
rmsprop_cache = {k: np.zeros_like(v) for k, v in model.items()} # rmsprop memory
actors = [PongEnv() for _ in range(batch_size)]
while True:
model_id = ray.put(model)
grads, reward_sums = [], []
actions = []
# Launch tasks to compute gradients from multiple rollouts in parallel.
for i in range(batch_size):
grad_id, reward_sum_id = compute_gradient.remote(model_id)
grads.append(grad_id)
reward_sums.append(reward_sum_id)
action_id = actors[i].compute_gradient(model_id)
actions.append(action_id)
for i in range(batch_size):
grad = ray.get(grads[i])
reward_sum = ray.get(reward_sums[i])
action_id, actions = ray.wait(actions)
grad, reward_sum = ray.get(action_id[0])
for k in model: grad_buffer[k] += grad[k] # accumulate grad over batch
running_reward = reward_sum if running_reward is None else running_reward * 0.99 + reward_sum * 0.01
print("Batch {}. episode reward total was {}. running mean: {}".format(batch_num, reward_sum, running_reward))