Added example for compute grads in ray tutorial (#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
This commit is contained in:
Wapaul1
2017-02-07 18:07:21 -08:00
committed by Philipp Moritz
parent 1fec94ef00
commit 1a7e1c47cb
3 changed files with 185 additions and 15 deletions
+147 -5
View File
@@ -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))
```
+5 -5
View File
@@ -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()})
+33 -5
View File
@@ -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)