Merge pull request #109 from amplab/lbfgs

Code for doing batch distributed L-BFGS

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()

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])

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()