mirror of
https://github.com/wassname/ray.git
synced 2026-07-14 11:17:54 +08:00
Resnet Example Uses tf.Datasets now (#960)
Change Resnet example to use tf.Datasets instead of queues.
This commit is contained in:
committed by
Robert Nishihara
parent
5c70faf76b
commit
c26c7553bc
@@ -31,24 +31,27 @@ def build_data(data_path, size, dataset):
|
||||
image_bytes = image_size * image_size * depth
|
||||
record_bytes = label_bytes + label_offset + image_bytes
|
||||
|
||||
data_files = tf.gfile.Glob(data_path)
|
||||
file_queue = tf.train.string_input_producer(data_files, shuffle=True)
|
||||
def load_transform(value):
|
||||
# Convert these examples to dense labels and processed images.
|
||||
record = tf.reshape(tf.decode_raw(value, tf.uint8), [record_bytes])
|
||||
label = tf.cast(tf.slice(record, [label_offset], [label_bytes]),
|
||||
tf.int32)
|
||||
# Convert from string to [depth * height * width] to
|
||||
# [depth, height, width].
|
||||
depth_major = tf.reshape(
|
||||
tf.slice(record, [label_bytes], [image_bytes]),
|
||||
[depth, image_size, image_size])
|
||||
# Convert from [depth, height, width] to [height, width, depth].
|
||||
image = tf.cast(tf.transpose(depth_major, [1, 2, 0]), tf.float32)
|
||||
return (image, label)
|
||||
# Read examples from files in the filename queue.
|
||||
reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)
|
||||
_, value = reader.read(file_queue)
|
||||
|
||||
# Convert these examples to dense labels and processed images.
|
||||
record = tf.reshape(tf.decode_raw(value, tf.uint8), [record_bytes])
|
||||
label = tf.cast(tf.slice(record, [label_offset], [label_bytes]), tf.int32)
|
||||
# Convert from string to [depth * height * width] to
|
||||
# [depth, height, width].
|
||||
depth_major = tf.reshape(tf.slice(record, [label_bytes], [image_bytes]),
|
||||
[depth, image_size, image_size])
|
||||
# Convert from [depth, height, width] to [height, width, depth].
|
||||
image = tf.cast(tf.transpose(depth_major, [1, 2, 0]), tf.float32)
|
||||
queue = tf.train.shuffle_batch([image, label], size, size, 0,
|
||||
num_threads=16)
|
||||
return queue
|
||||
data_files = tf.gfile.Glob(data_path)
|
||||
data = tf.contrib.data.FixedLengthRecordDataset(data_files,
|
||||
record_bytes=record_bytes)
|
||||
data = data.map(load_transform)
|
||||
data = data.batch(size)
|
||||
iterator = data.make_one_shot_iterator()
|
||||
return iterator.get_next()
|
||||
|
||||
|
||||
def build_input(data, batch_size, dataset, train):
|
||||
@@ -67,42 +70,35 @@ def build_input(data, batch_size, dataset, train):
|
||||
Raises:
|
||||
ValueError: When the specified dataset is not supported.
|
||||
"""
|
||||
images_constant = tf.constant(data[0])
|
||||
labels_constant = tf.constant(data[1])
|
||||
image_size = 32
|
||||
depth = 3
|
||||
num_classes = 10 if dataset == "cifar10" else 100
|
||||
image, label = tf.train.slice_input_producer([images_constant,
|
||||
labels_constant],
|
||||
capacity=16 * batch_size)
|
||||
if train:
|
||||
images, labels = data
|
||||
num_samples = images.shape[0] - images.shape[0] % batch_size
|
||||
dataset = tf.contrib.data.Dataset.from_tensor_slices(
|
||||
(images[:num_samples], labels[:num_samples]))
|
||||
|
||||
def map_train(image, label):
|
||||
image = tf.image.resize_image_with_crop_or_pad(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)
|
||||
image = tf.image.per_image_standardization(image)
|
||||
example_queue = tf.RandomShuffleQueue(
|
||||
capacity=16 * batch_size,
|
||||
min_after_dequeue=8 * batch_size,
|
||||
dtypes=[tf.float32, tf.int32],
|
||||
shapes=[[image_size, image_size, depth], [1]])
|
||||
num_threads = 16
|
||||
else:
|
||||
return (image, label)
|
||||
|
||||
def map_test(image, label):
|
||||
image = tf.image.resize_image_with_crop_or_pad(image, image_size,
|
||||
image_size)
|
||||
image = tf.image.per_image_standardization(image)
|
||||
example_queue = tf.FIFOQueue(
|
||||
3 * batch_size,
|
||||
dtypes=[tf.float32, tf.int32],
|
||||
shapes=[[image_size, image_size, depth], [1]])
|
||||
num_threads = 1
|
||||
return (image, label)
|
||||
|
||||
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))
|
||||
|
||||
# Read "batch" labels + images from the example queue.
|
||||
images, labels = example_queue.dequeue_many(batch_size)
|
||||
dataset = dataset.map(map_train if train else map_test)
|
||||
dataset = dataset.batch(batch_size)
|
||||
dataset = dataset.repeat()
|
||||
if train:
|
||||
dataset = dataset.shuffle(buffer_size=16 * batch_size)
|
||||
images, labels = dataset.make_one_shot_iterator().get_next()
|
||||
images = tf.reshape(images, [batch_size, image_size, image_size, depth])
|
||||
labels = tf.reshape(labels, [batch_size, 1])
|
||||
indices = tf.reshape(tf.range(0, batch_size, 1), [batch_size, 1])
|
||||
labels = tf.sparse_to_dense(
|
||||
|
||||
@@ -15,9 +15,11 @@ import tensorflow as tf
|
||||
import cifar_input
|
||||
import resnet_model
|
||||
|
||||
# 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 "
|
||||
# Tensorflow must be at least version 1.2.0 for the example to work.
|
||||
tf_major = int(tf.__version__.split(".")[0])
|
||||
tf_minor = int(tf.__version__.split(".")[1])
|
||||
if (tf_major < 1) or (tf_major == 1 and tf_minor < 2):
|
||||
raise Exception("Your Tensorflow version is less than 1.2.0. Please "
|
||||
"update Tensorflow to the latest version.")
|
||||
|
||||
parser = argparse.ArgumentParser(description="Run the ResNet example.")
|
||||
@@ -50,12 +52,9 @@ def get_data(path, size, dataset):
|
||||
# This only uses the cpu.
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
||||
with tf.device("/cpu:0"):
|
||||
queue = cifar_input.build_data(path, size, dataset)
|
||||
dataset = 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()
|
||||
images, labels = sess.run(dataset)
|
||||
sess.close()
|
||||
return images, labels
|
||||
|
||||
@@ -86,12 +85,9 @@ class ResNetTrainActor(object):
|
||||
# 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"):
|
||||
# Build the model.
|
||||
images, labels = cifar_input.build_input([input_images,
|
||||
input_labels],
|
||||
images, labels = cifar_input.build_input(data,
|
||||
hps.batch_size, dataset,
|
||||
False)
|
||||
self.model = resnet_model.ResNet(hps, images, labels, "train")
|
||||
@@ -100,8 +96,6 @@ class ResNetTrainActor(object):
|
||||
config.gpu_options.allow_growth = True
|
||||
sess = tf.Session(config=config)
|
||||
self.model.variables.set_session(sess)
|
||||
self.coord = tf.train.Coordinator()
|
||||
tf.train.start_queue_runners(sess, coord=self.coord)
|
||||
init = tf.global_variables_initializer()
|
||||
sess.run(init)
|
||||
self.steps = 10
|
||||
@@ -123,6 +117,7 @@ class ResNetTrainActor(object):
|
||||
@ray.remote
|
||||
class ResNetTestActor(object):
|
||||
def __init__(self, data, dataset, eval_batch_count, eval_dir):
|
||||
os.environ["CUDA_VISIBLE_DEVICES"] = ""
|
||||
hps = resnet_model.HParams(
|
||||
batch_size=100,
|
||||
num_classes=100 if dataset == "cifar100" else 10,
|
||||
@@ -134,12 +129,9 @@ class ResNetTestActor(object):
|
||||
relu_leakiness=0.1,
|
||||
optimizer="mom",
|
||||
num_gpus=0)
|
||||
input_images = data[0]
|
||||
input_labels = data[1]
|
||||
with tf.device("/cpu:0"):
|
||||
# Builds the testing network.
|
||||
images, labels = cifar_input.build_input([input_images,
|
||||
input_labels],
|
||||
images, labels = cifar_input.build_input(data,
|
||||
hps.batch_size, dataset,
|
||||
False)
|
||||
self.model = resnet_model.ResNet(hps, images, labels, "eval")
|
||||
@@ -148,8 +140,6 @@ class ResNetTestActor(object):
|
||||
config.gpu_options.allow_growth = True
|
||||
sess = tf.Session(config=config)
|
||||
self.model.variables.set_session(sess)
|
||||
self.coord = tf.train.Coordinator()
|
||||
tf.train.start_queue_runners(sess, coord=self.coord)
|
||||
init = tf.global_variables_initializer()
|
||||
sess.run(init)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user