Added example for compute grads in ray tutorial (ray-project#238)

* Added example for compute grads in ray

* Added formatting

* Removed need for placeholders in apply gradient

* Streamlined examples

* Fixed docs

* Added formatting

* Removed old references

* Simplified code some

* Addressed comments

* Changes to first code block

* Added test for training and updated code snippets

* Formatting

* Removed mean

* Removed all mention of mean

* Added comments

* Added comments

Author: Wapaul1

doc/using-ray-with-tensorflow.md

@@ -31,7 +31,8 @@ y = w * x_data + b
 
 loss = tf.reduce_mean(tf.square(y - y_data))
 optimizer = tf.train.GradientDescentOptimizer(0.5)
-train = optimizer.minimize(loss)
+grads = optimizer.compute_gradients(loss)
+train = optimizer.apply_gradients(grads)
 
 init = tf.global_variables_initializer()
 sess = tf.Session()
@@ -106,14 +107,15 @@ def net_vars_initializer():
     # Define the loss.
     loss = tf.reduce_mean(tf.square(y - y_data))
     optimizer = tf.train.GradientDescentOptimizer(0.5)
-    train = optimizer.minimize(loss)
+    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, train, loss, x_data, y_data, init
+  return variables, sess, grads, train, loss, x_data, y_data, init
 
 def net_vars_reinitializer(net_vars):
   return net_vars
@@ -125,15 +127,15 @@ ray.env.net_vars = ray.EnvironmentVariable(net_vars_initializer, net_vars_reinit
 # new weights.
 @ray.remote
 def step(weights, x, y):
-  variables, sess, train, _, x_data, y_data, _ = ray.env.net_vars
+  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, _, loss, x_data, y_data, init = ray.env.net_vars
+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.
@@ -176,3 +178,143 @@ for iteration in range(NUM_ITERS):
   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))
+```

python/ray/experimental/tfutils.py

@@ -63,13 +63,13 @@ def __init__(self, loss, sess=None):
     self.variables = OrderedDict()
     for v in [v for v in tf.global_variables() if v.op.node_def.name in variable_names]:
       self.variables[v.op.node_def.name] = v
-    self.assignment_placeholders = dict()
+    self.placeholders = dict()
     self.assignment_nodes = []
 
     # Create new placeholders to put in custom weights.
     for k, var in self.variables.items():
-      self.assignment_placeholders[k] = tf.placeholder(var.value().dtype, var.get_shape().as_list())
-      self.assignment_nodes.append(var.assign(self.assignment_placeholders[k]))
+      self.placeholders[k] = tf.placeholder(var.value().dtype, var.get_shape().as_list())
+      self.assignment_nodes.append(var.assign(self.placeholders[k]))
 
   def set_session(self, sess):
     """Modifies the current session used by the class."""
@@ -92,7 +92,7 @@ def set_flat(self, new_weights):
     self._check_sess()
     shapes = [v.get_shape().as_list() for v in self.variables.values()]
     arrays = unflatten(new_weights, shapes)
-    placeholders = [self.assignment_placeholders[k] for k, v in self.variables.items()]
+    placeholders = [self.placeholders[k] for k, v in self.variables.items()]
     self.sess.run(self.assignment_nodes, feed_dict=dict(zip(placeholders,arrays)))
 
   def get_weights(self):
@@ -103,4 +103,4 @@ def get_weights(self):
   def set_weights(self, new_weights):
     """Sets the weights to new_weights."""
     self._check_sess()
-    self.sess.run(self.assignment_nodes, feed_dict={self.assignment_placeholders[name]: value for (name, value) in new_weights.items()})
+    self.sess.run(self.assignment_nodes, feed_dict={self.placeholders[name]: value for (name, value) in new_weights.items()})

test/tensorflow_test.py

@@ -39,8 +39,10 @@ def train_vars_initializer():
     loss, init, x_data, y_data = make_linear_network()
     sess = tf.Session()
     variables = ray.experimental.TensorFlowVariables(loss, sess)
-    grad = tf.gradients(loss, list(variables.variables.values()))
-  return variables, init, sess, grad, [x_data, y_data]
+    optimizer = tf.train.GradientDescentOptimizer(0.9)
+    grads = optimizer.compute_gradients(loss)
+    train = optimizer.apply_gradients(grads)
+  return loss, variables, init, sess, grads, train, [x_data, y_data]
 
 
 class TensorFlowTest(unittest.TestCase):
@@ -200,16 +202,42 @@ def testRemoteTrainingStep(self):
 
     @ray.remote
     def training_step(weights):
-      variables, _, sess, grad, placeholders = ray.env.net
+      _, variables, _, sess, grads, _, placeholders = ray.env.net
       variables.set_weights(weights)
-      return sess.run(grad, feed_dict=dict(zip(placeholders, [[1]*100]*2)))
+      return sess.run([grad[0] for grad in grads], feed_dict=dict(zip(placeholders, [[1]*100]*2)))
 
-    variables, init, sess, _, _ = ray.env.net
+    _, variables, init, sess, _, _, _ = ray.env.net
 
     sess.run(init)
     ray.get(training_step.remote(variables.get_weights()))
 
     ray.worker.cleanup()
 
+
+  def testRemoteTrainingLoss(self):
+    ray.init(num_workers=2)
+
+    ray.env.net = ray.EnvironmentVariable(train_vars_initializer, net_vars_reinitializer)
+
+    @ray.remote
+    def training_step(weights):
+      _, variables, _, sess, grads, _, placeholders = ray.env.net
+      variables.set_weights(weights)
+      return sess.run([grad[0] for grad in grads], feed_dict=dict(zip(placeholders, [[1]*100, [2]*100])))
+
+    loss, variables, init, sess, grads, train, placeholders = ray.env.net
+
+    sess.run(init)
+    before_acc = sess.run(loss, feed_dict=dict(zip(placeholders, [[2]*100, [4]*100])))
+
+    for _ in range(3):
+      gradients_list = ray.get([training_step.remote(variables.get_weights()) for _ in range(2)])
+      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)
+    after_acc = sess.run(loss, feed_dict=dict(zip(placeholders, [[2]*100, [4]*100])))
+    self.assertTrue(before_acc < after_acc)
+    ray.worker.cleanup()
+
 if __name__ == "__main__":
   unittest.main(verbosity=2)