tensorflow / tensorflow/datasets
Unexpected memory overhead for `GeneratorBasedBuilder` when `disable_shuffling = True`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 4.6k
- Forks
- 1.6k
- Avg merge
- 3h 54m
- Merged PRs (30d)
- 1
Description
Description of issue
Implementing a GeneratorBasedBuilder subclass for a large dataset leads to a large memory overhead when disable_shuffling = True and the _generate_examples implementation yields a sequence of int as the keys.
The reason for this is that generated examples are distributed across temporary bucket files, and the tfds.core.shuffle.get_bucket_number function mapping keys to bucket file numbers partitions the space of all 128-bit integers evenly across all buckets, resulting in all integers in [0, 2**128 // BUCKETS_NUMBER] being mapped to the first bucket. Then, Shuffler._iter_buckets attempts to read and sort all examples in the dataset (because they are stored in a single bucket).
One way around this is to make sure that keys are spaced so as to fill the space of 128-bit numbers as evenly as possible, for instance by leaving gaps of size int(2 ** 128) // total_num_examples between keys.
When disable_shuffling = False this is not an issue, since the keys are first hashed before being mapped to bucket numbers.
This is not a bug per se, but users should be made aware of this implementation detail, especially since the documentation for GeneratorBasedBuilder encourages users to use image IDs or text file line numbers for keys.
Submit a pull request?
I'm happy to submit a PR adding a note to that effect in the GeneratorBasedBuilder documentation.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Read the GeneratorBasedBuilder documentation in tensorflow_datasets/core/dataset_builder.py around the guidance for image IDs or text file line numbers. Add a note explaining the key-spacing and memory implications when disable_shuffling is True, then verify that the documentation accurately reflects the shuffle behavior described in the issue.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, tensorflow
- Domain
- documentation
- Issue type
- Documentation
- Difficulty
- 1/5
- Estimated time
- Under an hour
- Activity status
- Stale
- Clarity
- Clearly specified
- Newbie friendliness
- 55/100