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)