Move TensorFlowVariables to ray.experimental.tf_utils. (#4145)

examples/hyperopt/objective.py

@@ -6,9 +6,11 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import ray
 import tensorflow as tf
 
+import ray
+import ray.experimental.tf_utils
+
 
 def get_batch(data, batch_index, batch_size):
     # This method currently drops data when num_data is not divisible by
@@ -34,8 +36,8 @@ def conv2d(x, W):
 
 
 def max_pool_2x2(x):
-    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1],
-                          padding="SAME")
+    return tf.nn.max_pool(
+        x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")
 
 
 def cnn_setup(x, y, keep_prob, lr, stddev):
@@ -59,8 +61,8 @@ def cnn_setup(x, y, keep_prob, lr, stddev):
     W_fc2 = weight([fc_hidden, 10], stddev)
     b_fc2 = bias([10])
     y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)
-    cross_entropy = tf.reduce_mean(-tf.reduce_sum(y * tf.log(y_conv),
-                                   reduction_indices=[1]))
+    cross_entropy = tf.reduce_mean(
+        -tf.reduce_sum(y * tf.log(y_conv), reduction_indices=[1]))
     correct_pred = tf.equal(tf.argmax(y_conv, 1), tf.argmax(y, 1))
     return (tf.train.AdamOptimizer(lr).minimize(cross_entropy),
             tf.reduce_mean(tf.cast(correct_pred, tf.float32)), cross_entropy)
@@ -69,8 +71,12 @@ def cnn_setup(x, y, keep_prob, lr, stddev):
 # Define a remote function that takes a set of hyperparameters as well as the
 # data, consructs and trains a network, and returns the validation accuracy.
 @ray.remote
-def train_cnn_and_compute_accuracy(params, steps, train_images, train_labels,
-                                   validation_images, validation_labels,
+def train_cnn_and_compute_accuracy(params,
+                                   steps,
+                                   train_images,
+                                   train_labels,
+                                   validation_images,
+                                   validation_labels,
                                    weights=None):
     # Extract the hyperparameters from the params dictionary.
     learning_rate = params["learning_rate"]
@@ -90,7 +96,8 @@ def train_cnn_and_compute_accuracy(params, steps, train_images, train_labels,
         with tf.Session() as sess:
             # Use the TensorFlowVariables utility. This is only necessary if we
             # want to set and get the weights.
-            variables = ray.experimental.TensorFlowVariables(loss, sess)
+            variables = ray.experimental.tf_utils.TensorFlowVariables(
+                loss, sess)
             # Initialize the network weights.
             sess.run(tf.global_variables_initializer())
             # If some network weights were passed in, set those.
@@ -102,12 +109,19 @@ def train_cnn_and_compute_accuracy(params, steps, train_images, train_labels,
                 image_batch = get_batch(train_images, i, batch_size)
                 label_batch = get_batch(train_labels, i, batch_size)
                 # Do one step of training.
-                sess.run(train_step, feed_dict={x: image_batch, y: label_batch,
-                                                keep_prob: keep})
+                sess.run(
+                    train_step,
+                    feed_dict={
+                        x: image_batch,
+                        y: label_batch,
+                        keep_prob: keep
+                    })
             # Training is done, so compute the validation accuracy and the
             # current weights and return.
-            totalacc = accuracy.eval(feed_dict={x: validation_images,
-                                                y: validation_labels,
-                                                keep_prob: 1.0})
+            totalacc = accuracy.eval(feed_dict={
+                x: validation_images,
+                y: validation_labels,
+                keep_prob: 1.0
+            })
             new_weights = variables.get_weights()
     return float(totalacc), new_weights

examples/lbfgs/driver.py

@@ -2,14 +2,16 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import ray
 import numpy as np
-import scipy.optimize
-import tensorflow as tf
 import os
+import scipy.optimize
 
+import tensorflow as tf
 from tensorflow.examples.tutorials.mnist import input_data
 
+import ray
+import ray.experimental.tf_utils
+
 
 class LinearModel(object):
     """Simple class for a one layer neural network.
@@ -55,7 +57,7 @@ class LinearModel(object):
         # In order to get and set the weights, we pass in the loss function to
         # Ray's TensorFlowVariables to automatically create methods to modify
         # the weights.
-        self.variables = ray.experimental.TensorFlowVariables(
+        self.variables = ray.experimental.tf_utils.TensorFlowVariables(
             cross_entropy, self.sess)
 
     def loss(self, xs, ys):

examples/parameter_server/model.py

@@ -6,17 +6,20 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import ray
+import time
+
 import tensorflow as tf
 from tensorflow.examples.tutorials.mnist import input_data
-import time
+
+import ray
+import ray.experimental.tf_utils
 
 
 def download_mnist_retry(seed=0, max_num_retries=20):
     for _ in range(max_num_retries):
         try:
-            return input_data.read_data_sets("MNIST_data", one_hot=True,
-                                             seed=seed)
+            return input_data.read_data_sets(
+                "MNIST_data", one_hot=True, seed=seed)
         except tf.errors.AlreadyExistsError:
             time.sleep(1)
     raise Exception("Failed to download MNIST.")
@@ -42,30 +45,29 @@ class SimpleCNN(object):
 
             with tf.name_scope('adam_optimizer'):
                 self.optimizer = tf.train.AdamOptimizer(learning_rate)
-                self.train_step = self.optimizer.minimize(
-                    self.cross_entropy)
+                self.train_step = self.optimizer.minimize(self.cross_entropy)
 
             with tf.name_scope('accuracy'):
-                correct_prediction = tf.equal(tf.argmax(self.y_conv, 1),
-                                              tf.argmax(self.y_, 1))
+                correct_prediction = tf.equal(
+                    tf.argmax(self.y_conv, 1), tf.argmax(self.y_, 1))
                 correct_prediction = tf.cast(correct_prediction, tf.float32)
             self.accuracy = tf.reduce_mean(correct_prediction)
 
-            self.sess = tf.Session(config=tf.ConfigProto(
-                intra_op_parallelism_threads=1,
-                inter_op_parallelism_threads=1))
+            self.sess = tf.Session(
+                config=tf.ConfigProto(
+                    intra_op_parallelism_threads=1,
+                    inter_op_parallelism_threads=1))
             self.sess.run(tf.global_variables_initializer())
 
             # Helper values.
 
-            self.variables = ray.experimental.TensorFlowVariables(
+            self.variables = ray.experimental.tf_utils.TensorFlowVariables(
                 self.cross_entropy, self.sess)
 
-            self.grads = self.optimizer.compute_gradients(
-                self.cross_entropy)
-            self.grads_placeholder = [
-                (tf.placeholder("float", shape=grad[1].get_shape()), grad[1])
-                for grad in self.grads]
+            self.grads = self.optimizer.compute_gradients(self.cross_entropy)
+            self.grads_placeholder = [(tf.placeholder(
+                "float", shape=grad[1].get_shape()), grad[1])
+                                      for grad in self.grads]
             self.apply_grads_placeholder = self.optimizer.apply_gradients(
                 self.grads_placeholder)
 
@@ -73,17 +75,24 @@ class SimpleCNN(object):
         # TODO(rkn): Computing the weights before and after the training step
         # and taking the diff is awful.
         weights = self.get_weights()[1]
-        self.sess.run(self.train_step, feed_dict={self.x: x,
-                                                  self.y_: y,
-                                                  self.keep_prob: 0.5})
+        self.sess.run(
+            self.train_step,
+            feed_dict={
+                self.x: x,
+                self.y_: y,
+                self.keep_prob: 0.5
+            })
         new_weights = self.get_weights()[1]
         return [x - y for x, y in zip(new_weights, weights)]
 
     def compute_gradients(self, x, y):
-        return self.sess.run([grad[0] for grad in self.grads],
-                             feed_dict={self.x: x,
-                                        self.y_: y,
-                                        self.keep_prob: 0.5})
+        return self.sess.run(
+            [grad[0] for grad in self.grads],
+            feed_dict={
+                self.x: x,
+                self.y_: y,
+                self.keep_prob: 0.5
+            })
 
     def apply_gradients(self, gradients):
         feed_dict = {}
@@ -92,10 +101,13 @@ class SimpleCNN(object):
         self.sess.run(self.apply_grads_placeholder, feed_dict=feed_dict)
 
     def compute_accuracy(self, x, y):
-        return self.sess.run(self.accuracy,
-                             feed_dict={self.x: x,
-                                        self.y_: y,
-                                        self.keep_prob: 1.0})
+        return self.sess.run(
+            self.accuracy,
+            feed_dict={
+                self.x: x,
+                self.y_: y,
+                self.keep_prob: 1.0
+            })
 
     def set_weights(self, variable_names, weights):
         self.variables.set_weights(dict(zip(variable_names, weights)))
@@ -175,8 +187,8 @@ def conv2d(x, W):
 
 def max_pool_2x2(x):
     """max_pool_2x2 downsamples a feature map by 2X."""
-    return tf.nn.max_pool(x, ksize=[1, 2, 2, 1],
-                          strides=[1, 2, 2, 1], padding='SAME')
+    return tf.nn.max_pool(
+        x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
 
 
 def weight_variable(shape):

examples/resnet/resnet_model.py

@@ -13,14 +13,17 @@ from __future__ import print_function
 
 from collections import namedtuple
 import numpy as np
-import ray
+
 import tensorflow as tf
 from tensorflow.python.training import moving_averages
 
-HParams = namedtuple('HParams',
-                     'batch_size, num_classes, min_lrn_rate, lrn_rate, '
-                     'num_residual_units, use_bottleneck, weight_decay_rate, '
-                     'relu_leakiness, optimizer, num_gpus')
+import ray
+import ray.experimental.tf_utils
+
+HParams = namedtuple(
+    'HParams', 'batch_size, num_classes, min_lrn_rate, lrn_rate, '
+    'num_residual_units, use_bottleneck, weight_decay_rate, '
+    'relu_leakiness, optimizer, num_gpus')
 
 
 class ResNet(object):
@@ -51,7 +54,8 @@ class ResNet(object):
             self._build_train_op()
         else:
             # Additional initialization for the test network.
-            self.variables = ray.experimental.TensorFlowVariables(self.cost)
+            self.variables = ray.experimental.tf_utils.TensorFlowVariables(
+                self.cost)
             self.summaries = tf.summary.merge_all()
 
     def _stride_arr(self, stride):
@@ -75,27 +79,24 @@ class ResNet(object):
             filters = [16, 16, 32, 64]
 
         with tf.variable_scope('unit_1_0'):
-            x = res_func(x, filters[0], filters[1],
-                         self._stride_arr(strides[0]),
-                         activate_before_residual[0])
+            x = res_func(x, filters[0], filters[1], self._stride_arr(
+                strides[0]), activate_before_residual[0])
         for i in range(1, self.hps.num_residual_units):
             with tf.variable_scope('unit_1_%d' % i):
                 x = res_func(x, filters[1], filters[1], self._stride_arr(1),
                              False)
 
         with tf.variable_scope('unit_2_0'):
-            x = res_func(x, filters[1], filters[2],
-                         self._stride_arr(strides[1]),
-                         activate_before_residual[1])
+            x = res_func(x, filters[1], filters[2], self._stride_arr(
+                strides[1]), activate_before_residual[1])
         for i in range(1, self.hps.num_residual_units):
             with tf.variable_scope('unit_2_%d' % i):
-                x = res_func(x, filters[2], filters[2],
-                             self._stride_arr(1), False)
+                x = res_func(x, filters[2], filters[2], self._stride_arr(1),
+                             False)
 
         with tf.variable_scope('unit_3_0'):
-            x = res_func(x, filters[2], filters[3],
-                         self._stride_arr(strides[2]),
-                         activate_before_residual[2])
+            x = res_func(x, filters[2], filters[3], self._stride_arr(
+                strides[2]), activate_before_residual[2])
         for i in range(1, self.hps.num_residual_units):
             with tf.variable_scope('unit_3_%d' % i):
                 x = res_func(x, filters[3], filters[3], self._stride_arr(1),
@@ -136,7 +137,8 @@ class ResNet(object):
         apply_op = optimizer.minimize(self.cost, global_step=self.global_step)
         train_ops = [apply_op] + self._extra_train_ops
         self.train_op = tf.group(*train_ops)
-        self.variables = ray.experimental.TensorFlowVariables(self.train_op)
+        self.variables = ray.experimental.tf_utils.TensorFlowVariables(
+            self.train_op)
 
     def _batch_norm(self, name, x):
         """Batch normalization."""
@@ -144,49 +146,65 @@ class ResNet(object):
             params_shape = [x.get_shape()[-1]]
 
             beta = tf.get_variable(
-                'beta', params_shape, tf.float32,
+                'beta',
+                params_shape,
+                tf.float32,
                 initializer=tf.constant_initializer(0.0, tf.float32))
             gamma = tf.get_variable(
-                'gamma', params_shape, tf.float32,
+                'gamma',
+                params_shape,
+                tf.float32,
                 initializer=tf.constant_initializer(1.0, tf.float32))
 
             if self.mode == 'train':
                 mean, variance = tf.nn.moments(x, [0, 1, 2], name='moments')
 
                 moving_mean = tf.get_variable(
-                    'moving_mean', params_shape, tf.float32,
+                    'moving_mean',
+                    params_shape,
+                    tf.float32,
                     initializer=tf.constant_initializer(0.0, tf.float32),
                     trainable=False)
                 moving_variance = tf.get_variable(
-                    'moving_variance', params_shape, tf.float32,
+                    'moving_variance',
+                    params_shape,
+                    tf.float32,
                     initializer=tf.constant_initializer(1.0, tf.float32),
                     trainable=False)
 
                 self._extra_train_ops.append(
-                    moving_averages.assign_moving_average(moving_mean, mean,
-                                                          0.9))
+                    moving_averages.assign_moving_average(
+                        moving_mean, mean, 0.9))
                 self._extra_train_ops.append(
-                    moving_averages.assign_moving_average(moving_variance,
-                                                          variance, 0.9))
+                    moving_averages.assign_moving_average(
+                        moving_variance, variance, 0.9))
             else:
                 mean = tf.get_variable(
-                    'moving_mean', params_shape, tf.float32,
+                    'moving_mean',
+                    params_shape,
+                    tf.float32,
                     initializer=tf.constant_initializer(0.0, tf.float32),
                     trainable=False)
                 variance = tf.get_variable(
-                    'moving_variance', params_shape, tf.float32,
+                    'moving_variance',
+                    params_shape,
+                    tf.float32,
                     initializer=tf.constant_initializer(1.0, tf.float32),
                     trainable=False)
                 tf.summary.histogram(mean.op.name, mean)
                 tf.summary.histogram(variance.op.name, variance)
             # elipson used to be 1e-5. Maybe 0.001 solves NaN problem in deeper
             # net.
-            y = tf.nn.batch_normalization(
-                x, mean, variance, beta, gamma, 0.001)
+            y = tf.nn.batch_normalization(x, mean, variance, beta, gamma,
+                                          0.001)
             y.set_shape(x.get_shape())
             return y
 
-    def _residual(self, x, in_filter, out_filter, stride,
+    def _residual(self,
+                  x,
+                  in_filter,
+                  out_filter,
+                  stride,
                   activate_before_residual=False):
         """Residual unit with 2 sub layers."""
         if activate_before_residual:
@@ -212,14 +230,18 @@ class ResNet(object):
             if in_filter != out_filter:
                 orig_x = tf.nn.avg_pool(orig_x, stride, stride, 'VALID')
                 orig_x = tf.pad(
-                    orig_x, [[0, 0], [0, 0], [0, 0],
-                             [(out_filter - in_filter) // 2,
-                             (out_filter - in_filter) // 2]])
+                    orig_x,
+                    [[0, 0], [0, 0], [0, 0], [(out_filter - in_filter) // 2,
+                                              (out_filter - in_filter) // 2]])
             x += orig_x
 
         return x
 
-    def _bottleneck_residual(self, x, in_filter, out_filter, stride,
+    def _bottleneck_residual(self,
+                             x,
+                             in_filter,
+                             out_filter,
+                             stride,
                              activate_before_residual=False):
         """Bottleneck residual unit with 3 sub layers."""
         if activate_before_residual:
@@ -271,7 +293,8 @@ class ResNet(object):
             n = filter_size * filter_size * out_filters
             kernel = tf.get_variable(
                 'DW', [filter_size, filter_size, in_filters, out_filters],
-                tf.float32, initializer=tf.random_normal_initializer(
+                tf.float32,
+                initializer=tf.random_normal_initializer(
                     stddev=np.sqrt(2.0 / n)))
             return tf.nn.conv2d(x, kernel, strides, padding='SAME')
 
@@ -285,8 +308,8 @@ class ResNet(object):
         w = tf.get_variable(
             'DW', [x.get_shape()[1], out_dim],
             initializer=tf.uniform_unit_scaling_initializer(factor=1.0))
-        b = tf.get_variable('biases', [out_dim],
-                            initializer=tf.constant_initializer())
+        b = tf.get_variable(
+            'biases', [out_dim], initializer=tf.constant_initializer())
         return tf.nn.xw_plus_b(x, w, b)
 
     def _global_avg_pool(self, x):