关于tensorflow：Dataset.map，Dataset.prefetch和Dataset.shuffle中buffer_size的含义


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

根据TensorFlow文档，tf.contrib.data.Dataset类的prefetch和map方法都有一个名为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对象，如下所示：

1
2
3
4
5
6
 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()中的buffer_size参数和tf.contrib.data.Dataset.map()中的output_buffer_size参数提供了一种调整输入管道性能的方法：两个参数都告诉TensorFlow创建最多包含buffer_size元素的缓冲区，以及一个背景线程在后台填充该缓冲区。
(请注意，当output_buffer_size自变量从tf.contrib.data移至tf.data时，已将其从Dataset.map()中移除。新代码应在map()之后使用Dataset.prefetch()来获得相同的行为。)

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

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

相关讨论

shuffle()中buffer_size的重要性

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

较低的buffer_size不仅会在某些情况下使您的混洗效果不佳：还会使您的整个训练变得混乱。

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

例如，假设您正在对图像进行cat分类器训练，并且您的数据是以以下方式组织的(每个类别中都有10000个图像)：

1
2
3
4
5
6
7
8
9
train/
    cat/
        filename_00001.jpg
        filename_00002.jpg
        ...
    not_cat/
        filename_10001.jpg
        filename_10002.jpg
        ...

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

1
2
3
4
5
6
7
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()。

1
2
3
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)。

相关讨论

码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
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]

相关讨论

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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样本。然后从缓冲区中一个一个批处理大小的样本中选取。每次从缓冲区中提取一批样本时，都会用一个新的样本替换它，该样本是从fake_data中按顺序提取的。我使用以下代码测试了这最后一件事：

1
2
3
4
5
6
7
8
9
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会将预采样的数量预取到内存中。我使用以下代码对此进行了测试：

1
2
3
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控制缓冲区大小。

实际上，@ olivier-moindrot的答案不正确。

您可以通过在提及时创建文件名和标签并打印随机值来对其进行验证。

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

1
2
3
4
5
6
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))