Cleanup setting and getting of tensorflow weights. (#385)

* Cleanup setting and getting of tensorflow weights.

* Add documentation for using TensorFlow.

* Group get_weights and set_weights in a function.

* Update readme.

README.md

@@ -3,7 +3,7 @@
 [![Build Status](https://travis-ci.org/amplab/ray.svg?branch=master)](https://travis-ci.org/amplab/ray)
 
 Ray is an experimental distributed extension of Python. It is under development
-and not ready for general use.
+and not ready to be used.
 
 The goal of Ray is to make it easy to write machine learning applications that
 run on a cluster while providing the development and debugging experience of
@@ -51,8 +51,9 @@ estimate of pi (waiting until the computation has finished if necessary).
 
 - Installation on [Ubuntu](doc/install-on-ubuntu.md), [Mac OS X](doc/install-on-macosx.md), [Windows](doc/install-on-windows.md), [Docker](doc/install-on-docker.md)
 - [Tutorial](doc/tutorial.md)
-- [About the System](doc/about-the-system.md)
-- [Using Ray on a Cluster](doc/using-ray-on-a-cluster.md)
+- Documentation
+  - [Using Ray with TensorFlow](doc/using-ray-wih-tensorflow.md)
+  - [Using Ray on a Cluster](doc/using-ray-on-a-cluster.md)
 
 ## Example Applications
 

doc/using-ray-with-tensorflow.md

@@ -0,0 +1,199 @@
+# 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)
+train = optimizer.minimize(loss)
+
+init = tf.initialize_all_variables()
+sess = tf.Session()
+```
+
+To extract the weights and set the weights, we need to write a couple lines of
+boilerplate code.
+
+```python
+def get_and_set_weights_methods():
+  assignment_placeholders = []
+  assignment_nodes = []
+  for var in tf.trainable_variables():
+    assignment_placeholders.append(tf.placeholder(var.value().dtype, var.get_shape().as_list()))
+    assignment_nodes.append(var.assign(assignment_placeholders[-1]))
+
+  def get_weights():
+    return [v.eval(session=sess) for v in tf.trainable_variables()]
+
+  def set_weights(new_weights):
+    sess.run(assignment_nodes, feed_dict={p: w for p, w in zip(assignment_placeholders, new_weights)})
+
+  return get_weights, set_weights
+
+get_weights, set_weights = get_and_set_weights_methods()
+```
+
+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 = get_weights()  # Returns a list of numpy arrays
+# Set the weights
+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
+reusable variables. Within the reusable variables, we would define `get_weights`
+and `set_weights` methods. We would then use those methods to ship the weights
+(as lists of numpy arrays) between the processes without shipping the actual
+TensorFlow graphs, which are much more complex Python objects.
+
+```python
+import tensorflow as tf
+import numpy as np
+import ray
+
+ray.init(start_ray_local=True, num_workers=5)
+
+BATCH_SIZE = 100
+NUM_BATCHES = 1
+NUM_ITERS = 201
+
+def net_vars_initializer():
+  # Seed TensorFlow to make the script deterministic.
+  tf.set_random_seed(0)
+
+  x_data = tf.placeholder(tf.float32, shape=[BATCH_SIZE])
+  y_data = tf.placeholder(tf.float32, shape=[BATCH_SIZE])
+
+  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)
+  train = optimizer.minimize(loss)
+
+  init = tf.initialize_all_variables()
+  sess = tf.Session()
+
+  # Additional code for setting and getting the weights.
+  def get_and_set_weights_methods():
+    assignment_placeholders = []
+    assignment_nodes = []
+    for var in tf.trainable_variables():
+      assignment_placeholders.append(tf.placeholder(var.value().dtype, var.get_shape().as_list()))
+      assignment_nodes.append(var.assign(assignment_placeholders[-1]))
+
+    def get_weights():
+      return [v.eval(session=sess) for v in tf.trainable_variables()]
+
+    def set_weights(new_weights):
+      sess.run(assignment_nodes, feed_dict={p: w for p, w in zip(assignment_placeholders, new_weights)})
+
+    return get_weights, set_weights
+
+  get_weights, set_weights = get_and_set_weights_methods()
+
+  return get_weights, set_weights, sess, train, loss, x_data, y_data, init
+
+def net_vars_reinitializer(net_vars):
+  return net_vars
+
+# Define a reusable variable for the network variables.
+ray.reusables.net_vars = ray.Reusable(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):
+  get_weights, set_weights, sess, train, _, x_data, y_data, _ = ray.reusables.net_vars
+  # Set the weights in the network.
+  set_weights(weights)
+  # Do one step of training.
+  sess.run(train, feed_dict={x_data: x, y_data: y})
+  # Return the new weights.
+  return get_weights()
+
+get_weights, set_weights, sess, _, loss, x_data, y_data, init = ray.reusables.net_vars
+# Initialize the network weights.
+sess.run(init)
+# Get the weights as a list of numpy arrays.
+weights = 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 list
+  # of weights, and we want to add up these lists component wise.
+  weights = [sum(weight_tuple) / NUM_BATCHES for weight_tuple in zip(*new_weights_list)]
+  # 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/alexnet/alexnet.py

@@ -77,13 +77,12 @@ def load_tarfiles_from_s3(bucket, s3_keys, size=[]):
 
   return [load_tarfile_from_s3.remote(bucket, s3_key, size) for s3_key in s3_keys]
 
-def setup_variables(params, placeholders, assigns, kernelshape, biasshape):
-  """Creates the variables for each layer and adds the variables and the components needed to feed them to various lists
+def setup_variables(params, placeholders, kernelshape, biasshape):
+  """Create the variables for each layer.
 
   Args:
     params (List): Network parameters used for creating feed_dicts
     placeholders (List): Placeholders used for feeding weights into
-    assigns (List): Assignments used for actually setting variables
     kernelshape (List): Shape of the kernel used for the conv layer
     biasshape (List): Shape of the bias used
 
@@ -99,7 +98,6 @@ def setup_variables(params, placeholders, assigns, kernelshape, biasshape):
   update_biases = biases.assign(biases_new)
   params += [kernel, biases]
   placeholders += [kernel_new, biases_new]
-  assigns += [update_kernel, update_biases]
 
 def conv_layer(parameters, prev_layer, shape, scope):
   """Constructs a convolutional layer for the network.
@@ -123,11 +121,10 @@ def net_initialization():
   images = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])
   y_true = tf.placeholder(tf.float32, shape=[None, 1000])
   parameters = []
-  assignment = []
   placeholders = []
   # conv1
   with tf.name_scope('conv1') as scope:
-    setup_variables(parameters, placeholders, assignment, [11, 11, 3, 96], [96])
+    setup_variables(parameters, placeholders, [11, 11, 3, 96], [96])
     conv1 = conv_layer(parameters, images, [1, 4, 4, 1], scope)
 
   # pool1
@@ -144,7 +141,7 @@ def net_initialization():
 
   # conv2
   with tf.name_scope('conv2') as scope:
-    setup_variables(parameters, placeholders, assignment, [5, 5, 96, 256], [256])
+    setup_variables(parameters, placeholders, [5, 5, 96, 256], [256])
     conv2 = conv_layer(parameters, pool1_lrn, [1, 1, 1, 1], scope)
 
   pool2 = tf.nn.max_pool(conv2,
@@ -160,17 +157,17 @@ def net_initialization():
 
   # conv3
   with tf.name_scope('conv3') as scope:
-    setup_variables(parameters, placeholders, assignment, [3, 3, 256, 384], [384])
+    setup_variables(parameters, placeholders, [3, 3, 256, 384], [384])
     conv3 = conv_layer(parameters, pool2_lrn, [1, 1, 1, 1], scope)
 
   # conv4
   with tf.name_scope('conv4') as scope:
-    setup_variables(parameters, placeholders, assignment, [3, 3, 384, 384], [384])
+    setup_variables(parameters, placeholders, [3, 3, 384, 384], [384])
     conv4 = conv_layer(parameters, conv3, [1, 1, 1, 1], scope)
 
   # conv5
   with tf.name_scope('conv5') as scope:
-    setup_variables(parameters, placeholders, assignment, [3, 3, 384, 256], [256])
+    setup_variables(parameters, placeholders, [3, 3, 384, 256], [256])
     conv5 = conv_layer(parameters, conv4, [1, 1, 1, 1], scope)
 
   # pool5
@@ -189,21 +186,21 @@ def net_initialization():
 
   with tf.name_scope('fc1') as scope:
     n_input = int(np.prod(pool5_lrn.get_shape().as_list()[1:]))
-    setup_variables(parameters, placeholders, assignment, [n_input, 4096], [4096])
+    setup_variables(parameters, placeholders, [n_input, 4096], [4096])
     fc_in = tf.reshape(pool5_lrn, [-1, n_input])
     fc_layer1 = tf.nn.tanh(tf.nn.bias_add(tf.matmul(fc_in, parameters[-2]), parameters[-1]))
     fc_out1 = tf.nn.dropout(fc_layer1, dropout)
 
   with tf.name_scope('fc2') as scope:
     n_input = int(np.prod(fc_out1.get_shape().as_list()[1:]))
-    setup_variables(parameters, placeholders, assignment, [n_input, 4096], [4096])
+    setup_variables(parameters, placeholders, [n_input, 4096], [4096])
     fc_in = tf.reshape(fc_out1, [-1, n_input])
     fc_layer2 = tf.nn.tanh(tf.nn.bias_add(tf.matmul(fc_in, parameters[-2]), parameters[-1]))
     fc_out2 = tf.nn.dropout(fc_layer2, dropout)
 
   with tf.name_scope('fc3') as scope:
     n_input = int(np.prod(fc_out2.get_shape().as_list()[1:]))
-    setup_variables(parameters, placeholders, assignment, [n_input, 1000], [1000])
+    setup_variables(parameters, placeholders, [n_input, 1000], [1000])
     fc_in = tf.reshape(fc_out2, [-1, n_input])
     fc_layer3 = tf.nn.softmax(tf.nn.bias_add(tf.matmul(fc_in, parameters[-2]), parameters[-1]))
 
@@ -221,10 +218,33 @@ def net_initialization():
 
   comp_grads = opt.compute_gradients(cross_entropy, parameters)
 
-  application = opt.apply_gradients(zip(placeholders,parameters))
+  application = opt.apply_gradients(zip(placeholders, parameters))
   sess = tf.Session()
   init_all_variables = tf.initialize_all_variables()
-  return comp_grads, sess, application, accuracy, images, y_true, dropout, placeholders, parameters, assignment, init_all_variables
+
+  # In order to set the weights of the TensorFlow graph on a worker, we add
+  # assignment nodes. To get the network weights (as a list of numpy arrays)
+  # and to set the network weights (from a list of numpy arrays), use the
+  # methods get_weights and set_weights. This can be done from within a remote
+  # function or on the driver.
+  def get_and_set_weights_methods():
+    assignment_placeholders = []
+    assignment_nodes = []
+    for var in tf.trainable_variables():
+      assignment_placeholders.append(tf.placeholder(var.value().dtype, var.get_shape().as_list()))
+      assignment_nodes.append(var.assign(assignment_placeholders[-1]))
+
+    def get_weights():
+      return [v.eval(session=sess) for v in tf.trainable_variables()]
+
+    def set_weights(new_weights):
+      sess.run(assignment_nodes, feed_dict={p: w for p, w in zip(assignment_placeholders, new_weights)})
+
+    return get_weights, set_weights
+
+  get_weights, set_weights = get_and_set_weights_methods()
+
+  return comp_grads, sess, application, accuracy, images, y_true, dropout, placeholders, init_all_variables, get_weights, set_weights
 
 
 def net_reinitialization(net_vars):
@@ -392,10 +412,9 @@ def compute_grad(X, Y, mean, weights):
   Returns:
     List of gradients for each variable
   """
-  comp_grads, sess, _, _, images, y_true, dropout, placeholders, _, assignment, _ = ray.reusables.net_vars
+  comp_grads, sess, _, _, images, y_true, dropout, placeholders, _, get_weights, set_weights = ray.reusables.net_vars
   # Set the network weights.
-  feed_dict = dict(zip(placeholders, weights))
-  sess.run(assignment, feed_dict=feed_dict)
+  set_weights(weights)
   # Choose a subset of the batch to compute on and crop the images.
   random_indices = np.random.randint(0, len(X), size=128)
   subset_X = crop_images(X[random_indices] - mean)
@@ -416,10 +435,9 @@ def compute_accuracy(X, Y, weights):
   Returns:
     The accuracy of the network on the given batch.
   """
-  _, sess, _, accuracy, images, y_true, dropout, placeholders, _, assignment, _ = ray.reusables.net_vars
+  _, sess, _, accuracy, images, y_true, dropout, placeholders, _, get_weights, set_weights = ray.reusables.net_vars
   # Set the network weights.
-  feed_dict = dict(zip(placeholders, weights))
-  sess.run(assignment, feed_dict=feed_dict)
+  set_weights(weights)
 
   one_hot_Y = np.asarray([one_hot(label) for label in Y])
   cropped_X = crop_images(X)

examples/alexnet/driver.py

@@ -64,7 +64,7 @@ if __name__ == "__main__":
   for i in range(num_shuffles):
     batches = alexnet.shuffle(batches)
 
-  _, sess, application, _, _, _, _, placeholders, parameters, assignment, init_all_variables = ray.reusables.net_vars
+  _, sess, application, _, _, _, _, placeholders, init_all_variables, get_weights, set_weights = ray.reusables.net_vars
   # Initialize the network and optimizer weights. This is only run once on the
   # driver. We initialize the weights manually on the workers.
   sess.run(init_all_variables)
@@ -73,7 +73,7 @@ if __name__ == "__main__":
   iteration = 0
   while True:
     # Extract weights from the local copy of the network.
-    weights = sess.run(parameters)
+    weights = get_weights()
     # Put weights in the object store.
     weights_id = ray.put(weights)
 

examples/lbfgs/driver.py

@@ -45,14 +45,31 @@ if __name__ == "__main__":
     cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))
     cross_entropy_grads = tf.gradients(cross_entropy, [w, b])
 
-    w_new = tf.placeholder(tf.float32, w_shape)
-    b_new = tf.placeholder(tf.float32, b_shape)
-    update_w = w.assign(w_new)
-    update_b = b.assign(b_new)
-
     sess = tf.Session()
 
-    return sess, update_w, update_b, cross_entropy, cross_entropy_grads, x, y_, w_new, b_new
+    # In order to set the weights of the TensorFlow graph on a worker, we add
+    # assignment nodes. To get the network weights (as a list of numpy arrays)
+    # and to set the network weights (from a list of numpy arrays), use the
+    # methods get_weights and set_weights. This can be done from within a remote
+    # function or on the driver.
+    def get_and_set_weights_methods():
+      assignment_placeholders = []
+      assignment_nodes = []
+      for var in tf.trainable_variables():
+        assignment_placeholders.append(tf.placeholder(var.value().dtype, var.get_shape().as_list()))
+        assignment_nodes.append(var.assign(assignment_placeholders[-1]))
+
+      def get_weights():
+        return [v.eval(session=sess) for v in tf.trainable_variables()]
+
+      def set_weights(new_weights):
+        sess.run(assignment_nodes, feed_dict={p: w for p, w in zip(assignment_placeholders, new_weights)})
+
+      return get_weights, set_weights
+
+    get_weights, set_weights = get_and_set_weights_methods()
+
+    return sess, cross_entropy, cross_entropy_grads, x, y_, get_weights, set_weights
 
   # By default, when a reusable variable is used by a remote function, the
   # initialization code will be rerun at the end of the remote task to ensure
@@ -70,20 +87,20 @@ if __name__ == "__main__":
 
   # Load the weights into the network.
   def load_weights(theta):
-    sess, update_w, update_b, _, _, _, _, w_new, b_new = ray.reusables.net_vars
-    sess.run([update_w, update_b], feed_dict={w_new: theta[:w_size].reshape(w_shape), b_new: theta[w_size:]})
+    sess, _, _, _, _, get_weights, set_weights = ray.reusables.net_vars
+    set_weights([theta[:w_size].reshape(w_shape), theta[w_size:].reshape(b_shape)])
 
   # Compute the loss on a batch of data.
   @ray.remote
   def loss(theta, xs, ys):
-    sess, _, _, cross_entropy, _, x, y_, _, _ = ray.reusables.net_vars
+    sess, cross_entropy, _, x, y_, _, _ = ray.reusables.net_vars
     load_weights(theta)
     return float(sess.run(cross_entropy, feed_dict={x: xs, y_: ys}))
 
   # Compute the gradient of the loss on a batch of data.
   @ray.remote
   def grad(theta, xs, ys):
-    sess, _, _, _, cross_entropy_grads, x, y_, _, _ = ray.reusables.net_vars
+    sess, _, cross_entropy_grads, x, y_, _, _ = ray.reusables.net_vars
     load_weights(theta)
     gradients = sess.run(cross_entropy_grads, feed_dict={x: xs, y_: ys})
     return np.concatenate([g.flatten() for g in gradients])