Dataset.map，Dataset.prefetch和Dataset.shuffle中buffer_size的含义

tensorflow

tensorflow-datasets

根据TensorFlow 文档， tf.contrib.data.Dataset类的prefetch和map方法都有一个名为buffer_size的参数。

对于prefetch方法，根据文档，该参数称为buffer_size ：

buffer_size：一个tf.int64标量tf.Tensor，表示预取时将要缓冲的最大元素数。

对于map方法，根据文档，该参数称为output_buffer_size ：

output_buffer_size：（可选。）tf.int64标量tf.Tensor，表示将要缓冲的最大已处理元素数。

同样，对于shuffle方法，根据文档显示相同的数量：

buffer_size：一个tf.int64标量tf.Tensor，表示此数据集中要从中采样新数据集的元素数。

这些参数之间有什么关系？

假设我创建一个Dataset对象，如下所示：

 tr_data = TFRecordDataset(trainfilenames)
    tr_data = tr_data.map(providefortraining, output_buffer_size=10 * trainbatchsize, num_parallel_calls\
=5)
    tr_data = tr_data.shuffle(buffer_size= 100 * trainbatchsize)
    tr_data = tr_data.prefetch(buffer_size = 10 * trainbatchsize)
    tr_data = tr_data.batch(trainbatchsize)

上面的代码片段中的buffer参数起着什么作用？

参考资料：

Stack Overflow

2020-05-29 21:14:01 更新

Ryan Eveline• 700

共 4 个回答

高赞时间活跃

TL; DR尽管名称相似，但这些参数具有完全不同的含义。 Dataset.shuffle()的buffer_size可能会影响数据集的随机性，从而影响元素生成的顺序。 Dataset.prefetch()的buffer_size仅影响生成下一个元素所需的时间。

该buffer_size在参数tf.data.Dataset.prefetch()和output_buffer_size在参数tf.contrib.data.Dataset.map()提供了一种方式来调整你的输入管线的性能：这两个参数告诉TensorFlow创造一个缓冲区最多有buffer_size元素，还有一个后台线程在后台填充该缓冲区。（请注意，当output_buffer_size参数从Dataset.map()移至tf.contrib.data时，已将其从Dataset.map()中tf.data 。新代码应在map()之后使用Dataset.prefetch() map()以获得相同的行为。）

通过将数据的预处理与下游计算重叠，添加预取缓冲区可以提高性能。通常，最有用的是在流水线的末尾添加一个小的预取缓冲区（可能只有一个元素），但是更复杂的流水线可以从附加的预取中受益，尤其是在产生单个元素的时间可以变化的情况下。

相比之下， tf.data.Dataset.shuffle()的buffer_size参数会影响转换的随机性 。我们设计了Dataset.shuffle()转换（例如它替换的tf.train.shuffle_batch()函数）来处理太大而无法容纳在内存中的数据集。它没有改组整个数据集，而是维护一个buffer_size元素的缓冲区，并从该缓冲区中随机选择下一个元素（如果可用，将其替换为下一个输入元素）。更改buffer_size的值会影响混洗的均匀性：如果buffer_size大于数据集中的元素数量，则会得到均匀的混洗；如果它是1那么您根本不会洗牌。对于非常大的数据集，典型的“足够好”的方法是在训练之前将数据随机分片到多个文件中，然后均匀地对文件名进行混洗，然后使用较小的混洗缓冲区。但是，适当的选择将取决于培训工作的确切性质。

2020-05-29 21:14:15 更新

笪高达• 676

在`shuffle()`中`buffer_size`重要性

我想跟进@mrry的先前答案，以强调tf.data.Dataset.shuffle()中buffer_size的重要性 。

较低的buffer_size会在某些情况下使您的改组效果 buffer_size还会使您的整个训练变得混乱。

一个实际的例子：猫分类器

例如，假设您正在图像上训练猫分类器，并且您的数据是以以下方式组织的（每个类别中有10000张图像）：

train/
    cat/
        filename_00001.jpg
        filename_00002.jpg
        ...
    not_cat/
        filename_10001.jpg
        filename_10002.jpg
        ...

使用tf.data输入数据的标准方法可以是拥有文件名列表和相应标签列表，并使用tf.data.Dataset.from_tensor_slices()创建数据集：

filenames = ["filename_00001.jpg", "filename_00002.jpg", ..., 
             "filename_10001.jpg", "filename_10002.jpg", ...]
labels = [1, 1, ..., 0, 0...]  # 1 for cat, 0 for not_cat

dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
dataset = dataset.shuffle(buffer_size=1000)  # 1000 should be enough right?
dataset = dataset.map(...)  # transform to images, preprocess, repeat, batch...

上面的代码的主要问题是，实际上不会以正确的方式对数据集进行洗牌。在大约前半段，我们只会看到猫的图像，而在下半段，只会看到非猫的图像。这将极大地伤害训练。
在训练开始时，数据集将采用前1000文件名并将其放入缓冲区，然后在其中随机选择一个。由于前1000张图片都是猫的图片，因此我们只会在开始时选择猫的图片。

此处的解决方法是确保buffer_size大于20000 ，或者预先filenames和labels （显然具有相同的索引）。

由于将所有文件名和标签存储在内存中不是问题，因此我们实际上可以使用buffer_size = len(filenames)来确保将所有内容混在一起。请确保在应用大量转换之前（例如读取图像，对其进行处理，批处理...），先调用tf.data.Dataset.shuffle() ）。

dataset = tf.data.Dataset.from_tensor_slices((filenames, labels))
dataset = dataset.shuffle(buffer_size=len(filenames)) 
dataset = dataset.map(...)  # transform to images, preprocess, repeat, batch...

要注意的是，总是要仔细检查改组将要做什么。捕获这些错误的一种好方法是绘制随时间变化的批次分布（确保批次包含与训练集大致相同的分布，在我们的示例中为一半cat和一半non cat）。

2020-05-29 21:14:14 更新

伯芦• 654

我发现@ olivier-moindrot确实是正确的，我使用@max指向的修改尝试了@Houtarou Oreki提供的代码。我使用的代码如下：

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())

输出的代码确实是从1到（buffer_size +（i * batch_size））的数字，其中i是您运行next_element的次数。我认为其工作方式如下。首先，从fake_data中按顺序选择buffer_size样本。然后从缓冲区中选取一个batch_size样本。每次从缓冲区中选取一批样本时，都会将其替换为一个新的样本，该样本将按顺序从fake_data获取 。我使用以下代码测试了这最后一件事：

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)

该代码产生的最大值为109。因此，您需要确保batch_size内的样本均衡，以确保训练期间进行均匀的采样。

我还测试了@mrry关于性能的内容，发现batch_size会将预采样的数量预取到内存中。我使用以下代码对此进行了测试：

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

更改dataset.prefetch（10）数量不会导致使用的内存（RAM）发生变化。当您的数据不适合RAM时，这一点很重要。我认为最好的方法是在将数据/文件名输入tf.dataset之前先对其进行混洗，然后使用buffer_size控制缓冲区大小。

2020-05-29 21:14:06 更新

Carl Agnes• 705

码

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()

输出量

[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]