def testPrefetchDictToDevice(self): host_dataset = dataset_ops.Dataset.range(10).map(lambda x: {"a": x}) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = dataset_ops.make_one_shot_iterator(device_dataset) next_element = iterator.get_next() self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(device_dataset))) self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(iterator))) self.assertEqual(dtypes.int64, next_element["a"].dtype) self.assertEqual([], next_element["a"].shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): for i in range(10): self.assertEqual({"a": i}, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchSparseTensorsToDevice(self): def make_tensor(i): return sparse_tensor.SparseTensorValue( indices=[[0, 0]], values=(i*[1]), dense_shape=[2, 2]) host_dataset = dataset_ops.Dataset.range(10).map(make_tensor) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) # NOTE(mrry): This device block creates the "host" dataset and iterator on # /cpu:0, and ensures that the prefetching is across devices. In typical use # this would not be necessary, because the GPU device would not support any # of the dataset-related ops. with ops.device("/cpu:0"): iterator = device_dataset.make_one_shot_iterator() self.assertTrue(dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(device_dataset))) self.assertTrue(dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(iterator))) next_element = iterator.get_next() self.assertEqual(dtypes.int64, next_element.dtype) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config) as sess: for i in range(10): actual = sess.run(next_element) self.assertAllEqual([i], actual.values) self.assertAllEqual([[0, 0]], actual.indices) self.assertAllEqual([2, 2], actual.dense_shape) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchDictToDevice(self): host_dataset = dataset_ops.Dataset.range(10).map(lambda x: {"a": x}) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = dataset_ops.make_one_shot_iterator(device_dataset) next_element = iterator.get_next() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) self.assertEqual(dtypes.int64, next_element["a"].dtype) self.assertEqual([], next_element["a"].shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): for i in range(10): self.assertEqual({"a": i}, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToDeviceWithReInit(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = dataset_ops.make_initializable_iterator(device_dataset) next_element = iterator.get_next() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): self.evaluate(iterator.initializer) for i in range(5): self.assertEqual(i, self.evaluate(next_element)) self.evaluate(iterator.initializer) for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToSameDevice(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device( "/job:localhost/replica:0/task:0/device:CPU:0")) with ops.device("/cpu:1"): iterator = dataset_ops.make_one_shot_iterator(device_dataset) next_element = iterator.get_next() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) with self.cached_session(): for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToSameDevice(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device( "/job:localhost/replica:0/task:0/device:CPU:0")) # NOTE(mrry): This device block creates the "host" dataset and iterator on # /cpu:0, and ensures that the prefetching is across devices. In typical use # this would not be necessary, because the GPU device would not support any # of the dataset-related ops. with ops.device("/cpu:0"): iterator = device_dataset.make_one_shot_iterator() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) next_element = iterator.get_next() self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) with self.cached_session() as sess: for i in range(10): self.assertEqual(i, sess.run(next_element)) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchToDeviceWithReInit(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = dataset_ops.make_initializable_iterator(device_dataset) next_element = iterator.get_next() self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(device_dataset))) self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(iterator))) self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): self.evaluate(iterator.initializer) for i in range(5): self.assertEqual(i, self.evaluate(next_element)) self.evaluate(iterator.initializer) for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testIteratorOnDeviceGraphModeInitializableIterator(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0")) iterator = dataset_ops.make_initializable_iterator(dataset) data = iterator.get_next() optional_data = iterator.get_next_as_optional() with ops.colocate_with(dataset._variant_tensor): dataset_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(dataset_device)) with ops.colocate_with(iterator._iterator_resource): iterator_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(iterator_device)) with ops.colocate_with(data): data_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(data_device)) with ops.colocate_with(optional_data.get_value()): get_value_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(get_value_device)) with ops.colocate_with(optional_data.has_value()): has_value_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(has_value_device))
def testPrefetchToSameDevice(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device( "/job:localhost/replica:0/task:0/device:CPU:0")) with ops.device("/cpu:1"): iterator = dataset_ops.make_one_shot_iterator(device_dataset) next_element = iterator.get_next() self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(device_dataset))) self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(iterator))) self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) with self.cached_session(): for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToSameDevice(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device( "/job:localhost/replica:0/task:0/device:CPU:0")) # NOTE(mrry): This device block creates the "host" dataset and iterator on # /cpu:0, and ensures that the prefetching is across devices. In typical use # this would not be necessary, because the GPU device would not support any # of the dataset-related ops. with ops.device("/cpu:0"): iterator = device_dataset.make_one_shot_iterator() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) next_element = iterator.get_next() self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) with self.cached_session() as sess: for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchSparseTensorsToDevice(self): def make_tensor(i): return sparse_tensor.SparseTensorValue( indices=[[0, 0]], values=(i*[1]), dense_shape=[2, 2]) host_dataset = dataset_ops.Dataset.range(10).map(make_tensor) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) # NOTE(mrry): This device block creates the "host" dataset and iterator on # /cpu:0, and ensures that the prefetching is across devices. In typical use # this would not be necessary, because the GPU device would not support any # of the dataset-related ops. with ops.device("/cpu:0"): iterator = device_dataset.make_one_shot_iterator() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) next_element = iterator.get_next() self.assertEqual(dtypes.int64, next_element.dtype) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config) as sess: for i in range(10): actual = self.evaluate(next_element) self.assertAllEqual([i], actual.values) self.assertAllEqual([[0, 0]], actual.indices) self.assertAllEqual([2, 2], actual.dense_shape) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchToDeviceGpuWithReInit(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device('/gpu:0')) iterator = dataset_ops.make_initializable_iterator(device_dataset) next_element = iterator.get_next() self.assertEqual(device_dataset._variant_tensor.device, '/device:GPU:0') self.assertEqual(next_element.device, '/device:GPU:0') with self.cached_session(config=config_pb2.ConfigProto( allow_soft_placement=False)): self.evaluate(iterator.initializer) for i in range(5): self.assertEqual(i, self.evaluate(next_element)) self.evaluate(iterator.initializer) for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testIteratorOnDeviceGraphModeOneShotIterator(self): self.skipTest("TODO(b/169429285): tf.data.Dataset.make_one_shot_iterator " "does not support GPU placement.") dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0")) iterator = dataset_ops.make_one_shot_iterator(dataset) data = iterator.get_next() optional_data = iterator.get_next_as_optional() with ops.colocate_with(dataset._variant_tensor): dataset_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(dataset_device)) with ops.colocate_with(iterator._iterator_resource): iterator_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(iterator_device)) with ops.colocate_with(data): data_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(data_device)) with ops.colocate_with(optional_data.get_value()): get_value_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(get_value_device)) with ops.colocate_with(optional_data.has_value()): has_value_device = test_ops.device_placement_op() self.assertIn(b"GPU:0", self.evaluate(has_value_device))
def testPrefetchDictToDevice(self): host_dataset = dataset_ops.Dataset.range(10).map(lambda x: {"a": x}) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) # NOTE(mrry): This device block creates the "host" dataset and iterator on # /cpu:0, and ensures that the prefetching is across devices. In typical use # this would not be necessary, because the GPU device would not support any # of the dataset-related ops. with ops.device("/cpu:0"): iterator = device_dataset.make_one_shot_iterator() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) next_element = iterator.get_next() self.assertEqual(dtypes.int64, next_element["a"].dtype) self.assertEqual([], next_element["a"].shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config) as sess: for i in range(10): self.assertEqual({"a": i}, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchSparseTensorsToDevice(self): def make_tensor(i): return sparse_tensor.SparseTensorValue(indices=[[0, 0]], values=(i * [1]), dense_shape=[2, 2]) host_dataset = dataset_ops.Dataset.range(10).map(make_tensor) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = dataset_ops.make_one_shot_iterator(device_dataset) next_element = iterator.get_next() self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(device_dataset))) self.assertTrue( dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(iterator))) self.assertEqual(dtypes.int64, next_element.dtype) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): for i in range(10): actual = self.evaluate(next_element) self.assertAllEqual([i], actual.values) self.assertAllEqual([[0, 0]], actual.indices) self.assertAllEqual([2, 2], actual.dense_shape) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchSparseTensorsToDevice(self): def make_tensor(i): return sparse_tensor.SparseTensorValue( indices=[[0, 0]], values=(i*[1]), dense_shape=[2, 2]) host_dataset = dataset_ops.Dataset.range(10).map(make_tensor) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = dataset_ops.make_one_shot_iterator(device_dataset) next_element = iterator.get_next() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) self.assertEqual(dtypes.int64, next_element.dtype) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): for i in range(10): actual = self.evaluate(next_element) self.assertAllEqual([i], actual.values) self.assertAllEqual([[0, 0]], actual.indices) self.assertAllEqual([2, 2], actual.dense_shape) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToSameDevice(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device( "/job:localhost/replica:0/task:0/device:CPU:0")) # NOTE(mrry): This device block creates the "host" dataset and iterator on # /cpu:0, and ensures that the prefetching is across devices. In typical use # this would not be necessary, because the GPU device would not support any # of the dataset-related ops. with ops.device("/cpu:0"): iterator = device_dataset.make_one_shot_iterator() self.assertTrue(dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(device_dataset))) self.assertTrue(dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(iterator))) next_element = iterator.get_next() self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): for i in range(10): self.assertEqual(i, sess.run(next_element)) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchDictToDevice(self): host_dataset = dataset_ops.Dataset.range(10).map(lambda x: {"a": x}) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) # NOTE(mrry): This device block creates the "host" dataset and iterator on # /cpu:0, and ensures that the prefetching is across devices. In typical use # this would not be necessary, because the GPU device would not support any # of the dataset-related ops. with ops.device("/cpu:0"): iterator = device_dataset.make_one_shot_iterator() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) next_element = iterator.get_next() self.assertEqual(dtypes.int64, next_element["a"].dtype) self.assertEqual([], next_element["a"].shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config) as sess: for i in range(10): self.assertEqual({"a": i}, sess.run(next_element)) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchToDevice(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = device_dataset.make_one_shot_iterator() next_element = iterator.get_next() self.assertEqual(host_dataset.output_types, device_dataset.output_types) self.assertEqual(host_dataset.output_types, iterator.output_types) self.assertEqual(host_dataset.output_shapes, device_dataset.output_shapes) self.assertEqual(host_dataset.output_shapes, iterator.output_shapes) self.assertEqual(host_dataset.output_classes, device_dataset.output_classes) self.assertEqual(host_dataset.output_classes, iterator.output_classes) self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToDeviceGpu(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) self.assertDatasetProduces(device_dataset, list(range(10)))
def testPrefetchToDeviceCorrectPlacement(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0")) self.assertIn("gpu:0", dataset._variant_tensor.device.lower())
def testPrefetchToDevicePlacement(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) self.assertEqual(device_dataset._variant_tensor.device, "/job:localhost/replica:0/task:0/device:GPU:0")
def testTensorsExplicitPrefetchToDevice(self): ds = Dataset.from_tensor_slices([0., 1.]) ds = ds.apply(prefetching_ops.prefetch_to_device(test.gpu_device_name())) with self.assertRaisesRegexp(TypeError, 'prefetch_to_device'): datasets.Iterator(ds) for i, x in enumerate(ds): with ops.device(test.gpu_device_name()): x = math_ops.add(x, x) self.assertEqual(float(i) + float(i), x.numpy())
def testIteratorOnDeviceEagerMode(self): dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0")) iterator = iter(dataset) data = next(iterator) optional_data = iterator.get_next_as_optional() self.assertIn("gpu:0", dataset._variant_tensor.device.lower()) self.assertIn("gpu:0", iterator._iterator_resource.device.lower()) self.assertIn("gpu:0", data.device.lower()) self.assertIn("gpu:0", optional_data.get_value().device.lower()) self.assertIn("gpu:0", optional_data.has_value().device.lower())
def testIteratorOnDeviceEagerMode(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0")) iterator = iter(dataset) data = next(iterator) self.assertIn("gpu:0", dataset._variant_tensor.device.lower()) self.assertIn("gpu:0", iterator._iterator_resource.device.lower()) self.assertIn("gpu:0", data.device.lower())
def testIteratorOnDeviceGraphModeInitializableIterator(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0")) iterator = dataset_ops.make_initializable_iterator(dataset) data = iterator.get_next() self.assertIn("gpu:0", dataset._variant_tensor.device.lower()) self.assertIn("gpu:0", iterator._iterator_resource.device.lower()) self.assertIn("gpu:0", data.device.lower())
def testTensorsExplicitPrefetchToDevice(self): ds = Dataset.from_tensor_slices([0., 1.]) ds = ds.apply( prefetching_ops.prefetch_to_device(test.gpu_device_name())) with self.assertRaisesRegexp(TypeError, 'prefetch_to_device'): datasets.Iterator(ds) for i, x in enumerate(ds): with ops.device(test.gpu_device_name()): x = math_ops.add(x, x) self.assertEqual(float(i) + float(i), x.numpy())
def testIteratorOnDeviceGraphModeOneShotIterator(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0")) iterator = dataset_ops.make_one_shot_iterator(dataset) data = iterator.get_next() optional_data = iterator.get_next_as_optional() self.assertIn("gpu:0", dataset._variant_tensor.device.lower()) self.assertIn("gpu:0", iterator._iterator_resource.device.lower()) self.assertIn("gpu:0", data.device.lower()) self.assertIn("gpu:0", optional_data.get_value().device.lower()) self.assertIn("gpu:0", optional_data.has_value().device.lower())
def testPrefetchToDeviceGpu(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) iterator = device_dataset.make_one_shot_iterator() next_element = iterator.get_next() with self.cached_session(): for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def prefetch_to_device(device, buffer_size=None): """A transformation that prefetches dataset values to the given `device`. NOTE: Although the transformation creates a `tf.data.Dataset`, the transformation must be the final `Dataset` in the input pipeline. Args: device: A string. The name of a device to which elements will be prefetched. buffer_size: (Optional.) The number of elements to buffer on `device`. Defaults to an automatically chosen value. Returns: A `Dataset` transformation function, which can be passed to `tf.data.Dataset.apply`. """ return prefetching_ops.prefetch_to_device(device, buffer_size)
def testPrefetchToDeviceGpu(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) iterator = device_dataset.make_one_shot_iterator() next_element = iterator.get_next() with self.cached_session() as sess: for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def prefetch_to_device(device, buffer_size=None): """A transformation that prefetches dataset values to the given `device`. NOTE: Although the transformation creates a `tf.data.Dataset`, the transformation must be the final `Dataset` in the input pipeline. Args: device: A string. The name of a device to which elements will be prefetched. buffer_size: (Optional.) The number of elements to buffer on `device`. Defaults to an automatically chosen value. Returns: A `Dataset` transformation function, which can be passed to `tf.data.Dataset.apply`. """ return prefetching_ops.prefetch_to_device(device, buffer_size)
def testIteratorOnDeviceEagerMode(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) host_iterator = iter(host_dataset) device_iterator = iter(device_dataset) host_tensor = next(host_iterator) device_tensor = next(device_iterator) self.assert_dataset_placement(host_dataset, host_iterator, host_tensor, device_dataset, device_iterator, device_tensor)
def testIteratorOnDeviceGraphModeInitializableIterator(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) host_iterator = dataset_ops.make_initializable_iterator(host_dataset) device_iterator = dataset_ops.make_initializable_iterator(device_dataset) host_tensor = host_iterator.get_next() device_tensor = device_iterator.get_next() self.assert_dataset_placement(host_dataset, host_iterator, host_tensor, device_dataset, device_iterator, device_tensor)
def testPrefetchToDeviceGpu(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) iterator = dataset_ops.make_initializable_iterator(device_dataset) next_element = iterator.get_next() with self.cached_session( config=config_pb2.ConfigProto(allow_soft_placement=False)): self.evaluate(iterator.initializer) for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToDeviceGpuWithReInit(self): if not test_util.is_gpu_available(): self.skipTest("No GPU available") host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/gpu:0")) iterator = device_dataset.make_initializable_iterator() next_element = iterator.get_next() with self.cached_session() as sess: sess.run(iterator.initializer) for i in range(5): self.assertEqual(i, sess.run(next_element)) sess.run(iterator.initializer) for i in range(10): self.assertEqual(i, sess.run(next_element)) with self.assertRaises(errors.OutOfRangeError): sess.run(next_element)
def testPrefetchToDevice(self): host_dataset = dataset_ops.Dataset.range(10) device_dataset = host_dataset.apply( prefetching_ops.prefetch_to_device("/cpu:1")) with ops.device("/cpu:1"): iterator = dataset_ops.make_one_shot_iterator(device_dataset) next_element = iterator.get_next() self.assertTrue(dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(device_dataset))) self.assertTrue(dataset_ops.get_structure(host_dataset).is_compatible_with( dataset_ops.get_structure(iterator))) self.assertEqual(dtypes.int64, next_element.dtype) self.assertEqual([], next_element.shape) worker_config = config_pb2.ConfigProto(device_count={"CPU": 2}) with self.test_session(config=worker_config): for i in range(10): self.assertEqual(i, self.evaluate(next_element)) with self.assertRaises(errors.OutOfRangeError): self.evaluate(next_element)
def testPrefetchToDevice(self): if not context.num_gpus(): self.skipTest("No GPU available") dataset = dataset_ops.Dataset.range(10) dataset = dataset.apply(prefetching_ops.prefetch_to_device("/gpu:0"))