def testLimitedRetracing(self):
if not test_util.is_gpu_available():
self.skipTest("No GPU available")
trace_count = [0]
@def_function.function
def f(iterator):
trace_count[0] += 1
counter = np.int64(0)
for _ in range(5):
elem = next(iterator)
counter += elem[0]
counter += elem[1]
return counter
dataset = dataset_ops.Dataset.range(10)
dataset2 = dataset_ops.Dataset.range(20)
for _ in range(10):
multi_device_iterator = (
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, ["/cpu:0", "/gpu:0"]))
self.assertEqual(self.evaluate(f(multi_device_iterator)), 45)
multi_device_iterator2 = (
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset2, ["/cpu:0", "/gpu:0"]))
self.assertEqual(self.evaluate(f(multi_device_iterator2)), 45)
self.assertEqual(trace_count[0], 1)
def testLimitedRetracing(self):
trace_count = [0]
@def_function.function
def f(iterator):
trace_count[0] += 1
counter = np.int64(0)
for _ in range(5):
elem = next(iterator)
counter += elem[0]
counter += elem[1]
return counter
dataset = dataset_ops.Dataset.range(10)
dataset2 = dataset_ops.Dataset.range(20)
for _ in range(10):
multi_device_iterator = (
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, [self._devices[1], self._devices[2]]))
self.assertEqual(self.evaluate(f(multi_device_iterator)), 45)
multi_device_iterator2 = (
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset2, [self._devices[1], self._devices[2]]))
self.assertEqual(self.evaluate(f(multi_device_iterator2)), 45)
self.assertEqual(trace_count[0], 1)
def testMissingInput(self):
with self.assertRaisesRegex(
ValueError,
"When `dataset` is not provided, both `components` and `element_spec` "
"must be specified."):
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset=None, devices=[self._devices[1], self._devices[2]])
def testGraphMode(self):
dataset = dataset_ops.Dataset.range(1000)
with self.assertRaisesRegex(
RuntimeError,
"OwnedMultiDeviceIterator is only supported inside of tf.function or "
"when eager execution is enabled."):
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, devices=[self._devices[1], self._devices[2]])
def fn():
with ops.device(self._devices[0]):
dataset = dataset_ops.Dataset.range(10)
iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, [self._devices[1], self._devices[2]])
for _ in range(5):
el0, el1 = next(iterator)
queue.enqueue(el0)
queue.enqueue(el1)
def fn():
with ops.device("/cpu:0"):
dataset = dataset_ops.Dataset.range(10)
iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, ["/cpu:0", "/gpu:0"])
for _ in range(5):
el0, el1 = next(iterator)
queue.enqueue(el0)
queue.enqueue(el1)
def testMultipleInitializations(self):
dataset = dataset_ops.Dataset.range(1000)
for _ in range(5):
multi_device_iterator = (
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, [self._devices[1], self._devices[2]]))
for i, el in enumerate(multi_device_iterator):
self.assertEqual([i * 2, i * 2 + 1], [el[0].numpy(), el[1].numpy()])
def testExtraElementSpecInput(self):
dataset = dataset_ops.Dataset.range(1000)
with self.assertRaisesRegex(
ValueError,
"When `dataset` is provided, `element_spec` and `components` must "
"not be specified."):
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, devices=[self._devices[1], self._devices[2]],
element_spec=dataset.element_spec)
def fn():
dataset = dataset_ops._GeneratorDataset(
1,
init_fn,
next_fn,
finalize_fn,
output_signature=tensor_spec.TensorSpec([], dtypes.int64))
iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, [self._devices[1], self._devices[2]])
next(iterator)
def testBasic(self, max_buffer_size, prefetch_buffer_size):
dataset = dataset_ops.Dataset.range(1000)
mdi = multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, [self._devices[1], self._devices[2]],
max_buffer_size=max_buffer_size,
prefetch_buffer_size=prefetch_buffer_size)
for i, el in enumerate(mdi):
self.assertEqual([i * 2, i * 2 + 1], [el[0].numpy(), el[1].numpy()])
def testBasic(self):
if not test_util.is_gpu_available():
self.skipTest("No GPU available")
with ops.device("/cpu:0"):
dataset = dataset_ops.Dataset.range(1000)
mdi = multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, ["/cpu:0", "/gpu:0"])
for i, el in enumerate(mdi):
self.assertEqual([i * 2, i * 2 + 1],
[el[0].numpy(), el[1].numpy()])
def testMultipleInitializations(self):
if not test_util.is_gpu_available():
self.skipTest("No GPU available")
with ops.device("/cpu:0"):
dataset = dataset_ops.Dataset.range(1000)
for _ in range(5):
multi_device_iterator = (
multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, ["/cpu:0", "/gpu:0"]))
for i, el in enumerate(multi_device_iterator):
self.assertEqual([i * 2, i * 2 + 1],
[el[0].numpy(), el[1].numpy()])
def testMissingDevices(self):
dataset = dataset_ops.Dataset.range(1000)
with self.assertRaisesRegex(ValueError, "`devices` must be provided."):
multi_device_iterator_ops.OwnedMultiDeviceIterator(dataset)
def fn():
dataset = dataset_ops._GeneratorDataset(1, init_fn, next_fn,
finalize_fn)
iterator = multi_device_iterator_ops.OwnedMultiDeviceIterator(
dataset, ["/cpu:0", "/gpu:0"])
next(iterator)
