diff --git a/examples/lbfgs/driver.py b/examples/lbfgs/driver.py new file mode 100644 index 000000000..701dd94c0 --- /dev/null +++ b/examples/lbfgs/driver.py @@ -0,0 +1,55 @@ +import numpy as np +import scipy.optimize +import os +import time +import ray +import ray.services as services +import ray.worker as worker + +import ray.arrays.remote as ra +import ray.arrays.distributed as da + +import functions + +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)] + +if __name__ == "__main__": + test_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py") + test_path = os.path.join("worker.py") + services.start_singlenode_cluster(return_drivers=False, num_workers_per_objstore=16, worker_path=test_path) + + x_batches = [ray.push(batches[i][0]) for i in range(num_batches)] + y_batches = [ray.push(batches[i][1]) for i in range(num_batches)] + + # 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, + # full_grad is a function which takes some parameters theta, and computes the + # gradient of the loss. Internally, these functions use Ray to distribute the + # computation of the loss and the gradient over the data that is represented + # by the remote object references is x_batches and y_batches and which is + # 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.push(theta) + val_ref = ra.sum_list(*[functions.loss(theta_ref, x_batches[i], y_batches[i]) for i in range(num_batches)]) + return ray.pull(val_ref) + + def full_grad(theta): + theta_ref = ray.push(theta) + grad_ref = ra.sum_list(*[functions.grad(theta_ref, x_batches[i], y_batches[i]) for i in range(num_batches)]) + return ray.pull(grad_ref).astype("float64") # This conversion is necessary for use with fmin_l_bfgs_b. + + theta_init = np.zeros(functions.dim) + + start_time = time.time() + result = scipy.optimize.fmin_l_bfgs_b(full_loss, theta_init, maxiter=10, fprime=full_grad, disp=True) + end_time = time.time() + print "Elapsed time = {}".format(end_time - start_time) + + services.cleanup() diff --git a/examples/lbfgs/functions.py b/examples/lbfgs/functions.py new file mode 100644 index 000000000..e20078b7a --- /dev/null +++ b/examples/lbfgs/functions.py @@ -0,0 +1,43 @@ +import numpy as np +import ray + +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 new file mode 100644 index 000000000..1e1d44ff7 --- /dev/null +++ b/examples/lbfgs/worker.py @@ -0,0 +1,29 @@ +import argparse + +import ray +import ray.worker as worker + +import ray.arrays.remote as ra +import ray.arrays.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() + 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) + + worker.main_loop()