I have tried in colab with TF 2.1.0 and Nightly versions. Please, find the gist here. Thanks!

I am working on a "TFIndexedDataset" RFC for externally shuffling dataset whose limit is beyond memory size. Let me know your preferred use case and I will consider adding it to the RFC. Thanks!

@byronyi Basically a way to use either TFDS or TFRecord files where shuffling considers the whole dataset. To be more specific: I expect a dataset to know its size and provide random access to its elements where this is possible (and I assume for most datasets it is possible as they use images, videos or whatever files or lines which can be counted and ordered first, I even thought TFRecord files were made for that). With this 2 properties a random shuffle operation should exist, which produces every element from the dataset exactly once in a completely random order. So basically a dataset.shuffle(dataset.size) equivalent which doesn't preload everything in memory.

Feature request:
Add a return_shuffled argument to the .cache() method on the tf.data.Datasets. Whether the dataset is cached in memory or on disk it should be doable to create a pointer index while caching and sample randomly from the index when reading from the cache. Obviously, you can't do that from a TFRecord file ccurrently and implementing it would be so much work. It also makes sense that data would be stored on a magnetic drive but cached on an SSD.

@Flamefire The recommended approach is to write your dataset out into multiple TFRecord files and then when loading the data doing it like so:

files = tf.data.Dataset.from_tensor_slices(filenames)
files = files.shuffle(len(filenames))
ds = files.interleave(lambda x: tf.data.TFRecordDataset(x, compression_type='GZIP').prefetch(1),
                      num_parallel_calls=tf.data.experimental.AUTOTUNE,
                      deterministic=False, cycle_length=10)

That way you are reading from 10 files at once (still serially in each file though), feeding it to the shuffle buffer at the appropriate point in your pipeline. And each epoch you shuffle the files you are reading from and if you have saved the dataset in more than 10 files that will change the order that you get the files (i.e. if you have 20 files then the expectation is that 5 files will be the same in the first half of epochs 1 and 2).

@jsimsa I really think adding a shuffle capability to the cache of TF Dataset would be a slam-dunk. It would definitely let me remove a lot of code from my projects that only exists to deal with Tensorflows limited capability in shuffling ❤️

@grofte am I correct to assume that you would want the functionality of return_shuffled to be equivalent to cache().shuffle(buffer_size=NUM_ELEMENTS) while avoiding the extra buffer that shuffle introduces? In particular, the expected behavior during the first epoch (where the cache is constructed) would be that the first element is returned only after the entire cache has been constructed.

I don't think the return_shuffled is general enough (e.g. you would also need to specify the RNG seed management to achieve reproducibility and deterministic reshuffling across epochs). What I can imagine though is to introduce tf.data.Dataset.at(i) transformation which could be used to index into "indexable" datasets (such as cache or from_tensor_slices). The way one could use it then would be:

ds = ... # the dataset to be cached
ds = ds.cache()

indices = tf.data.Dataset.range(NUM_ELEMENTS)
indices = indices.shuffle(NUM_ELEMENTS)
ds = indices.flat_map(lambda index: ds.at(index))

This is a more general API, which handles the aspect of shuffling mentioned above and could be used with any datasets that supports indexing. The main disadvantage is that it requires the number of elements of the caches dataset to be known.

You are 100% correct that an .at method would be an excellent extension. Probably very useful for researchers that need their data with some kind of deterministic properties. But for both applications you are going to need to read the whole dataset once without training - analogous to how the Keras preprocessing layers have an .adapt method. Because you need to link the indices to memory/disk adresses either way. Say you didn't and your first three indices were 100, 3, and 2000. Then you would have to read the first 100 elements/rows from the file, then 3, and then 1997. So 2100 reads to get three elements/rows. That's very ineffiencient.

As I understand it caching currently happens during the first training epoch for efficiency reasons. However, if you expect to train your model for many epochs then performing the caching before is only a small percentage increase in time and a large potential increase in model quality (and faster convergence which reduces training time). So you wouldn't have to know the number of elements in advance.

I am working on a "TFIndexedDataset" RFC for externally shuffling dataset whose limit is beyond memory size. Let me know your preferred use case and I will consider adding it to the RFC. Thanks!

The proposed IndexedDataset is stored with the very same on-disk format as that of cached datasets, but it could be stored in the remote storage without reading it through first. The dataset.at(index) API looks perfect, and if people in this thread still interest, I could continue work on this proposal.

@grofte I am a little bit confused by your example. I would assume that if the first three indices were 100, 3, and 2000, then you would need to read 100 elements to get the first results, and then either 0 (if the read elements are cached) or 3, and then either 1900 (if the read elements are cached) or 2000. So you would either need to read 2000 elements if the read elements are cached or 2103 elements if there is no caching -- neither of which matches your description.

Note that at will generally not require reading the whole dataset. For dataset that cannot be efficiently indexed into, it simply won't be supported (i.e. it will throw an informative error). The existing skip transformation could be used to perform this inefficient indexing that you are alluding to.

I was asking you about what would you expect the behavior of cache(..., returned_shuffled=True) be on the first epoch (e.g. when the data is not cached yet). In order for the method to be able to return any permutation of the input dataset, it needs to a) known the cardinality of the input dataset so that it knows the range of indices to consider and b) possibly require reading the entire input dataset if the first element to return after shuffling is the last element of its input.

@byronyi it makes sense for us to chat over VC about your proposal -- we have internal WIP proposal for indexed datasets as well and it would be great if we could align the two and I would be more than happy for you to lead the effort. I will reach out to you via email.