mirror of
https://github.com/wassname/ray.git
synced 2026-08-11 11:24:51 +08:00
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:
@@ -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))
|
||||
```
|
||||
|
||||
@@ -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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user