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

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

根据 TensorFlow documentationtf.contrib.data.Dataset class 的 prefetchmap 方法都有一个名为 buffer_size 的参数。

对于 prefetch 方法,参数称为 buffer_size 并且根据文档:

buffer_size: A tf.int64 scalar tf.Tensor, representing the maximum number elements that will be buffered when prefetching.

对于 map 方法,参数称为 output_buffer_size 并且根据文档:

output_buffer_size: (Optional.) A tf.int64 scalar tf.Tensor, representing the maximum number of processed elements that will be buffered.

shuffle 方法类似,根据文档显示相同的数量:

buffer_size: A tf.int64 scalar tf.Tensor, representing the number of elements from this dataset from which the new dataset will sample.

这些参数之间的关系是什么?

假设我创建一个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 参数起什么作用?

TL;DR 尽管名称相似,但这些参数的含义却大相径庭。 Dataset.shuffle() 中的 buffer_size 会影响数据集的随机性,从而影响元素的生成顺序。 Dataset.prefetch()中的buffer_size只影响产生下一个元素的时间。


tf.data.Dataset.prefetch() and the output_buffer_size argument in tf.contrib.data.Dataset.map() 中的 buffer_size 参数提供了一种调整输入管道的 性能 的方法:两个参数都告诉 TensorFlow 创建一个缓冲区大多数 buffer_size 元素,以及一个后台线程来填充后台缓冲区。 (请注意,当 Dataset.map()tf.contrib.data 移动到 tf.data 时,我们从 Dataset.map() 中删除了 output_buffer_size 参数。新代码应该在 map() 之后使用 Dataset.prefetch() 来得到相同的行为。)

添加预取缓冲区可以通过将数据预处理与下游计算重叠来提高性能。通常,在管道的最后添加一个小的预取缓冲区(可能只有一个元素)是最有用的,但更复杂的管道可以从额外的预取中受益,尤其是当生成单个元素的时间可能不同时。

相比之下,tf.data.Dataset.shuffle() affects the randomness of the transformation. We designed the Dataset.shuffle() transformation (like the tf.train.shuffle_batch() 函数的 buffer_size 参数被它替换)来处理太大而无法放入内存的数据集。它不是打乱整个数据集,而是维护一个 buffer_size 元素的缓冲区,并从该缓冲区中随机选择下一个元素(如果可用,则用下一个输入元素替换它)。更改 buffer_size 的值会影响混洗的均匀程度:如果 buffer_size 大于数据集中的元素数量,则会得到均匀混洗;如果它是 1 那么你根本不会洗牌。对于非常大的数据集,一种典型的 "good enough" 方法是在训练前将数据随机分片一次到多个文件中,然后统一打乱文件名,然后使用较小的打乱缓冲区。但是,适当的选择将取决于您的培训工作的确切性质。


buffer_sizeshuffle()

中的重要性

我想跟进@mrry 之前的回答,以强调 tf.data.Dataset.shuffle()buffer_size 重要性

buffer_sizebuffer_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,或者提前洗牌filenameslabels(显然具有相同的索引) .

由于将所有文件名和标签存储在内存中不是问题,我们实际上可以使用 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...

要点是始终仔细检查洗牌将做什么。捕捉这些错误的一个好方法可能是绘制批次随时间的分布(确保批次包含与训练集大致相同的分布,在我们的示例中一半是猫,一半是非猫)。

实际上@olivier-moindrot 的回答是不正确的。

您可以通过创建文件名和标签来验证它 he/she 提及并打印随机播放值。

您会看到每个随机播放过程都会随机生成大小等于数据集缓冲区大小的样本。

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

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

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 时,这一点很重要。我认为最好的方法是在将 data/file_names 送入 tf.dataset 之前洗牌,然后使用 buffer_size.[=13 控制缓冲区大小=]

代码

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]

以下代码片段演示了 ds.shufflebuffer_size 的效果:

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 是一个“动作”(对于熟悉 Spark 的人来说),它将 buffer_size 的数据读入内存并在内存中随机播放。之后,根据批次大小将混洗后的数据分成批次。请注意,5 是如何进入第一批的(而后半部分数据中没有其他内容)。

这提出了其他答案中涉及的所有问题,例如您是否有足够的内存来在内存中随机播放整个数据集,或者您最好随机播放文件名或磁盘上的数据,或两者都在内存中和磁盘上。