wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-08 11:25:28 +08:00
Code Issues Packages Projects Releases Wiki Activity

Improved the Resnet Example. (#551)

Browse Source 
* Initial updates

* Mostly done

* Now works with no arguments

* Changed version check
This commit is contained in:
Wapaul1
2017-05-15 22:40:41 -07:00
committed by Robert Nishihara
parent 08e988aee5
commit 31bf0e8da4
4 changed files with 99 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/source/example-resnet.rst
+12 -10
View File
@@ -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
examples/resnet/cifar_input.py
+10 -5
View File
@@ -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))
examples/resnet/resnet_main.py
+75 -56
View File
@@ -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()
examples/resnet/resnet_model.py
+2 -8
View File
@@ -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]),