From b1cb48159a96a5740f8008c38abe7ec488849ef8 Mon Sep 17 00:00:00 2001 From: Wapaul1 Date: Sat, 11 Mar 2017 15:30:31 -0800 Subject: [PATCH] Examples updated with actors. (#358) * Updated examples with actors * Small changes, and convert documentation from MD to RST. --- doc/source/example-lbfgs.md | 150 ----------- doc/source/example-lbfgs.rst | 156 +++++++++++ doc/source/example-rl-pong.md | 113 -------- doc/source/example-rl-pong.rst | 70 +++++ doc/source/index.rst | 6 +- doc/source/using-ray-with-tensorflow.md | 320 ---------------------- doc/source/using-ray-with-tensorflow.rst | 322 +++++++++++++++++++++++ examples/a3c/driver.py | 6 +- examples/lbfgs/driver.py | 71 +++-- examples/resnet/resnet_main.py | 3 +- examples/rl_pong/driver.py | 98 ++++--- 11 files changed, 631 insertions(+), 684 deletions(-) delete mode 100644 doc/source/example-lbfgs.md create mode 100644 doc/source/example-lbfgs.rst delete mode 100644 doc/source/example-rl-pong.md create mode 100644 doc/source/example-rl-pong.rst delete mode 100644 doc/source/using-ray-with-tensorflow.md create mode 100644 doc/source/using-ray-with-tensorflow.rst diff --git a/doc/source/example-lbfgs.md b/doc/source/example-lbfgs.md deleted file mode 100644 index b9659d148..000000000 --- a/doc/source/example-lbfgs.md +++ /dev/null @@ -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) -``` diff --git a/doc/source/example-lbfgs.rst b/doc/source/example-lbfgs.rst new file mode 100644 index 000000000..1d4703718 --- /dev/null +++ b/doc/source/example-lbfgs.rst @@ -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) diff --git a/doc/source/example-rl-pong.md b/doc/source/example-rl-pong.md deleted file mode 100644 index 6bf3de5a8..000000000 --- a/doc/source/example-rl-pong.md +++ /dev/null @@ -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) -``` diff --git a/doc/source/example-rl-pong.rst b/doc/source/example-rl-pong.rst new file mode 100644 index 000000000..3ccba7a11 --- /dev/null +++ b/doc/source/example-rl-pong.rst @@ -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) diff --git a/doc/source/index.rst b/doc/source/index.rst index 1cbbf2141..824d374f3 100644 --- a/doc/source/index.rst +++ b/doc/source/index.rst @@ -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 diff --git a/doc/source/using-ray-with-tensorflow.md b/doc/source/using-ray-with-tensorflow.md deleted file mode 100644 index 7e7dc2fa3..000000000 --- a/doc/source/using-ray-with-tensorflow.md +++ /dev/null @@ -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)) -``` diff --git a/doc/source/using-ray-with-tensorflow.rst b/doc/source/using-ray-with-tensorflow.rst new file mode 100644 index 000000000..0a2bdefcd --- /dev/null +++ b/doc/source/using-ray-with-tensorflow.rst @@ -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)) diff --git a/examples/a3c/driver.py b/examples/a3c/driver.py index 9193e1193..06a1058f1 100644 --- a/examples/a3c/driver.py +++ b/examples/a3c/driver.py @@ -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) diff --git a/examples/lbfgs/driver.py b/examples/lbfgs/driver.py index 4427e6f35..68639ca37 100644 --- a/examples/lbfgs/driver.py +++ b/examples/lbfgs/driver.py @@ -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. diff --git a/examples/resnet/resnet_main.py b/examples/resnet/resnet_main.py index 2ffe46793..a6fc42b3d 100644 --- a/examples/resnet/resnet_main.py +++ b/examples/resnet/resnet_main.py @@ -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: diff --git a/examples/rl_pong/driver.py b/examples/rl_pong/driver.py index 37c30c5ac..af2e6002e 100644 --- a/examples/rl_pong/driver.py +++ b/examples/rl_pong/driver.py @@ -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))