def datasets_interleave(datasets, block_length=None, cycle_length=None):
datasets = tuple(datasets)
if cycle_length is not None:
return datasets_concatenate([
datasets_interleave(datasets[i:i + cycle_length],
block_length=block_length)
for i in range(0, len(datasets), cycle_length)
])
choices = Dataset.range(len(datasets))
if block_length is not None:
if not is_listing(block_length):
block_length = tuple(block_length for _ in range(len(datasets)))
choices = datasets_concatenate([
Dataset.from_tensors(tf.convert_to_tensor(
i, dtype=tf.int64)).repeat(block_length[i])
for i in range(len(datasets))
])
return tf.data.experimental.choose_from_datasets(datasets,
choices.cache().repeat())
def testTensorsPlacedOnDevice(self):
ds = Dataset.from_tensors([0., 1.])
with ops.device(test.gpu_device_name()):
x = datasets.Iterator(ds).next()
x = math_ops.add(x, x)
self.assertAllEqual([0., 2.], x.numpy())
def testGpuDefinedDataset(self):
with ops.device(test.gpu_device_name()):
ds = Dataset.from_tensors([0., 1.])
for x in ds:
y = math_ops.add(x, x)
self.assertAllEqual([0., 2.], y.numpy())
def testGpuDefinedDataset(self):
with ops.device(test.gpu_device_name()):
ds = Dataset.from_tensors([0., 1.])
for x in ds:
y = math_ops.add(x, x)
self.assertAllEqual([0., 2.], y.numpy())
