diff --git a/doc/using-ray-with-tensorflow.md b/doc/using-ray-with-tensorflow.md index d917509a4..7e7dc2fa3 100644 --- a/doc/using-ray-with-tensorflow.md +++ b/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,7 +127,7 @@ 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. @@ -133,7 +135,7 @@ def step(weights, x, 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)) +``` diff --git a/python/ray/experimental/tfutils.py b/python/ray/experimental/tfutils.py index 258d60bbc..365b897dc 100644 --- a/python/ray/experimental/tfutils.py +++ b/python/ray/experimental/tfutils.py @@ -63,13 +63,13 @@ class TensorFlowVariables(object): 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 @@ class TensorFlowVariables(object): 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 @@ class TensorFlowVariables(object): 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()}) diff --git a/test/tensorflow_test.py b/test/tensorflow_test.py index a4dae42e2..0d9b9fd27 100644 --- a/test/tensorflow_test.py +++ b/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 @@ class TensorFlowTest(unittest.TestCase): @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)