update lbfgs app (#301)

This commit is contained in:
Robert Nishihara
2016-07-26 23:45:14 -07:00
committed by Philipp Moritz
parent b5215f1e6a
commit a97574d471
5 changed files with 100 additions and 116 deletions
+1 -1
View File
@@ -9,7 +9,7 @@ Then from the directory `ray/examples/hyperopt/` run the following.
```
source ../../setup-env.sh
python driver.py # This will take a minute to first download the MNIST dataset.
python driver.py
```
Machine learning algorithms often have a number of *hyperparameters* whose
+10 -24
View File
@@ -31,7 +31,6 @@ built in methods for loading the data.
```python
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
batch_size = 100
num_batches = mnist.train.num_examples / batch_size
batches = [mnist.train.next_batch(batch_size) for _ in range(num_batches)]
@@ -70,39 +69,26 @@ functions, along with an initial choice of model parameters, into
`scipy.optimize.fmin_l_bfgs_b`.
```python
theta_init = np.zeros(functions.dim)
theta_init = 1e-2 * np.random.normal(size=dim)
result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, fprime=full_grad)
```
### The distributed version
The extent to which this computation can be parallelized depends on the
specifics of the optimization problem, but for large datasets, the computation
of the gradient itself is often embarrassingly parallel.
To run this example in Ray, we use three files.
- [driver.py](driver.py) - This is the script that gets run. It launches the
remote tasks and retrieves the results. The application can be run with
`python driver.py`.
- [functions.py](functions.py) - This is the file that defines the remote
functions (in this case, just `loss` and `grad`).
- [worker.py](worker.py) - This is the Python code that each worker process
runs. It imports the relevant modules and tells the scheduler what functions
it knows how to execute. Then it enters a loop that waits to receive tasks
from the scheduler.
In this example, the computation of the gradient itself can be done in parallel
on a number of workers or machines.
First, let's turn the data into a collection of remote objects.
```python
batch_refs = [ray.put(xs), ray.put(ys) for (xs, ys) in batches]
batch_refs = [(ray.put(xs), ray.put(ys)) for (xs, ys) in batches]
```
MNIST easily fits on a single machine, but for larger data sets, we will need to
We can load the data on the driver and distribute it this way because MNIST
easily fits on a single machine. However, for larger data sets, we will need to
use remote functions to distribute the loading of the data.
Now, lets turn `loss` and `grad` into remote functions. In the example
application, this is done in [functions.py](functions.py).
Now, lets turn `loss` and `grad` into remote functions.
```python
@ray.remote([np.ndarray, np.ndarray, np.ndarray], [float])
@@ -139,14 +125,14 @@ Note that we turn `theta` into a remote object with the line `theta_ref =
ray.put(theta)` before passing it into the remote functions. If we had written
```python
[loss(theta, xs_ref, ys_ref) for ... in batch_refs]
[loss(theta, xs_ref, ys_ref) for (xs_ref, ys_ref) in batch_refs]
```
instead of
```python
theta_ref = ray.put(theta)
[loss(theta_ref, xs_ref, ys_ref) for ... in batch_refs]
[loss(theta_ref, xs_ref, ys_ref) for (xs_ref, ys_ref) in batch_refs]
```
then each task that got sent to the scheduler (one for every element of
@@ -162,6 +148,6 @@ identical to the behavior in the serial version.
We can now optimize the objective with the same function call as before.
```python
theta_init = np.zeros(functions.dim)
theta_init = 1e-2 * np.random.normal(size=dim)
result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, fprime=full_grad)
```
+89 -20
View File
@@ -1,24 +1,92 @@
import numpy as np
import scipy.optimize
import os
import ray
import functions
import numpy as np
import scipy.optimize
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
if __name__ == "__main__":
worker_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py")
ray.services.start_ray_local(num_workers=16, worker_path=worker_path)
ray.services.start_ray_local(num_workers=16)
print "Downloading and loading MNIST data..."
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
# Define the dimensions of the data and of the model.
image_dimension = 784
label_dimension = 10
w_shape = [image_dimension, label_dimension]
w_size = np.prod(w_shape)
b_shape = [label_dimension]
b_size = np.prod(b_shape)
dim = w_size + b_size
batch_size = 100
num_batches = mnist.train.num_examples / batch_size
batches = [mnist.train.next_batch(batch_size) for _ in range(num_batches)]
# Define a function for initializing the network. Note that this code does not
# call initialize the network weights. If it did, the weights would be
# randomly initialized on each worker and would differ from worker to worker.
# We pass the weights into the remote functions loss and grad so that the
# weights are the same on each worker.
def net_initialization():
x = tf.placeholder(tf.float32, [None, image_dimension])
w = tf.Variable(tf.zeros(w_shape))
b = tf.Variable(tf.zeros(b_shape))
y = tf.nn.softmax(tf.matmul(x, w) + b)
y_ = tf.placeholder(tf.float32, [None, label_dimension])
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1]))
cross_entropy_grads = tf.gradients(cross_entropy, [w, b])
batch_refs = [(ray.put(xs), ray.put(ys)) for (xs, ys) in batches]
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
# 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
# that the state of the variable is not changed by the remote task. However,
# the initialization code may be expensive. This case is one example, because
# a TensorFlow network is constructed. In this case, we pass in a special
# reinitialization function which gets run instead of the original
# initialization code. As users, if we pass in custom reinitialization code,
# we must ensure that no state is leaked between tasks.
def net_reinitialization(net_vars):
return net_vars
# Create a reusable variable for the network.
ray.reusables.net_vars = ray.Reusable(net_initialization, net_reinitialization)
# 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:]})
# Compute the loss on a batch of data.
@ray.remote([np.ndarray, np.ndarray, np.ndarray], [float])
def loss(theta, xs, ys):
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([np.ndarray, np.ndarray, np.ndarray], [np.ndarray])
def grad(theta, xs, ys):
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])
# Compute the loss on the entire dataset.
def full_loss(theta):
theta_ref = ray.put(theta)
loss_refs = [loss(theta_ref, xs_ref, ys_ref) for (xs_ref, ys_ref) in batch_refs]
return sum([ray.get(loss_ref) for loss_ref in loss_refs])
# Compute the gradient of the loss on the entire dataset.
def full_grad(theta):
theta_ref = ray.put(theta)
grad_refs = [grad(theta_ref, xs_ref, ys_ref) for (xs_ref, ys_ref) in batch_refs]
return sum([ray.get(grad_ref) for grad_ref in grad_refs]).astype("float64") # This conversion is necessary for use with fmin_l_bfgs_b.
# From the perspective of scipy.optimize.fmin_l_bfgs_b, full_loss is simply a
# function which takes some parameters theta, and computes a loss. Similarly,
@@ -29,15 +97,16 @@ if __name__ == "__main__":
# potentially distributed over a cluster. However, these details are hidden
# from scipy.optimize.fmin_l_bfgs_b, which simply uses it to run the L-BFGS
# algorithm.
def full_loss(theta):
theta_ref = ray.put(theta)
loss_refs = [functions.loss(theta_ref, xs_ref, ys_ref) for (xs_ref, ys_ref) in batch_refs]
return sum([ray.get(loss_ref) for loss_ref in loss_refs])
def full_grad(theta):
theta_ref = ray.put(theta)
grad_refs = [functions.grad(theta_ref, xs_ref, ys_ref) for (xs_ref, ys_ref) in batch_refs]
return sum([ray.get(grad_ref) for grad_ref in grad_refs]).astype("float64") # This conversion is necessary for use with fmin_l_bfgs_b.
# Load the mnist data and turn the data into remote objects.
print "Downloading the MNIST dataset. This may take a minute."
mnist = input_data.read_data_sets("MNIST_data", one_hot=True)
batch_size = 100
num_batches = mnist.train.num_examples / batch_size
batches = [mnist.train.next_batch(batch_size) for _ in range(num_batches)]
batch_refs = [(ray.put(xs), ray.put(ys)) for (xs, ys) in batches]
theta_init = np.zeros(functions.dim)
# Initialize the weights for the network to the vector of all zeros.
theta_init = 1e-2 * np.random.normal(size=dim)
# Use L-BFGS to minimize the loss function.
result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, maxiter=10, fprime=full_grad, disp=True)
-43
View File
@@ -1,43 +0,0 @@
import ray
import numpy as np
import tensorflow as tf
image_dimension = 784
label_dimension = 10
w_shape = [image_dimension, label_dimension]
w_size = np.prod(w_shape)
b_shape = [label_dimension]
b_size = np.prod(b_shape)
dim = w_size + b_size
x = tf.placeholder(tf.float32, [None, image_dimension])
w = tf.Variable(tf.zeros(w_shape))
b = tf.Variable(tf.zeros(b_shape))
y = tf.nn.softmax(tf.matmul(x, w) + b)
y_ = tf.placeholder(tf.float32, [None, label_dimension])
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)
init = tf.initialize_all_variables()
sess = tf.Session()
sess.run(init)
def load_weights(theta):
sess.run([update_w, update_b], feed_dict={w_new: theta[:w_size].reshape(w_shape), b_new: theta[w_size:]})
@ray.remote([np.ndarray, np.ndarray, np.ndarray], [float])
def loss(theta, xs, ys):
load_weights(theta)
return float(sess.run(cross_entropy, feed_dict={x: xs, y_: ys}))
@ray.remote([np.ndarray, np.ndarray, np.ndarray], [np.ndarray])
def grad(theta, xs, ys):
load_weights(theta)
gradients = sess.run(cross_entropy_grads, feed_dict={x: xs, y_: ys})
return np.concatenate([g.flatten() for g in gradients])
-28
View File
@@ -1,28 +0,0 @@
import argparse
import ray
import ray.array.remote as ra
import ray.array.distributed as da
import functions
parser = argparse.ArgumentParser(description="Parse addresses for the worker to connect to.")
parser.add_argument("--scheduler-address", default="127.0.0.1:10001", type=str, help="the scheduler's address")
parser.add_argument("--objstore-address", default="127.0.0.1:20001", type=str, help="the objstore's address")
parser.add_argument("--worker-address", default="127.0.0.1:40001", type=str, help="the worker's address")
if __name__ == "__main__":
args = parser.parse_args()
ray.worker.connect(args.scheduler_address, args.objstore_address, args.worker_address)
ray.register_module(functions)
ray.register_module(ra)
ray.register_module(ra.random)
ray.register_module(ra.linalg)
ray.register_module(da)
ray.register_module(da.random)
ray.register_module(da.linalg)
ray.worker.main_loop()