Interleaving tf.data.Datasets

MattScarpino is on the right track in his comment. You can use Dataset.zip() along with Dataset.flat_map() to flatten a multi-element dataset:

ds0 = tf.data.Dataset.range(0, 10, 2)
ds1 = tf.data.Dataset.range(1, 10, 2)

# Zip combines an element from each input into a single element, and flat_map
# enables you to map the combined element into two elements, then flattens the
# result.
dataset = tf.data.Dataset.zip((ds0, ds1)).flat_map(
    lambda x0, x1: tf.data.Dataset.from_tensors(x0).concatenate(
        tf.data.Dataset.from_tensors(x1)))

iter = dataset.make_one_shot_iterator()
val = iter.get_next()

Having said this, your intuition about using Dataset.interleave() is pretty sensible. We're investigating ways that you can do this more easily.

PS. As an alternative, you can use Dataset.interleave() to solve the problem if you change how ds0 and ds1 are defined:

dataset = tf.data.Dataset.range(2).interleave(
    lambda x: tf.data.Dataset.range(x, 10, 2), cycle_length=2, block_length=1)

tf.data.experimental.sample_from_datasets method could be also useful if you do not need to preserve the strict order for the items you want to interleave.

In my case I had to interleave a real life data with some synthetic data, so the order was not an issue for me. Then this can be easily done as follows

dataset = tf.data.experimental.sample_from_datasets([ds0, ds1])

Note that the result will be non-deterministic and some items could be taken from same dataset twice, but in general it will be very similar to regular interleave.

The advantages of this approach:

• you can use multiple datasets in one method call
• you can specify fraction of the samples for each dataset using weights parameter (e.g. I wanted to have only small fraction of the data to be generated so I used weights=[0.9, 0.1])

Pavel's answer works great if you don't mind the order of interleaving. If you do care...

Option 1

A variant on mrry's answer that works with an arbitrary number of input datasets:

ds0 = tf.data.Dataset.range(0, 10, 3)
ds1 = tf.data.Dataset.range(1, 10, 3)
ds2 = tf.data.Dataset.range(2, 10, 3)
datasets = (ds0, ds1, ds2)

# Note: `datasets` should be a *tuple*, not a list.
datasets_zipped = tf.data.Dataset.zip(datasets)
# Each element of the dataset will now be a tuple, e.g. (0, 1, 2).
datasets_zipped_tensor = datasets_zipped.map(lambda *args: tf.stack(args))
# Each element will now be a Tensor, e.g. Tensor([0, 1, 2]).
datasets_interleaved = datasets_zipped_tensor.unbatch()

However, note that because of the way zip works, the dataset this produces limited by the length of the shortest input dataset. For example, using the above code with

datasets = [
    tf.data.Dataset.from_tensor_slices([1, 2, 3, 4, 5]),
    tf.data.Dataset.from_tensor_slices([10, 20]),
)

would yield a dataset comprising just [1, 10, 2, 20].

Option 2

Dataset.interleave doesn't suffer from this problem. In some cases you can use interleave with:

# Note: `datasets` should be a *list*, not a tuple
tf.data.Dataset.from_tensor_slices(datasets).interleave(lambda x: x)

But this doesn't seem to work for all kinds of dataset; calling from_tensor_slices on your datasets may not work.

Option 3

If option 2 doesn't work, you might be able to use interleave at an earlier stage of your dataset pipeline. For example, you might be able to change from calling interleave on pre-existing datasets to calling interleave on the file names from which the individual datasets were created:

filenames = ['foo', 'bar']
filesnames_dataset = tf.data.Dataset.from_tensor_slices(filenames)

def read_dataset(filename): ...

interleaved_dataset = filenames_dataset.interleave(read_dataset)

But this will only work if your read_dataset function accepts a Tensor argument.

Option 4

If none of the other options work for you, I think the only solution is to implement the interleaving yourself, with something like:

element_spec = datasets[0].element_spec
assert all(dataset.element_spec == element_spec for dataset in datasets)

def interleave_generator():
  iters_not_exhausted = [iter(dataset) for dataset in datasets]
  while iters_not_exhausted:
    for dataset_iter in iters_not_exhausted:
      try:
        x = next(dataset_iter)
      except StopIteration:
        iters_not_exhausted.remove(dataset_iter)
      else:
        yield x

datasets_interleaved = tf.data.Dataset.from_generator(
    interleave_generator,
    output_signature=element_spec,
)