tensorflow / tensorflow/datasets

Unexpected memory overhead for `GeneratorBasedBuilder` when `disable_shuffling = True`

Open
#3,491 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

contributions welcome enhancement
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.