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())