wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-06 13:31:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

Alexnet example for Imagenet (#319)

Browse Source
This commit is contained in:
Robert Nishihara
2016-07-28 16:59:19 -07:00
committed by GitHub
parent 1961deeffb
commit f79494c8c6
9 changed files with 630 additions and 171 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -57,3 +57,4 @@ estimate of pi (waiting until the computation has finished if necessary).
- [Hyperparameter Optimization](examples/hyperopt/README.md)
- [Batch L-BFGS](examples/lbfgs/README.md)
- [Train AlexNet](examples/alexnet/README.md)
examples/alexnet/README.md
+108
View File
@@ -0,0 +1,108 @@
# Training AlexNet
WARNING: Running this application is fairly involved. In particular, it requires
you to download the ImageNet dataset and put it on S3.
This document walks through how to load the ImageNet dataset from S3 and train
AlexNet using data parallel stochastic gradient descent.
## Running the Application
The instructions in this section must be done before you can run the
application.
### Install the Dependencies
Install the following dependencies.
- [TensorFlow](https://www.tensorflow.org/)
In addition, install the following dependencies.
**On Ubuntu**
```
sudo apt-get install libjpeg8-dev awscli
sudo pip install boto3 botocore pillow
```
**On Mac OSX**
```
brew install libjpeg awscli
sudo pip install boto3 botocore pillow
```
### Put ImageNet on S3
1. To run this application, first put the ImageNet tar files on S3 (e.g., in the
directory `ILSVRC2012_img_train`). Also, put the file `train.txt` on S3. We will
use `$S3_BUCKET` to refer to the name of your S3 bucket.
2. Use `aws configure` to enable Boto to connect to S3. If you are using
multiple machines, this must be done on all machines.
### Run the Application
From the directory `ray/examples/alexnet/` run the following
```
source ../../setup-env.sh
python driver.py --s3-bucket $S3_BUCKET
```
## Parallel Data Loading
To speed up data loading, we will pull data from S3 in parallel with a number of
workers. At the core of our loading code is the remote function
`load_tarfile_from_s3`. When executed, this function connects to S3 and
retrieves the appropriate object.
```python
@ray.remote([str, str, List], [np.ndarray, List])
def load_tarfile_from_s3(bucket, s3_key, size=[]):
# Pull the object with the given key and bucket from S3, untar the contents,
# and return it.
return images, labels
```
To load data in parallel, we simply call this function multiple times with the
keys of all the objects that we want to retrieve. This returns a list of pairs
of object references, where the first object reference in each pair refers to a
batch of images and the second refers to the corresponding batch of labels.
```python
batches = [load_tarfile_from_s3(bucket, s3_key, size) for s3_key in s3_keys]
```
By default, this will only fetch objects whose keys have prefix
`ILSVRC2012_img_train/n015` (this is 13 tar files). To fetch all of the data,
pass in `--key-prefix ILSVRC2012_img_train/n`.
## Data Parallel Training
The other parallel component of this application is the training procedure. This
is built on top of the remote function `compute_grad`.
```python
@ray.remote([np.ndarray, np.ndarray, np.ndarray, List], [List])
def compute_grad(X, Y, mean, weights):
# Load the weights into the network.
# Subtract the mean and crop the images.
# Compute the gradients.
return gradients
```
This function takes training inputs and outputs, the mean image (to subtract off
of the input images), and the current network weights.
We can parallelize the computation of the gradient over multiple batches by
calling `compute_grad` multiple times in parallel.
```python
gradient_refs = []
for i in range(num_workers):
# Choose a random batch and use it to compute the gradient of the loss.
x_ref, y_ref = batches[np.random.randint(len(batches))]
gradient_refs.append(compute_grad(x_ref, y_ref, mean_ref, weights_ref))
```
examples/alexnet/alexnet.py
+428
View File
@@ -0,0 +1,428 @@
# The code for AlexNet is copied and adapted from the TensorFlow repository
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/models/image/alexnet/alexnet_benchmark.py.
import ray
import numpy as np
import tarfile, io
import boto3
import PIL.Image as Image
import tensorflow as tf
from typing import List, Tuple
import ray.array.remote as ra
STDDEV = 0.001 # The standard deviation of the network weight initialization.
def load_chunk(tarfile, size=None):
"""Load a number of images from a single imagenet .tar file.
This function also converts the image from grayscale to RGB if necessary.
Args:
tarfile (tarfile.TarFile): The archive from which the files get loaded.
size (Optional[Tuple[int, int]]): Resize the image to this size if provided.
Returns:
numpy.ndarray: Contains the image data in format [batch, w, h, c]
"""
result = []
filenames = []
for member in tarfile.getmembers():
filename = member.path
content = tarfile.extractfile(member)
img = Image.open(content)
rgbimg = Image.new("RGB", img.size)
rgbimg.paste(img)
if size != None:
rgbimg = rgbimg.resize(size, Image.ANTIALIAS)
result.append(np.array(rgbimg).reshape(1, rgbimg.size[0], rgbimg.size[1], 3))
filenames.append(filename)
return np.concatenate(result), filenames
@ray.remote([str, str, List], [np.ndarray, List])
def load_tarfile_from_s3(bucket, s3_key, size=[]):
"""Load an imagenet .tar file.
Args:
bucket (str): Bucket holding the imagenet .tar.
s3_key (str): s3 key from which the .tar file is loaded.
size (List[int]): Resize the image to this size if size != []; len(size) == 2 required.
Returns:
np.ndarray: The image data (see load_chunk).
"""
s3 = boto3.client("s3")
response = s3.get_object(Bucket=bucket, Key=s3_key)
output = io.BytesIO()
chunk = response["Body"].read(1024 * 8)
while chunk:
output.write(chunk)
chunk = response["Body"].read(1024 * 8)
output.seek(0) # go to the beginning of the .tar file
tar = tarfile.open(mode="r", fileobj=output)
return load_chunk(tar, size=size if size != [] else None)
def load_tarfiles_from_s3(bucket, s3_keys, size=[]):
"""Load a number of imagenet .tar files.
Args:
bucket (str): Bucket holding the imagenet .tars.
s3_keys (List[str]): List of s3 keys from which the .tar files are being loaded.
size (List[int]): Resize the image to this size if size != []; len(size) == 2 required.
Returns:
np.ndarray: Contains object references to the chunks of the images (see load_chunk).
"""
return [load_tarfile_from_s3(bucket, s3_key, size) for s3_key in s3_keys]
def setup_variables(params, placeholders, assigns, kernelshape, biasshape):
"""Creates the variables for each layer and adds the variables and the components needed to feed them to various lists
Args:
params (List): Network parameters used for creating feed_dicts
placeholders (List): Placeholders used for feeding weights into
assigns (List): Assignments used for actually setting variables
kernelshape (List): Shape of the kernel used for the conv layer
biasshape (List): Shape of the bias used
Returns:
None
"""
kernel = tf.Variable(tf.truncated_normal(kernelshape, stddev=STDDEV))
biases = tf.Variable(tf.constant(0.0, shape=biasshape, dtype=tf.float32),
trainable=True, name='biases')
kernel_new = tf.placeholder(tf.float32, shape=kernel.get_shape())
biases_new = tf.placeholder(tf.float32, shape=biases.get_shape())
update_kernel = kernel.assign(kernel_new)
update_biases = biases.assign(biases_new)
params += [kernel, biases]
placeholders += [kernel_new, biases_new]
assigns += [update_kernel, update_biases]
def conv_layer(parameters, prev_layer, shape, scope):
"""Constructs a convolutional layer for the network.
Args:
parameters (List): Parameters used in constructing layer.
prevlayer (Tensor): The previous layer to connect the network together.
shape (List): The strides used for convolution
scope (Scope): Current scope of tensorflow
Returns:
Tensor: Activation of layer
"""
kernel = parameters[-2]
bias = parameters[-1]
conv = tf.nn.conv2d(prev_layer, kernel, shape, padding='SAME')
add_bias = tf.nn.bias_add(conv, bias)
return tf.nn.relu(add_bias, name=scope)
def net_initialization():
images = tf.placeholder(tf.float32, shape=[None, 224, 224, 3])
y_true = tf.placeholder(tf.float32, shape=[None, 1000])
parameters = []
assignment = []
placeholders = []
# conv1
with tf.name_scope('conv1') as scope:
setup_variables(parameters, placeholders, assignment, [11, 11, 3, 96], [96])
conv1 = conv_layer(parameters, images, [1, 4, 4, 1], scope)
# pool1
pool1 = tf.nn.max_pool(conv1,
ksize=[1, 3, 3, 1],
strides=[1, 2, 2, 1],
padding='VALID',
name='pool1')
# lrn1
pool1_lrn = tf.nn.lrn(pool1, depth_radius=5, bias=1.0,
alpha=0.0001, beta=0.75,
name="LocalResponseNormalization")
# conv2
with tf.name_scope('conv2') as scope:
setup_variables(parameters, placeholders, assignment, [5, 5, 96, 256], [256])
conv2 = conv_layer(parameters, pool1_lrn, [1, 1, 1, 1], scope)
pool2 = tf.nn.max_pool(conv2,
ksize=[1, 3, 3, 1],
strides=[1, 2, 2, 1],
padding='VALID',
name='pool2')
# lrn2
pool2_lrn = tf.nn.lrn(pool2, depth_radius=5, bias=1.0,
alpha=0.0001, beta=0.75,
name="LocalResponseNormalization")
# conv3
with tf.name_scope('conv3') as scope:
setup_variables(parameters, placeholders, assignment, [3, 3, 256, 384], [384])
conv3 = conv_layer(parameters, pool2_lrn, [1, 1, 1, 1], scope)
# conv4
with tf.name_scope('conv4') as scope:
setup_variables(parameters, placeholders, assignment, [3, 3, 384, 384], [384])
conv4 = conv_layer(parameters, conv3, [1, 1, 1, 1], scope)
# conv5
with tf.name_scope('conv5') as scope:
setup_variables(parameters, placeholders, assignment, [3, 3, 384, 256], [256])
conv5 = conv_layer(parameters, conv4, [1, 1, 1, 1], scope)
# pool5
pool5 = tf.nn.max_pool(conv5,
ksize=[1, 3, 3, 1],
strides=[1, 2, 2, 1],
padding='VALID',
name='pool5')
# lrn5
pool5_lrn = tf.nn.lrn(pool5, depth_radius=5, bias=1.0,
alpha=0.0001, beta=0.75,
name="LocalResponseNormalization")
dropout = tf.placeholder(tf.float32)
with tf.name_scope('fc1') as scope:
n_input = int(np.prod(pool5_lrn.get_shape().as_list()[1:]))
setup_variables(parameters, placeholders, assignment, [n_input, 4096], [4096])
fc_in = tf.reshape(pool5_lrn, [-1, n_input])
fc_layer1 = tf.nn.tanh(tf.nn.bias_add(tf.matmul(fc_in, parameters[-2]), parameters[-1]))
fc_out1 = tf.nn.dropout(fc_layer1, dropout)
with tf.name_scope('fc2') as scope:
n_input = int(np.prod(fc_out1.get_shape().as_list()[1:]))
setup_variables(parameters, placeholders, assignment, [n_input, 4096], [4096])
fc_in = tf.reshape(fc_out1, [-1, n_input])
fc_layer2 = tf.nn.tanh(tf.nn.bias_add(tf.matmul(fc_in, parameters[-2]), parameters[-1]))
fc_out2 = tf.nn.dropout(fc_layer2, dropout)
with tf.name_scope('fc3') as scope:
n_input = int(np.prod(fc_out2.get_shape().as_list()[1:]))
setup_variables(parameters, placeholders, assignment, [n_input, 1000], [1000])
fc_in = tf.reshape(fc_out2, [-1, n_input])
fc_layer3 = tf.nn.softmax(tf.nn.bias_add(tf.matmul(fc_in, parameters[-2]), parameters[-1]))
y_pred = fc_layer3 / tf.reduce_sum(fc_layer3,
reduction_indices=len(fc_layer3.get_shape()) - 1,
keep_dims=True)
# manual computation of crossentropy
y_pred = tf.clip_by_value(y_pred, tf.cast(1e-10, dtype=tf.float32),
tf.cast(1. - 1e-10, dtype=tf.float32))
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_true * tf.log(y_pred),
reduction_indices=len(y_pred.get_shape()) - 1))
opt = tf.train.MomentumOptimizer(learning_rate=0.01, momentum=0.9) # Any other optimizier can be placed here
correct_pred = tf.equal(tf.argmax(y_pred, 1), tf.argmax(y_true, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
comp_grads = opt.compute_gradients(cross_entropy, parameters)
application = opt.apply_gradients(zip(placeholders,parameters))
sess = tf.Session()
init_all_variables = tf.initialize_all_variables()
return comp_grads, sess, application, accuracy, images, y_true, dropout, placeholders, parameters, assignment, init_all_variables
def net_reinitialization(net_vars):
return net_vars
@ray.remote([List], [int])
def num_images(batches):
"""Counts number of images in batches.
Args:
batches (List): Collection of batches of images and labels.
Returns:
int: The number of images
"""
shape_refs = [ra.shape(batch) for batch in batches]
return sum([ray.get(shape_ref)[0] for shape_ref in shape_refs])
@ray.remote([List], [np.ndarray])
def compute_mean_image(batches):
"""Computes the mean image given a list of batches of images.
Args:
batches (List[ObjRef]): A list of batches of images.
Returns:
ndarray: The mean image
"""
if len(batches) == 0:
raise Exception("No images were passed into `compute_mean_image`.")
sum_image_refs = [ra.sum(batch, axis=0) for batch in batches]
sum_images = [ray.get(ref) for ref in sum_image_refs]
n_images = num_images(batches)
return np.sum(sum_images, axis=0).astype("float64") / ray.get(n_images)
@ray.remote([np.ndarray, np.ndarray, np.ndarray, np.ndarray], [np.ndarray, np.ndarray, np.ndarray, np.ndarray])
def shuffle_arrays(first_images, first_labels, second_images, second_labels):
"""Shuffles the images and labels from two batches.
Args:
first_images (ndarray): First batch of images.
first_labels (ndarray): First batch of labels.
second_images (ndarray): Second batch of images.
second_labels (ndarray): Second batch of labels.
Returns:
ndarray: First batch of shuffled images.
ndarray: First batch of shuffled labels.
ndarray: Second bach of shuffled images.
ndarray: Second batch of shuffled labels.
"""
images = np.concatenate((first_images, second_images))
labels = np.concatenate((first_labels, second_labels))
total_length = len(images)
first_len = len(first_images)
random_indices = np.random.permutation(total_length)
new_first_images = images[random_indices[0:first_len]]
new_first_labels = labels[random_indices[0:first_len]]
new_second_images = images[random_indices[first_len:total_length]]
new_second_labels = labels[random_indices[first_len:total_length]]
return new_first_images, new_first_labels, new_second_images, new_second_labels
def shuffle_pair(first_batch, second_batch):
"""Shuffle two batches of data.
Args:
first_batch (Tuple[ObjRef. ObjRef]): The first batch to be shuffled. The
first component is the object reference of a batch of images, and the
second component is the object reference of the corresponding batch of
labels.
second_batch (Tuple[ObjRef, ObjRef]): The second batch to be shuffled. The
first component is the object reference of a batch of images, and the
second component is the object reference of the corresponding batch of
labels.
Returns:
Tuple[ObjRef, Objref]: The first batch of shuffled data.
Tuple[ObjRef, Objref]: Two second bach of shuffled data.
"""
images1, labels1, images2, labels2 = shuffle_arrays(first_batch[0], first_batch[1], second_batch[0], second_batch[1])
return (images1, labels1), (images2, labels2)
@ray.remote([list, dict], [np.ndarray])
def filenames_to_labels(filenames, filename_label_dict):
"""Converts filename strings to integer labels.
Args:
filenames (List[str]): The filenames of the images.
filename_label_dict (Dict[str, int]): A dictionary mapping filenames to
integer labels.
Returns:
ndarray: Integer labels
"""
return np.asarray([int(filename_label_dict[filename]) for filename in filenames])
def one_hot(x):
"""Converts integer labels to one hot vectors.
Args:
x (int): Index to be set to one
Returns:
ndarray: One hot vector.
"""
zero = np.zeros([1000])
zero[x] = 1.0
return zero
def crop_images(images):
"""Randomly crop a batch of images.
This is used to generate many slightly different images from each training
example.
Args:
images (ndarray): A batch of images to crop. The shape of images should be
batch_size x height x width x channels.
Returns:
ndarray: A batch of cropped images.
"""
original_height = 256
original_width = 256
cropped_height = 224
cropped_width = 224
height_offset = np.random.randint(original_height - cropped_height + 1)
width_offset = np.random.randint(original_width - cropped_width + 1)
return images[:, height_offset:(height_offset + cropped_height), width_offset:(width_offset + cropped_width), :]
def shuffle(batches):
"""Shuffle the data.
This method groups the batches together in pairs and within each pair shuffles
the data between the two members.
Args:
batches (List[Tuple[ObjRef, ObjRef]]): This is a list of tuples, where each
tuple consists of two object references. The first component is an object
reference for a batch of images, and the second component is an object
reference for the corresponding batch of labels.
Returns:
List[Tuple[ObjRef, ObjRef]]: The shuffled data.
"""
# Randomly permute the order of the batches.
permuted_batches = np.random.permutation(batches)
new_batches = []
for i in range(len(batches) / 2):
# Swap data between consecutive batches.
shuffled_batch1, shuffled_batch2 = shuffle_pair(permuted_batches[2 * i], permuted_batches[2 * i + 1])
new_batches += [shuffled_batch1, shuffled_batch2]
if len(batches) % 2 == 1:
# If there is an odd number of batches, don't forget the last one.
new_batches.append(permuted_batches[-1])
return new_batches
@ray.remote([np.ndarray, np.ndarray, np.ndarray, List], [List])
def compute_grad(X, Y, mean, weights):
"""Computes the gradient of the network.
Args:
X (ndarray): Numpy array of images in the form of [224, 224,3]
Y (ndarray): Labels corresponding to each image
mean (ndarray): Mean image to subtract from images
weights (List[ndarray]): The network weights.
Returns:
List of gradients for each variable
"""
comp_grads, sess, _, _, images, y_true, dropout, placeholders, _, assignment, _ = ray.reusables.net_vars
# Set the network weights.
feed_dict = dict(zip(placeholders, weights))
sess.run(assignment, feed_dict=feed_dict)
# Choose a subset of the batch to compute on and crop the images.
random_indices = np.random.randint(0, len(X), size=128)
subset_X = crop_images(X[random_indices] - mean)
subset_Y = np.asarray([one_hot(label) for label in Y[random_indices]])
# Compute the gradients.
return sess.run([g for (g, v) in comp_grads], feed_dict={images: subset_X, y_true: subset_Y, dropout: 0.5})
@ray.remote([np.ndarray, np.ndarray, List], [np.float32])
def compute_accuracy(X, Y, weights):
"""Returns the accuracy of the network
Args:
X (ndarray): A batch of images.
Y (ndarray): A batch of labels.
weights (List[ndarray]): The network weights.
Returns:
The accuracy of the network on the given batch.
"""
_, sess, _, accuracy, images, y_true, dropout, placeholders, _, assignment, _ = ray.reusables.net_vars
# Set the network weights.
feed_dict = dict(zip(placeholders, weights))
sess.run(assignment, feed_dict=feed_dict)
one_hot_Y = np.asarray([one_hot(label) for label in Y])
cropped_X = crop_images(X)
return sess.run(accuracy, feed_dict={images: cropped_X, y_true: one_hot_Y, dropout: 1.0})
examples/alexnet/driver.py
+93
View File
@@ -0,0 +1,93 @@
import numpy as np
import ray
import os
import argparse
import boto3
import alexnet
# Arguments to specify where the imagenet data is stored.
parser = argparse.ArgumentParser(description="Parse information for data loading.")
parser.add_argument("--s3-bucket", required=True, type=str, help="Name of the bucket that contains the image data.")
parser.add_argument("--key-prefix", default="ILSVRC2012_img_train/n015", type=str, help="Prefix for files to fetch.")
parser.add_argument("--label-file", default="train.txt", type=str, help="File containing labels")
if __name__ == "__main__":
args = parser.parse_args()
num_workers = 4
ray.services.start_ray_local(num_workers=num_workers)
# Note we do not do sess.run(tf.initialize_all_variables()) because that would
# result in a different initialization on each worker. Instead, we initialize
# the weights on the driver and load the weights on the workers every time we
# compute a gradient.
ray.reusables.net_vars = ray.Reusable(alexnet.net_initialization, alexnet.net_reinitialization)
# Prepare keys for downloading the data.
s3_resource = boto3.resource("s3")
imagenet_bucket = s3_resource.Bucket(args.s3_bucket)
objects = imagenet_bucket.objects.filter(Prefix=args.key_prefix)
image_tar_files = [str(obj.key) for obj in objects.all()]
print "Images will be downloaded from {} files.".format(len(image_tar_files))
# Downloading the label file, and create a dictionary mapping the filenames of
# the images to their labels.
s3_client = boto3.client("s3")
label_file = s3_client.get_object(Bucket=args.s3_bucket, Key=args.label_file)
filename_label_str = label_file["Body"].read().strip().split("\n")
filename_label_pairs = [line.split(" ") for line in filename_label_str]
filename_label_dict = dict([(os.path.basename(name), label) for name, label in filename_label_pairs])
filename_label_dict_ref = ray.put(filename_label_dict)
print "Labels extracted"
# Download the imagenet dataset.
imagenet_data = alexnet.load_tarfiles_from_s3(args.s3_bucket, image_tar_files, [256, 256])
# Convert the parsed filenames to integer labels and create batches.
batches = [(images, alexnet.filenames_to_labels(filenames, filename_label_dict_ref)) for images, filenames in imagenet_data]
# Compute the mean image.
mean_ref = alexnet.compute_mean_image([images for images, labels in batches])
# The data does not start out shuffled. Images of the same class all appear
# together, so we shuffle it ourselves here. Each shuffle pairs up the batches
# and swaps data within a pair.
num_shuffles = 5
for i in range(num_shuffles):
batches = alexnet.shuffle(batches)
_, sess, application, _, _, _, _, placeholders, parameters, assignment, init_all_variables = ray.reusables.net_vars
# Initialize the network and optimizer weights. This is only run once on the
# driver. We initialize the weights manually on the workers.
sess.run(init_all_variables)
print "Initialized network weights."
iteration = 0
while True:
# Extract weights from the local copy of the network.
weights = sess.run(parameters)
# Put weights in the object store.
weights_ref = ray.put(weights)
# Compute the accuracy on a random training batch.
x_ref, y_ref = batches[np.random.randint(len(batches))]
accuracy = alexnet.compute_accuracy(x_ref, y_ref, weights_ref)
# Launch tasks in parallel to compute the gradients for some batches.
gradient_refs = []
for i in range(num_workers - 1):
# Choose a random batch and use it to compute the gradient of the loss.
x_ref, y_ref = batches[np.random.randint(len(batches))]
gradient_refs.append(alexnet.compute_grad(x_ref, y_ref, mean_ref, weights_ref))
# Print the accuracy on a random training batch.
print "Iteration {}: accuracy = {:.3}%".format(iteration, 100 * ray.get(accuracy))
# Fetch the gradients. This blocks until the gradients have been computed.
gradient_sets = [ray.get(ref) for ref in gradient_refs]
# Average the gradients over all of the tasks.
mean_gradients = [np.mean([gradient_set[i] for gradient_set in gradient_sets], axis=0) for i in range(len(weights))]
# Use the gradients to update the network.
sess.run(application, feed_dict=dict(zip(placeholders, mean_gradients)))
iteration += 1
examples/imagenet/README.md
-14
View File
@@ -1,14 +0,0 @@
Dependencies for Imagenet
**On Ubuntu**
```
sudo apt-get install libjpeg8-dev
sudo pip install boto3 botocore pillow
```
**On Mac OSX**
```
brew install libjpeg
sudo pip install boto3 botocore pillow
```
examples/imagenet/driver.py
-35
View File
@@ -1,35 +0,0 @@
import argparse
import boto3
import os
import numpy as np
import ray
import imagenet
import functions
parser = argparse.ArgumentParser(description="Parse information for data loading.")
parser.add_argument("--s3-bucket", type=str, required=True, help="Name of the bucket that contains the image data.")
parser.add_argument("--key-prefix", default="ILSVRC2012_img_train/n015", type=str, help="Prefix for files to fetch.")
parser.add_argument("--drop-ipython", default=False, type=bool, help="Drop into IPython at the end?")
if __name__ == "__main__":
args = parser.parse_args()
worker_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py")
ray.services.start_ray_local(num_workers=5, worker_path=worker_path)
s3 = boto3.resource("s3")
imagenet_bucket = s3.Bucket(args.s3_bucket)
objects = imagenet_bucket.objects.filter(Prefix=args.key_prefix)
images = [obj.key for obj in objects.all()]
x = imagenet.load_tarfiles_from_s3(args.s3_bucket, map(str, images), [256, 256]) # TODO(pcm): implement unicode serialization
mean_image = functions.compute_mean_image(x)
mean_image = ray.get(mean_image)
print "The mean image is:"
print mean_image
if args.drop_ipython:
import IPython
IPython.embed()
examples/imagenet/functions.py
-18
View File
@@ -1,18 +0,0 @@
import numpy as np
from typing import List, Tuple
import ray
import ray.array.remote as ra
@ray.remote([List[Tuple[ray.ObjRef, ray.ObjRef]]], [int])
def num_images(batches):
shape_refs = [ra.shape(batch[0]) for batch in batches]
return sum([ray.get(shape_ref)[0] for shape_ref in shape_refs])
@ray.remote([List[Tuple[ray.ObjRef, ray.ObjRef]]], [np.ndarray])
def compute_mean_image(batches):
if len(batches) == 0:
raise Exception("No images were passed into `compute_mean_image`.")
sum_image_refs = [ra.sum(batch[0], axis=0) for batch in batches]
sum_images = [ray.get(ref) for ref in sum_image_refs]
n_images = num_images(batches)
return np.sum(sum_images, axis=0).astype("float64") / ray.get(n_images)
examples/imagenet/imagenet.py
-73
View File
@@ -1,73 +0,0 @@
import tarfile, io
from typing import List, Tuple
import PIL.Image
import numpy as np
import boto3
import ray
s3 = boto3.client("s3")
def load_chunk(tarfile, size=None):
"""Load a number of images from a single imagenet .tar file.
This function also converts the image from grayscale to RGB if neccessary.
Args:
tarfile (tarfile.TarFile): The archive from which the files get loaded.
size (Optional[Tuple[int, int]]): Resize the image to this size if provided.
Returns:
numpy.ndarray: Contains the image data in format [batch, w, h, c]
"""
result = []
filenames = []
for member in tarfile.getmembers():
filename = member.path
content = tarfile.extractfile(member)
img = PIL.Image.open(content)
rgbimg = PIL.Image.new("RGB", img.size)
rgbimg.paste(img)
if size != None:
rgbimg = rgbimg.resize(size, PIL.Image.ANTIALIAS)
result.append(np.array(rgbimg).reshape(1, rgbimg.size[0], rgbimg.size[1], 3))
filenames.append(filename)
return np.concatenate(result), filenames
@ray.remote([str, str, List[int]], [np.ndarray, List])
def load_tarfile_from_s3(bucket, s3_key, size=[]):
"""Load an imagenet .tar file.
Args:
bucket (str): Bucket holding the imagenet .tar.
s3_key (str): s3 key from which the .tar file is loaded.
size (List[int]): Resize the image to this size if size != []; len(size) == 2 required.
Returns:
np.ndarray: The image data (see load_chunk).
"""
response = s3.get_object(Bucket=bucket, Key=s3_key)
output = io.BytesIO()
chunk = response["Body"].read(1024 * 8)
while chunk:
output.write(chunk)
chunk = response["Body"].read(1024 * 8)
output.seek(0) # go to the beginning of the .tar file
tar = tarfile.open(mode="r", fileobj=output)
return load_chunk(tar, size=size if size != [] else None)
@ray.remote([str, List[str], List[int]], [List[Tuple[ray.ObjRef, ray.ObjRef]]])
def load_tarfiles_from_s3(bucket, s3_keys, size=[]):
"""Load a number of imagenet .tar files.
Args:
bucket (str): Bucket holding the imagenet .tars.
s3_keys (List[str]): List of s3 keys from which the .tar files are being loaded.
size (List[int]): Resize the image to this size if size != []; len(size) == 2 required.
Returns:
np.ndarray: Contains object references to the chunks of the images (see load_chunk).
"""
return [load_tarfile_from_s3(bucket, s3_key, size) for s3_key in s3_keys]
examples/imagenet/worker.py
-31
View File
@@ -1,31 +0,0 @@
import sys
import argparse
import numpy as np
import imagenet
import ray
import ray.array.remote as ra
import ray.array.distributed as da
import functions
parser = argparse.ArgumentParser(description="Parse addresses for the worker to connect to.")
parser.add_argument("--scheduler-address", default="127.0.0.1:10001", type=str, help="the scheduler's address")
parser.add_argument("--objstore-address", default="127.0.0.1:20001", type=str, help="the objstore's address")
parser.add_argument("--worker-address", default="127.0.0.1:40001", type=str, help="the worker's address")
if __name__ == "__main__":
args = parser.parse_args()
ray.worker.connect(args.scheduler_address, args.objstore_address, args.worker_address)
ray.register_module(imagenet)
ray.register_module(functions)
ray.register_module(ra)
ray.register_module(ra.random)
ray.register_module(ra.linalg)
ray.register_module(da)
ray.register_module(da.random)
ray.register_module(da.linalg)
ray.worker.main_loop()