From a97574d4718d13ca56a5dbcee19501f88f2dc2ce Mon Sep 17 00:00:00 2001 From: Robert Nishihara Date: Tue, 26 Jul 2016 23:45:14 -0700 Subject: [PATCH] update lbfgs app (#301) --- examples/hyperopt/README.md | 2 +- examples/lbfgs/README.md | 34 ++++------- examples/lbfgs/driver.py | 109 +++++++++++++++++++++++++++++------- examples/lbfgs/functions.py | 43 -------------- examples/lbfgs/worker.py | 28 --------- 5 files changed, 100 insertions(+), 116 deletions(-) delete mode 100644 examples/lbfgs/functions.py delete mode 100644 examples/lbfgs/worker.py diff --git a/examples/hyperopt/README.md b/examples/hyperopt/README.md index f2d43a1bc..578b9a5ae 100644 --- a/examples/hyperopt/README.md +++ b/examples/hyperopt/README.md @@ -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 diff --git a/examples/lbfgs/README.md b/examples/lbfgs/README.md index 677b35eea..2b48bdea1 100644 --- a/examples/lbfgs/README.md +++ b/examples/lbfgs/README.md @@ -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) ``` diff --git a/examples/lbfgs/driver.py b/examples/lbfgs/driver.py index 7f4b307cb..0b1f883cc 100644 --- a/examples/lbfgs/driver.py +++ b/examples/lbfgs/driver.py @@ -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) diff --git a/examples/lbfgs/functions.py b/examples/lbfgs/functions.py deleted file mode 100644 index eece0f9db..000000000 --- a/examples/lbfgs/functions.py +++ /dev/null @@ -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]) diff --git a/examples/lbfgs/worker.py b/examples/lbfgs/worker.py deleted file mode 100644 index 1fc8c9364..000000000 --- a/examples/lbfgs/worker.py +++ /dev/null @@ -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()