load imagenet

examples/imagenet/driver.py

@@ -0,0 +1,38 @@
+import argparse
+import boto3
+import os
+import numpy as np
+import ray
+import ray.services as services
+import ray.datasets.imagenet as imagenet
+
+import functions
+
+parser = argparse.ArgumentParser(description="Parse information for data loading.")
+parser.add_argument("--s3-bucket", 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("--drop-ipython", default=False, type=bool, help="Drop into IPython at the end?")
+
+if __name__ == "__main__":
+  args = parser.parse_args()
+  test_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "worker.py")
+  services.start_singlenode_cluster(return_drivers=False, num_workers_per_objstore=5, worker_path=test_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.pull(mean_image)
+
+  print "The mean image is:"
+  print mean_image
+
+  if args.drop_ipython:
+    import IPython
+    IPython.embed()
+
+  services.cleanup()

examples/imagenet/functions.py

@@ -0,0 +1,18 @@
+import numpy as np
+from typing import List
+import ray
+import ray.arrays.remote as ra
+
+@ray.remote([List[ray.ObjRef]], [int])
+def num_images(batches):
+  shape_refs = [ra.shape(batch) for batch in batches]
+  return sum([ray.pull(shape_ref)[0] for shape_ref in shape_refs])
+
+@ray.remote([List[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, axis=0) for batch in batches]
+  sum_images = [ray.pull(ref) for ref in sum_image_refs]
+  n_images = num_images(batches)
+  return np.sum(sum_images, axis=0).astype("float64") / ray.pull(n_images)

examples/imagenet/worker.py

@@ -0,0 +1,33 @@
+import sys
+import argparse
+import numpy as np
+
+import ray.datasets.imagenet
+
+import ray
+import ray.services as services
+import ray.worker as worker
+import ray.arrays.remote as ra
+import ray.arrays.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()
+  worker.connect(args.scheduler_address, args.objstore_address, args.worker_address)
+
+  ray.register_module(ray.datasets.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)
+
+  worker.main_loop()