From 31bf0e8da4188b228374e1893cba541d1d7b3d84 Mon Sep 17 00:00:00 2001 From: Wapaul1 Date: Mon, 15 May 2017 22:40:41 -0700 Subject: [PATCH] Improved the Resnet Example. (#551) * Initial updates * Mostly done * Now works with no arguments * Changed version check --- doc/source/example-resnet.rst | 22 +++--- examples/resnet/cifar_input.py | 15 ++-- examples/resnet/resnet_main.py | 131 ++++++++++++++++++-------------- examples/resnet/resnet_model.py | 10 +-- 4 files changed, 99 insertions(+), 79 deletions(-) diff --git a/doc/source/example-resnet.rst b/doc/source/example-resnet.rst index 3b0313c0a..0aff3c88b 100644 --- a/doc/source/example-resnet.rst +++ b/doc/source/example-resnet.rst @@ -4,7 +4,7 @@ ResNet This code adapts the `TensorFlow ResNet example`_ to do data parallel training across multiple GPUs using Ray. View the `code for this example`_. -To run the example, you will need to install `TensorFlow with GPU support`_ (at +To run the example, you will need to install `TensorFlow`_ (at least version ``1.0.0``). Then you can run the example as follows. First download the CIFAR-10 or CIFAR-100 dataset. @@ -55,7 +55,9 @@ The core of the script is the actor definition. @ray.remote(num_gpus=1) class ResNetTrainActor(object): - def __init__(self, path, num_gpus): + def __init__(self, data, dataset, num_gpus): + # data is the preprocessed images and labels extracted from the dataset. + # Thus, every actor has its own copy of the data. # Set the CUDA_VISIBLE_DEVICES environment variable in order to restrict # which GPUs TensorFlow uses. Note that this only works if it is done before # the call to tf.Session. @@ -63,24 +65,24 @@ The core of the script is the actor definition. with tf.Graph().as_default(): with tf.device('/gpu:0'): # We omit the code here that actually constructs the residual network - # and initializes it. + # and initializes it. Uses the definition in the Tensorflow Resnet Example. def compute_steps(self, weights): # This method sets the weights in the network, runs some training steps, - # and returns the new weights. - steps = 10 + # and returns the new weights. self.model.variables is a TensorFlowVariables + # class that we pass the train operation into. self.model.variables.set_weights(weights) - for i in range(steps): + for i in range(self.steps): self.model.variables.sess.run(self.model.train_op) return self.model.variables.get_weights() -The main script first creates one actor for each GPU. +The main script first creates one actor for each GPU, or a single actor if `num_gpus` is zero. .. code-block:: python - train_actors = [ResNetTrainActor.remote(train_data, num_gpus) for _ in range(num_gpus)] + train_actors = [ResNetTrainActor.remote(train_data, dataset, num_gpus) for _ in range(num_gpus)] -Then after initializing the actors with the same weights, the main loop performs +Then the main loop passes the same weights to every model, performs updates on each model, averages the updates, and puts the new weights in the object store. @@ -92,5 +94,5 @@ object store. weight_id = ray.put(mean_weights) .. _`TensorFlow ResNet example`: https://github.com/tensorflow/models/tree/master/resnet -.. _`TensorFlow with GPU support`: https://www.tensorflow.org/install/ +.. _`TensorFlow`: https://www.tensorflow.org/install/ .. _`code for this example`: https://github.com/ray-project/ray/tree/master/examples/resnet diff --git a/examples/resnet/cifar_input.py b/examples/resnet/cifar_input.py index bf7c18e80..6e251c5de 100644 --- a/examples/resnet/cifar_input.py +++ b/examples/resnet/cifar_input.py @@ -10,6 +10,16 @@ import numpy as np import tensorflow as tf def build_data(data_path, size, dataset): + """Creates the queue and preprocessing operations for the dataset. + + Args: + data_path: Filename for cifar10 data. + size: The number of images in the dataset. + dataset: The dataset we are using. + + Returns: + queue: A Tensorflow queue for extracting the images and labels. + """ image_size = 32 if dataset == 'cifar10': label_bytes = 1 @@ -66,10 +76,6 @@ def build_input(data, batch_size, dataset, train): image, image_size+4, image_size+4) image = tf.random_crop(image, [image_size, image_size, 3]) image = tf.image.random_flip_left_right(image) - # Brightness/saturation/constrast provides small gains .2%~.5% on cifar. - # image = tf.image.random_brightness(image, max_delta=63. / 255.) - # image = tf.image.random_saturation(image, lower=0.5, upper=1.5) - # image = tf.image.random_contrast(image, lower=0.2, upper=1.8) image = tf.image.per_image_standardization(image) example_queue = tf.RandomShuffleQueue( capacity=16 * batch_size, @@ -87,7 +93,6 @@ def build_input(data, batch_size, dataset, train): shapes=[[image_size, image_size, depth], [1]]) num_threads = 1 - example_enqueue_op = example_queue.enqueue([image, label]) tf.train.add_queue_runner(tf.train.queue_runner.QueueRunner( example_queue, [example_enqueue_op] * num_threads)) diff --git a/examples/resnet/resnet_main.py b/examples/resnet/resnet_main.py index 356752e08..ca9de7c4a 100644 --- a/examples/resnet/resnet_main.py +++ b/examples/resnet/resnet_main.py @@ -6,6 +6,7 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +import argparse import os import numpy as np import ray @@ -14,35 +15,37 @@ import tensorflow as tf import cifar_input import resnet_model -FLAGS = tf.app.flags.FLAGS -tf.app.flags.DEFINE_string('dataset', 'cifar10', 'cifar10 or cifar100.') -tf.app.flags.DEFINE_string('train_data_path', '', - 'Filepattern for training data.') -tf.app.flags.DEFINE_string('eval_data_path', '', - 'Filepattern for eval data') -tf.app.flags.DEFINE_string('eval_dir', '', - 'Directory to keep eval outputs.') -tf.app.flags.DEFINE_integer('eval_batch_count', 50, - 'Number of batches to eval.') -tf.app.flags.DEFINE_integer('num_gpus', 0, - 'Number of gpus used for training.') +# Tensorflow must be at least version 1.0.0 for the example to work. +if int(tf.__version__.split('.')[0]) < 1: + raise Exception('Your Tensorflow version is less than 1.0.0. Please update Tensorflow to the latest version.') + +parser = argparse.ArgumentParser(description="Run the hyperparameter optimization example.") +parser.add_argument("--dataset", default='cifar10', type=str, help="Dataset to use: cifar10 or cifar100.") +parser.add_argument("--train_data_path", default='cifar-10-batches-bin/data_batch*', type=str, help="Data path for the training data.") +parser.add_argument("--eval_data_path", default='cifar-10-batches-bin/test_batch.bin', type=str, help="Data path for the testing data.") +parser.add_argument("--eval_dir", default='/tmp/resnet-model/eval', type=str, help="Data path for the tensorboard logs.") +parser.add_argument("--eval_batch_count", default=50, type=int, help="Number of batches to evaluate over.") +parser.add_argument("--num_gpus", default=0, type=int, help="Number of GPUs to use for training.") + +FLAGS = parser.parse_args() + +# Determines if the actors require a gpu or not. use_gpu = 1 if int(FLAGS.num_gpus) > 0 else 0 -@ray.remote(num_return_vals=4) +@ray.remote def get_data(path, size, dataset): - os.environ['CUDA_VISIBLE_DEVICES'] = '' - with tf.device('/cpu:0'): - queue = cifar_input.build_data(path, size, dataset) - sess = tf.Session() - coord = tf.train.Coordinator() - tf.train.start_queue_runners(sess, coord=coord) - images, labels = sess.run(queue) - coord.request_stop() - sess.close() - return (images[:int(size / 3), :], - images[int(size / 3):int(2 * size / 3), :], - images[int(2 * size / 3):, :], - labels) + # Retrieves all preprocessed images and labels using a tensorflow queue. + # This only uses the cpu. + os.environ['CUDA_VISIBLE_DEVICES'] = '' + with tf.device('/cpu:0'): + queue = cifar_input.build_data(path, size, dataset) + sess = tf.Session() + coord = tf.train.Coordinator() + tf.train.start_queue_runners(sess, coord=coord) + images, labels = sess.run(queue) + coord.request_stop() + sess.close() + return images, labels @ray.remote(num_gpus=use_gpu) class ResNetTrainActor(object): @@ -50,7 +53,7 @@ class ResNetTrainActor(object): if num_gpus > 0: os.environ['CUDA_VISIBLE_DEVICES'] = ','.join([str(i) for i in ray.get_gpu_ids()]) hps = resnet_model.HParams(batch_size=128, - num_classes=10 if dataset == 'cifar10' else 100, + num_classes=100 if dataset == 'cifar100' else 10, min_lrn_rate=0.0001, lrn_rate=0.1, num_residual_units=5, @@ -59,14 +62,19 @@ class ResNetTrainActor(object): relu_leakiness=0.1, optimizer='mom', num_gpus=num_gpus) - data = ray.get(data) - total_images = np.concatenate([data[0], data[1], data[2]]) + + # We seed each actor differently so that each actor operates on a different subset of data. if num_gpus > 0: tf.set_random_seed(ray.get_gpu_ids()[0] + 1) else: + # Only a single actor in this case. tf.set_random_seed(1) + + input_images = data[0] + input_labels = data[1] with tf.device('/gpu:0' if num_gpus > 0 else '/cpu:0'): - images, labels = cifar_input.build_input([total_images, data[3]], hps.batch_size, dataset, True) + # Build the model. + images, labels = cifar_input.build_input([input_images, input_labels], hps.batch_size, dataset, False) self.model = resnet_model.ResNet(hps, images, labels, 'train') self.model.build_graph() config = tf.ConfigProto(allow_soft_placement=True) @@ -76,24 +84,26 @@ class ResNetTrainActor(object): tf.train.start_queue_runners(sess, coord=self.coord) init = tf.global_variables_initializer() sess.run(init) + self.steps = 10 def compute_steps(self, weights): - # This method sets the weights in the network, runs some training steps, + # This method sets the weights in the network, trains the network self.steps times, # and returns the new weights. - steps = 10 self.model.variables.set_weights(weights) - for i in range(steps): + for i in range(self.steps): self.model.variables.sess.run(self.model.train_op) return self.model.variables.get_weights() def get_weights(self): + # Note that the driver cannot directly access fields of the class, + # so helper methods must be created. return self.model.variables.get_weights() @ray.remote class ResNetTestActor(object): def __init__(self, data, dataset, eval_batch_count, eval_dir): hps = resnet_model.HParams(batch_size=100, - num_classes=10 if dataset == 'cifar10' else 100, + num_classes=100 if dataset == 'cifar100' else 10, min_lrn_rate=0.0001, lrn_rate=0.1, num_residual_units=5, @@ -102,10 +112,11 @@ class ResNetTestActor(object): relu_leakiness=0.1, optimizer='mom', num_gpus=0) - data = ray.get(data) - total_images = np.concatenate([data[0], data[1], data[2]]) + input_images = data[0] + input_labels = data[1] with tf.device('/cpu:0'): - images, labels = cifar_input.build_input([total_images, data[3]], hps.batch_size, dataset, False) + # Builds the testing network. + images, labels = cifar_input.build_input([input_images, input_labels], hps.batch_size, dataset, False) self.model = resnet_model.ResNet(hps, images, labels, 'eval') self.model.build_graph() config = tf.ConfigProto(allow_soft_placement=True) @@ -115,13 +126,17 @@ class ResNetTestActor(object): tf.train.start_queue_runners(sess, coord=self.coord) init = tf.global_variables_initializer() sess.run(init) + + # Initializing parameters for tensorboard. self.best_precision = 0.0 self.eval_batch_count = eval_batch_count self.summary_writer = tf.summary.FileWriter(eval_dir, sess.graph) - self.summary_writer + # The IP address where tensorboard logs will be on. self.ip_addr = ray.services.get_node_ip_address() def accuracy(self, weights, train_step): + # Sets the weights, computes the accuracy and other metrics + # over eval_batches, and outputs to tensorboard. self.model.variables.set_weights(weights) total_prediction, correct_prediction = 0, 0 model = self.model @@ -153,38 +168,42 @@ class ResNetTestActor(object): return precision def get_ip_addr(self): + # As above, a helper method must be created to access the field from the driver. return self.ip_addr def train(): - """Training loop.""" - num_gpus = int(FLAGS.num_gpus) + num_gpus = FLAGS.num_gpus ray.init(num_gpus=num_gpus, redirect_output=True) train_data = get_data.remote(FLAGS.train_data_path, 50000, FLAGS.dataset) test_data = get_data.remote(FLAGS.eval_data_path, 10000, FLAGS.dataset) - if num_gpus > 0: + # Creates an actor for each gpu, or one if only using the cpu. Each actor has its own copy of the dataset. + if FLAGS.num_gpus > 0: train_actors = [ResNetTrainActor.remote(train_data, FLAGS.dataset, num_gpus) for _ in range(num_gpus)] else: - train_actors = [ResNetTrainActor.remote(train_data, num_gpus, 0)] + train_actors = [ResNetTrainActor.remote(train_data, FLAGS.dataset, 0)] test_actor = ResNetTestActor.remote(test_data, FLAGS.dataset, FLAGS.eval_batch_count, FLAGS.eval_dir) print('The log files for tensorboard are stored at ip {}.'.format(ray.get(test_actor.get_ip_addr.remote()))) step = 0 weight_id = train_actors[0].get_weights.remote() acc_id = test_actor.accuracy.remote(weight_id, step) + # Correction for dividing the weights by the number of gpus. if num_gpus == 0: num_gpus = 1 - print("Starting computation.") - while True: - all_weights = ray.get([actor.compute_steps.remote(weight_id) for actor in train_actors]) - mean_weights = {k: sum([weights[k] for weights in all_weights]) / num_gpus for k in all_weights[0]} - weight_id = ray.put(mean_weights) - step += 10 - if step % 200 == 0: - acc = ray.get(acc_id) - acc_id = test_actor.accuracy.remote(weight_id, step) - print('Step {0}: {1:.6f}'.format(step - 200, acc)) - -def main(_): - train() + print("Starting training loop. Use Ctrl-C to exit.") + try: + while True: + all_weights = ray.get([actor.compute_steps.remote(weight_id) for actor in train_actors]) + mean_weights = {k: sum([weights[k] for weights in all_weights]) / num_gpus for k in all_weights[0]} + weight_id = ray.put(mean_weights) + step += 10 + if step % 200 == 0: + # Retrieves the previously computed accuracy and launches a new + # testing task with the current weights every 200 steps. + acc = ray.get(acc_id) + acc_id = test_actor.accuracy.remote(weight_id, step) + print('Step {0}: {1:.6f}'.format(step - 200, acc)) + except KeyboardInterrupt: + pass if __name__ == '__main__': - tf.app.run() + train() diff --git a/examples/resnet/resnet_model.py b/examples/resnet/resnet_model.py index 75cbf3840..9687e064e 100644 --- a/examples/resnet/resnet_model.py +++ b/examples/resnet/resnet_model.py @@ -49,6 +49,7 @@ class ResNet(object): if self.mode == 'train': self._build_train_op() else: + # Additional initialization for the test network. self.variables = ray.experimental.TensorFlowVariables(self.cost) self.summaries = tf.summary.merge_all() @@ -60,8 +61,7 @@ class ResNet(object): """Build the core model within the graph.""" with tf.variable_scope('init'): - x = self._images - x = self._conv('init_conv', x, 3, 3, 16, self._stride_arr(1)) + x = self._conv('init_conv', self._images, 3, 3, 16, self._stride_arr(1)) strides = [1, 2, 2] activate_before_residual = [True, False, False] @@ -71,12 +71,6 @@ class ResNet(object): else: res_func = self._residual filters = [16, 16, 32, 64] - # Uncomment the following codes to use w28-10 wide residual network. - # It is more memory efficient than very deep residual network and has - # comparably good performance. - # https://arxiv.org/pdf/1605.07146v1.pdf - # filters = [16, 160, 320, 640] - # Update hps.num_residual_units to 9 with tf.variable_scope('unit_1_0'): x = res_func(x, filters[0], filters[1], self._stride_arr(strides[0]),