Added helper class for getting tf variables from loss function (ray-project#184)

* Added helper class for getting tf variables from loss function

* Updated usage and documentation

* Removed try-catches

* Added futures

* Added documentation

* fixes and tests

* more tests

* install tensorflow in travis

Author: Wapaul1

.travis.yml

@@ -72,6 +72,7 @@ script:
 
   - python test/runtest.py
   - python test/array_test.py
+  - python test/tensorflow_test.py
   - python test/failure_test.py
   - python test/microbenchmarks.py
   - python test/stress_tests.py

.travis/install-dependencies.sh

@@ -20,7 +20,7 @@ fi
 if [[ "$PYTHON" == "2.7" ]] && [[ "$platform" == "linux" ]]; then
   sudo apt-get update
   sudo apt-get install -y cmake build-essential autoconf curl libtool python-dev python-numpy python-pip libboost-all-dev unzip
-  sudo pip install funcsigs colorama psutil redis
+  sudo pip install funcsigs colorama psutil redis tensorflow
   sudo pip install --upgrade git+git://github.com/cloudpipe/cloudpickle.git@0d225a4695f1f65ae1cbb2e0bbc145e10167cce4
 elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "linux" ]]; then
   sudo apt-get update
@@ -29,7 +29,7 @@ elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "linux" ]]; then
   wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
   bash miniconda.sh -b -p $HOME/miniconda
   export PATH="$HOME/miniconda/bin:$PATH"
-  pip install numpy funcsigs colorama psutil redis
+  pip install numpy funcsigs colorama psutil redis tensorflow
   pip install --upgrade git+git://github.com/cloudpipe/cloudpickle.git@0d225a4695f1f65ae1cbb2e0bbc145e10167cce4
 elif [[ "$PYTHON" == "2.7" ]] && [[ "$platform" == "macosx" ]]; then
   # check that brew is installed
@@ -43,7 +43,7 @@ elif [[ "$PYTHON" == "2.7" ]] && [[ "$platform" == "macosx" ]]; then
   fi
   brew install cmake automake autoconf libtool boost
   sudo easy_install pip
-  sudo pip install numpy funcsigs colorama psutil redis --ignore-installed six
+  sudo pip install numpy funcsigs colorama psutil redis tensorflow --ignore-installed six
   sudo pip install --upgrade git+git://github.com/cloudpipe/cloudpickle.git@0d225a4695f1f65ae1cbb2e0bbc145e10167cce4
 elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "macosx" ]]; then
   # check that brew is installed
@@ -60,7 +60,7 @@ elif [[ "$PYTHON" == "3.5" ]] && [[ "$platform" == "macosx" ]]; then
   wget https://repo.continuum.io/miniconda/Miniconda3-latest-MacOSX-x86_64.sh -O miniconda.sh
   bash miniconda.sh -b -p $HOME/miniconda
   export PATH="$HOME/miniconda/bin:$PATH"
-  pip install numpy funcsigs colorama psutil redis
+  pip install numpy funcsigs colorama psutil redis tensorflow
   pip install --upgrade git+git://github.com/cloudpipe/cloudpickle.git@0d225a4695f1f65ae1cbb2e0bbc145e10167cce4
 else
   echo "Unrecognized environment."

doc/using-ray-with-tensorflow.md

@@ -37,38 +37,24 @@ init = tf.initialize_all_variables()
 sess = tf.Session()
 ```
 
-To extract the weights and set the weights, we need to write a couple lines of
-boilerplate code.
+To extract the weights and set the weights, you can call
 
 ```python
-def get_and_set_weights_methods():
-  assignment_placeholders = []
-  assignment_nodes = []
-  for var in tf.trainable_variables():
-    assignment_placeholders.append(tf.placeholder(var.value().dtype, var.get_shape().as_list()))
-    assignment_nodes.append(var.assign(assignment_placeholders[-1]))
-  # Define a function for getting the network weights.
-  def get_weights():
-    return [v.eval(session=sess) for v in tf.trainable_variables()]
-  # Define a function for setting the network weights.
-  def set_weights(new_weights):
-    sess.run(assignment_nodes, feed_dict={p: w for p, w in zip(assignment_placeholders, new_weights)})
-  # Return the methods.
-  return get_weights, set_weights
-
-get_weights, set_weights = get_and_set_weights_methods()
+variables = ray.experimental.TensorFlowVariables(loss, sess)
 ```
 
+which gives you methods to set and get the weights as well as collecting all of the variables in the model.
+
 Now we can use these methods to extract the weights, and place them back in the
 network as follows.
 
 ```python
 # First initialize the weights.
 sess.run(init)
 # Get the weights
-weights = get_weights()  # Returns a list of numpy arrays
+weights = variables.get_weights()  # Returns a dictionary of numpy arrays
 # Set the weights
-set_weights(weights)
+variables.set_weights(weights)
 ```
 
 **Note:** If we were to set the weights using the `assign` method like below,
@@ -117,20 +103,9 @@ def net_vars_initializer():
   init = tf.initialize_all_variables()
   sess = tf.Session()
   # Additional code for setting and getting the weights.
-  def get_and_set_weights_methods():
-    assignment_placeholders = []
-    assignment_nodes = []
-    for var in tf.trainable_variables():
-      assignment_placeholders.append(tf.placeholder(var.value().dtype, var.get_shape().as_list()))
-      assignment_nodes.append(var.assign(assignment_placeholders[-1]))
-    def get_weights():
-      return [v.eval(session=sess) for v in tf.trainable_variables()]
-    def set_weights(new_weights):
-      sess.run(assignment_nodes, feed_dict={p: w for p, w in zip(assignment_placeholders, new_weights)})
-    return get_weights, set_weights
-  get_weights, set_weights = get_and_set_weights_methods()
+  variables = ray.experimental.TensorFlowVariables(loss, sess)
   # Return all of the data needed to use the network.
-  return get_weights, set_weights, sess, train, loss, x_data, y_data, init
+  return variables, sess, train, loss, x_data, y_data, init
 
 def net_vars_reinitializer(net_vars):
   return net_vars
@@ -142,19 +117,19 @@ ray.reusables.net_vars = ray.Reusable(net_vars_initializer, net_vars_reinitializ
 # new weights.
 @ray.remote
 def step(weights, x, y):
-  get_weights, set_weights, sess, train, _, x_data, y_data, _ = ray.reusables.net_vars
+  variables, sess, train, _, x_data, y_data, _ = ray.reusables.net_vars
   # Set the weights in the network.
-  set_weights(weights)
+  variables.set_weights(weights)
   # Do one step of training.
   sess.run(train, feed_dict={x_data: x, y_data: y})
   # Return the new weights.
-  return get_weights()
+  return variables.get_weights()
 
-get_weights, set_weights, sess, _, loss, x_data, y_data, init = ray.reusables.net_vars
+variables, sess, _, loss, x_data, y_data, init = ray.reusables.net_vars
 # Initialize the network weights.
 sess.run(init)
 # Get the weights as a list of numpy arrays.
-weights = get_weights()
+weights = variables.get_weights()
 
 # Define a remote function for generating fake data.
 @ray.remote(num_return_vals=2)

lib/python/ray/experimental/__init__.py

@@ -3,3 +3,4 @@
 from __future__ import print_function
 
 from .utils import copy_directory
+from .tfutils import TensorFlowVariables

lib/python/ray/experimental/tfutils.py

@@ -0,0 +1,37 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+class TensorFlowVariables(object):
+  """An object used to extract variables from a loss function, and provide 
+     methods for getting and setting the weights of said variables.
+
+  Attributes:
+    sess (tf.Session): The tensorflow session used to run assignment.
+    loss: The loss function passed in by the user.
+    variables (List[tf.Variable]): Extracted variables from the loss.
+    assignment_placeholders (List[tf.placeholders]): The nodes that weights get passed to.
+    assignment_nodes (List[tf.Tensor]): The nodes that assign the weights.
+  """
+  def __init__(self, loss, sess):
+    """Creates a TensorFlowVariables instance."""
+    import tensorflow as tf
+    self.sess = sess
+    self.loss = loss
+    variable_names = [op.node_def.name for op in loss.graph.get_operations() if op.node_def.op == "Variable"]
+    self.variables = [v for v in tf.trainable_variables() if v.op.node_def.name in variable_names]
+    self.assignment_placeholders = dict()
+    self.assignment_nodes = []
+    
+    # Create new placeholders to put in custom weights.
+    for var in self.variables:
+      self.assignment_placeholders[var.op.node_def.name] = tf.placeholder(var.value().dtype, var.get_shape().as_list())
+      self.assignment_nodes.append(var.assign(self.assignment_placeholders[var.op.node_def.name]))
+
+  def get_weights(self):
+    """Returns the weights of the variables of the loss function in a list."""
+    return {v.op.node_def.name: v.eval(session=self.sess) for v in self.variables}
+
+  def set_weights(self, new_weights):
+    """Sets the weights to new_weights."""
+    self.sess.run(self.assignment_nodes, feed_dict={self.assignment_placeholders[name]: value for (name, value) in new_weights.items()})

test/tensorflow_test.py

@@ -0,0 +1,53 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import unittest
+import tensorflow as tf
+import ray
+
+class TensorFlowTest(unittest.TestCase):
+
+  def testTensorFlowVariables(self):
+    ray.init(start_ray_local=True, num_workers=2)
+
+    x_data = tf.placeholder(tf.float32, shape=[100])
+    y_data = tf.placeholder(tf.float32, shape=[100])
+
+    w = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
+    b = tf.Variable(tf.zeros([1]))
+    y = w * x_data + b
+    loss = tf.reduce_mean(tf.square(y - y_data))
+
+    sess = tf.Session()
+    sess.run(tf.global_variables_initializer())
+
+    variables = ray.experimental.TensorFlowVariables(loss, sess)
+    weights = variables.get_weights()
+
+    for (name, val) in weights.items():
+      weights[name] += 1.0
+
+    variables.set_weights(weights)
+    self.assertEqual(weights, variables.get_weights())
+
+    w2 = tf.Variable(tf.random_uniform([1], -1.0, 1.0), name="w")
+    b2 = tf.Variable(tf.zeros([1]), name="b")
+    y2 = w2 * x_data + b2
+    loss2 = tf.reduce_mean(tf.square(y2 - y_data))
+
+    sess.run(tf.global_variables_initializer())
+
+    variables2 = ray.experimental.TensorFlowVariables(loss2, sess)
+    weights2 = variables2.get_weights()
+
+    for (name, val) in weights2.items():
+      weights2[name] += 2.0
+
+    variables2.set_weights(weights2)
+    self.assertEqual(weights2, variables2.get_weights())
+
+    ray.worker.cleanup()
+
+if __name__ == "__main__":
+  unittest.main(verbosity=2)