diff --git a/README.md b/README.md index 6c26c4a29..eb479c9f8 100644 --- a/README.md +++ b/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 diff --git a/doc/using-ray-with-tensorflow.md b/doc/using-ray-with-tensorflow.md new file mode 100644 index 000000000..fbf713d52 --- /dev/null +++ b/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) +``` diff --git a/examples/alexnet/alexnet.py b/examples/alexnet/alexnet.py index 5b0a707a8..9f656763f 100644 --- a/examples/alexnet/alexnet.py +++ b/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) diff --git a/examples/alexnet/driver.py b/examples/alexnet/driver.py index a9cdd6e86..17afe13c3 100644 --- a/examples/alexnet/driver.py +++ b/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) diff --git a/examples/lbfgs/driver.py b/examples/lbfgs/driver.py index 9662e8cca..dcda76f7f 100644 --- a/examples/lbfgs/driver.py +++ b/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])