Meaning of buffer_size in Dataset.map , Dataset.prefetch and Dataset.shuffle

TL;DR Despite their similar names, these arguments have quite difference meanings. The buffer_size in Dataset.shuffle() can affect the randomness of your dataset, and hence the order in which elements are produced. The buffer_size in Dataset.prefetch() only affects the time it takes to produce the next element.

The buffer_size argument in tf.data.Dataset.prefetch() and the output_buffer_size argument in tf.contrib.data.Dataset.map() provide a way to tune the performance of your input pipeline: both arguments tell TensorFlow to create a buffer of at most buffer_size elements, and a background thread to fill that buffer in the background. (Note that we removed the output_buffer_size argument from Dataset.map() when it moved from tf.contrib.data to tf.data. New code should use Dataset.prefetch() after map() to get the same behavior.)

Adding a prefetch buffer can improve performance by overlapping the preprocessing of data with downstream computation. Typically it is most useful to add a small prefetch buffer (with perhaps just a single element) at the very end of the pipeline, but more complex pipelines can benefit from additional prefetching, especially when the time to produce a single element can vary.

By contrast, the buffer_size argument to tf.data.Dataset.shuffle() affects the randomness of the transformation. We designed the Dataset.shuffle() transformation (like the tf.train.shuffle_batch() function that it replaces) to handle datasets that are too large to fit in memory. Instead of shuffling the entire dataset, it maintains a buffer of buffer_size elements, and randomly selects the next element from that buffer (replacing it with the next input element, if one is available). Changing the value of buffer_size affects how uniform the shuffling is: if buffer_size is greater than the number of elements in the dataset, you get a uniform shuffle; if it is 1 then you get no shuffling at all. For very large datasets, a typical "good enough" approach is to randomly shard the data into multiple files once before training, then shuffle the filenames uniformly, and then use a smaller shuffle buffer. However, the appropriate choice will depend on the exact nature of your training job.

Code

import tensorflow as tf
def shuffle():
    ds = list(range(0,1000))
    dataset = tf.data.Dataset.from_tensor_slices(ds)
    dataset=dataset.shuffle(buffer_size=500)
    dataset = dataset.batch(batch_size=1)
    iterator = dataset.make_initializable_iterator()
    next_element=iterator.get_next()
    init_op = iterator.initializer
    with tf.Session() as sess:
        sess.run(init_op)
        for i in range(100):
            print(sess.run(next_element), end='')

shuffle()

Output

[298][326][2][351][92][398][72][134][404][378][238][131][369][324][35][182][441][370][372][144][77][11][199][65][346][418][493][343][444][470][222][83][61][81][366][49][295][399][177][507][288][524][401][386][89][371][181][489][172][159][195][232][160][352][495][241][435][127][268][429][382][479][519][116][395][165][233][37][486][553][111][525][170][571][215][530][47][291][558][21][245][514][103][45][545][219][468][338][392][54][139][339][448][471][589][321][223][311][234][314]

I found that @olivier-moindrot is indeed correct, I tried the code provided by @Houtarou Oreki, using the modifications pointed by @max. The code I used was the following:

fake_data = np.concatenate((np.arange(1,500,1),np.zeros(500)))

dataset = tf.data.Dataset.from_tensor_slices(fake_data)
dataset=dataset.shuffle(buffer_size=100)
dataset = dataset.batch(batch_size=10)
iterator = dataset.make_initializable_iterator()
next_element=iterator.get_next()

init_op = iterator.initializer

with tf.Session() as sess:
    sess.run(init_op)
    for i in range(50):
        print(i)
        salida = np.array(sess.run(next_element))
        print(salida)
        print(salida.max())

The code output was indeed a number ranging from 1 to (buffer_size+(i*batch_size)), where i is the number of times you ran next_element. I think the way it is working is the following. First, buffer_size samples are picked in order from the fake_data. Then one by one the batch_size samples are picked from the buffer. Each time a batch sample is picked from the buffer it is replaced by a new one, taken in order from fake_data. I tested this last thing using the following code:

aux = 0
for j in range (10000):
    with tf.Session() as sess:
        sess.run(init_op)
        salida = np.array(sess.run(next_element))
        if salida.max() > aux:
            aux = salida.max()

print(aux)

The maximum value produced by the code was 109. So you need to assure a balanced sample within your batch_size to ensure a uniform sampling during training.

I also tested what @mrry said about performance, I found that the batch_size will prefetch that amount of samples into memory. I tested this using the following code:

dataset = dataset.shuffle(buffer_size=20)
dataset = dataset.prefetch(10)
dataset = dataset.batch(batch_size=5)

Changing the dataset.prefetch(10) amount resulted in no change in memory (RAM) used. This is important when your data does no fit into RAM. I think the best way is to shuffle your data/file_names before feeding them to tf.dataset, and then control the buffer size using buffer_size.

Actually the answer by @olivier-moindrot is not correct.

You can verify it by creating filenames and labels as he/she mention and print the shuffle values.

You will see each shuffle procedure will generate sample randomly with the size equals to buffer size from the dataset.

dataset = dataset.shuffle(buffer_size=1000)
iterator = dataset.make_one_shot_iterator()
next_element = iterator.get_next()
with tf.Session() as sess:
    for i in range(1000):
        print(sess.run(next_element))

The following code snippet demonstrates the effect of buffer_size in ds.shuffle:

t = tf.range(10)
ds = tf.data.Dataset.from_tensor_slices(t)
for batch in ds.shuffle(buffer_size=2, seed=42).batch(5):
  print(batch)

tf.Tensor([1 2 0 3 5], shape=(5,), dtype=int32)
tf.Tensor([4 6 7 8 9], shape=(5,), dtype=int32)

Shuffle is an "action" (for who is familiar with Spark), which reads data of buffer_size into memory and shuffles it in-memory. After that the shuffled data is cut into batches according to the batch size. Note, how 5 has made it into the first batch (and nothing else from the second half of the data).

This brings up all the questions touched in other answers, like do you have enough memory to shuffle the whole dataset in-memory, or you better shuffle file-names, or shuffle data on disk, or both in-memory and on disk.