load imagenet

test/datasets_test.py

@@ -0,0 +1,23 @@
+import os
+import tarfile
+import unittest
+import ray
+import ray.services as services
+import ray.datasets.imagenet as imagenet
+
+class ImageNetTest(unittest.TestCase):
+
+  def testImageNetLoading(self):
+    test_dir = os.path.dirname(os.path.abspath(__file__))
+    test_path = os.path.join(test_dir, "test_worker.py")
+    services.start_singlenode_cluster(return_drivers=False, num_workers_per_objstore=5, worker_path=test_path)
+
+    chunk_name = os.path.join(test_dir, "..", "data", "mini.tar")
+    tar = tarfile.open(chunk_name, mode= "r")
+    chunk = imagenet.load_chunk(tar, size=(256, 256))
+    self.assertEqual(chunk.shape, (2, 256, 256, 3))
+
+    services.cleanup()
+
+if __name__ == "__main__":
+    unittest.main()