本篇主要介绍怎么使用 tf.data API 来构建高性能的输入 pipeline。
tf.data官方教程详见前面的博客<<<<<<<<<<tf.data官方教程
目录
模型单个训练 step 时间的减少 依赖于 GPU、TPU 的使用。最优性能不仅依赖于高速的计算硬件，也要求有一个高效的输入管道（Input Pipeline Performance Guide），这个管道在当前step完成前，进行下一个 step 需要的数据的准备。tf.data API 有助于去构建灵活且高效的输入管道。这个文档解释了 tf.data API 的特征和怎么去构建高性能的 TensorFlow 输入管道 over 各种模型 及 硬件加速器。
这个指南主要包含以下几部分：
- 说明 TensorFlow 输入管道本质上是一个 ETL 过程（Extract，Transform，Load）。
- Describes common performance optimizations in the context of the tf.data API.
- Discusses the performance implications of the order in which you apply transformations.
- Summarizes the best practices for designing performant TensorFlow input pipelines.
1. 输入管道结构
一个典型的 TensorFlow 训练输入管道能够抽象为一个 ETL 过程（Extract，Transform，Load）：
- Extract：从永久存储上读取数据——可以是本地（HDD 或 SSD），也可以是网盘（GCS 或 HDFS）
- Transform：使用 CPU 去解析、预处理数据——比如：图像解码、数据增强、变换（比如：随机裁剪、翻转、颜色变换）、打乱、batching。
- Load：将 Transform 后的数据加载到 计算设备（accelerator device(s)）——例如：GPU、TPU 等执行机器学习模型的设备。
同时保留加速器用于重启训练你的模型
这个模式在利用了 GPU 强大算力的同时，有效地利用了 CPU。另外，将输入管道看作一个 ETL 过程，十分有利于性能的优化。
当使用 tf.estimator.Estimator API 时，传给 tf.estimator.Estimator 的 input_fn 包括了前两个阶段（Extract 和 Transform）。在代码中，这可能看起来像下面（简易，顺序）的实现：
def parse_fn(example):
"Parse TFExample records and perform simple data augmentation."
example_fmt = {
"image": tf.FixedLengthFeature((), tf.string, ""),
"label": tf.FixedLengthFeature((), tf.int64, -1)
}
parsed = tf.parse_single_example(example, example_fmt)
image = tf.image.decode_image(parsed["image"])
image = _augment_helper(image) # augments image using slice, reshape, resize_bilinear
return image, parsed["label"]
def input_fn():
files = tf.data.Dataset.list_files("/path/to/dataset/train-*.tfrecord")
dataset = files.interleave(tf.data.TFRecordDataset)
dataset = dataset.shuffle(buffer_size=FLAGS.shuffle_buffer_size)
dataset = dataset.map(map_func=parse_fn)
dataset = dataset.batch(batch_size=FLAGS.batch_size)
return dataset下一部分建立在这个输入管道上。
2. 优化输入管道的性能
因为新计算设备使得训练网络的速度快了很多，所以 CPU 上的预处理很可能会变成瓶颈。tf.data API 提供了很多构建块来设计能够有效利用 CPU 的输入管道，通过优化 ETL 过程的各个步骤。
2.1 Pipelining——prefetch(n)
为了运行一个训练 step，你必须首先 Extract、Transform 训练数据，然后将它 feed 到计算设备上去。但是，在一个简易的同步实现中，当 CPU 在准备数据时，计算设备处于闲置状态。相反，当计算设备在训练模型时，CPU 处于闲置状态。因此，训练 step 的时间是 CPU 的预处理时间 和 计算设备的训练时间的总和。
Pipelining 将一个训练 step 中的 预处理 和 模型执行 重叠起来。当计算设备在执行第 N 个训练 step 时，CPU 为第 N+1 个训练 step 准备数据。通过这个重叠将 step 的时间由原来的 总和 变为了 两个部分（执行训练、数据准备）的最大值。
没有 pipelining，CPU 和 GPU / TPU 很大一部分时间都是闲置的：
使用 pipelining 后，空闲时间显著减少：
tf.data API 通过 tf.data.Dataset.prefetch 变换提供了一个 software pipelining 机制，这个机制解耦了 数据产生的时间 和 数据消耗的时间。尤其是，这个机制使用一个后台线程和一个内部缓存区，在数据被请求前，去从数据数据集中预加载一些数据。因此，为了实现上述的 pipelining 效果，你可以添加 prefetch(1) 作为你数据集管道的最终变换（或者 prefetch(n)，如果你一个训练 step 消耗 n 个元素）。
为了应用这个特性，将：
dataset = dataset.batch(batch_size=FLAGS.batch_size)
return dataset改为：
dataset = dataset.batch(batch_size=FLAGS.batch_size)
dataset = dataset.prefetch(buffer_size=FLAGS.prefetch_buffer_size)
return dataset注意：如果 数据产生器 和 数据消耗器 的工作可能重合，那么 prefetch 变换将在任何时间都能产生性能提升。前面的建议仅仅是最简单的应用。
2.2 并行数据变换——多线程进行 map 变换，map和batch的融合
当准备一个 batch 时，输入元素可能需要去进行预处理。为了这个目的，tf.data API 提供了 tf.data.Dataset.map 变换，这个变换应用一个用户自定义函数到输入数据集的每一个元素。因为输入元素之间是独立的，所以能够在多个 CPU 核心上并行地进行预处理。为了使这成为可能，map 变换提供了一个 num_parallel_calls 参数去指定并行的级别。例如，下面的框图说明了 num_parallel_calls=2 时，map 变换的效果：
num_parallel_calls 参数的最优值取决于你的硬件，训练数据的特点（比如：它的 size、shape），map 函数的计算量 和 CPU 上同时进行的其它处理。一个简单的原则是：将 num_parallel_calls 设置为 CPU 的核心数。例如，如果 CPU 有四个核，将 num_parallel_calls 设置为 4 将会很高效。另一方面，设置 num_parallel_calls 大于 CPU 的核心数，能够导致低效的调度，导致输入管道速度下降。
为了应用这个特性，将：
dataset = dataset.map(map_func=parse_fn)改为：
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)进一步，如果你的 batch size 为成百上千，你的输入管道将很可能受益于并行地 batch 处理。为了这个目的，tf.data API 提供了 tf.contrib.data.map_and_batch 变化，它有效地融合了 map 和 batch 变化。
为了使用这个变换，将：
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)
dataset = dataset.batch(batch_size=FLAGS.batch_size)改为：
dataset = dataset.apply(tf.contrib.data.map_and_batch(
map_func=parse_fn, batch_size=FLAGS.batch_size))2.3 并行数据提取——并行地读取并解析多个数据文件
在一个实际配置中，输入数据可能被存储在网盘（例如，GCS 或 HDFS）。要么因为输入数据不适合本地，要么因为训练是分布式的，在每台机器上复制输入数据是没有意义的。在本地能够很好的读取数据的数据集管道 可能会卡在 I/O 瓶颈上，因为 本地 和 远程存储 有以下区别：
- Time-to-first-byte（读取第一个bytes的时间）：从远程存储读取文件的第一个字节的时间比本地存储长一个数量级。
- Read throughput（读取吞吐量）：虽然远程存储通常提供大的聚合带宽，但是读取单个文件可能仅能利用该带宽的一小部分。
另外，一旦原始字节被读取到内存中，也可能需要对数据进行反序列化或解密（例如：protobuf），这将导致额外的负载。不管数据是本地存储还是远程存储，该开销都存在，但如果数据未被高效地预加载，则远程情况下可能更糟。
为了减轻各种数据提取开销的影响，tf.data API提供了 tf.contrib.data.parallel_interleave 转换。使用此变换可以将 从多个文件中提取数据并解析 这一过程并行化。同时读取的文件的数目可以通过参数 cycle_length 来指定。
下面的框图说明了将 parallel_interleave 变化 cycle_length=2 时的效果：
为了应用这个特性，将：
dataset = files.interleave(tf.data.TFRecordDataset)改为：
dataset = files.apply(tf.contrib.data.parallel_interleave(
tf.data.TFRecordDataset, cycle_length=FLAGS.num_parallel_readers))由于负载或网络事件，远程存储系统的吞吐量会随时间而变化。
为了缓解这种变化，parallel_interleave 变换能够可选地使用 prefetching （详情见：tf.contrib.data.parallel_interleave）
默认情况下，parallel_interleave 变换为元素提供一个确定性顺序，以方便再现。作为 prefetching 的一个替代方案（这在某些情况下，可能不高效），parallel_interleave 变换也提供了一个选项去提高性能（代价是元素的顺序的确定性）。尤其是，如果 sloppy 参数被设置为 True，变换可能偏离设定的顺序，通过临时跳过在下一个元素被请求时元素不可用的文件。
3. 性能考量（Performance Considerations）
tf.data API 是围绕可组合的变换设计的（为用户提供灵活性）。虽然这些变换中的很多变换的次序是可交换的，但某些变换的次序对性能有影响。
3.1 Map and Batch——用户自定义函数向量化
将用户自定义的函数传给 map 变换 会产生调度、执行用户自定义函数的负载。一般情况下，这个负载与自定义函数的计算量相比很小。但是，如果 map 的函数的计算量很小，这个负载将是主要开销。在这种情况下，我们推荐使用向量化的自定义函数（它一次对一个batch进行变换），并且在 map 变换前使用 batch 变换。
3.2 Map and Cache——缓存数据集
tf.data.Dataset.cache 变化能够在内存或本地存储器上缓存一个数据集。如果传递给 map 变换的用户自定义函数的计算量很大，只要得到的数据集仍然适合内存或本地存储，就可以在 map 转换之后应用 cache 转换。
如果用户定义函数导致存储数据集需要的空间超过了 cache 的容量，考虑提前对数据集进行预处理，以减少资源的使用。
3.3 Map and Interleave / Prefetch / Shuffle——变换的顺序
很多变变化（包括 interleave，prefetch，shuffle）维护元素的内部缓存。如果传给 map 变换的 用户自定义函数 改变了元素的 size，那么 map 变换的次序影响内存的使用量。通常情况下，我们建议选择内存使用量更低的次序，除非不同的次序能够产生性能上的提高（例如，为了使用融合的 tf.contrib.data.map_and_batch）。
3.4 Repeat and Shuffle——repeat 和 shuffle 的次序
tf.data.Dataset.repeat 变换重复输入数据有限次（或无限次）；数据的每一次重复称为一个 epoch。tf.data.Dataset.shuffle 变换随机打乱数据集 example 的次序。
如果 repeat 变换被放在 shuffle 变换之前，那么 epoch 边界将变得模糊。也就是说，某些元素可以在其他元素出现一次之前重复。另一方面，如果在 repeat 变换之前应用 shuffle 变换，那么在每个 epoch 开始时，性能可能会下降（因为这时，也需要进行 shuffle 变化的初始化）。换句话说，将 repeat 放置在 shuffle 之前，提供了更好的性能，将 shuffle 放置在 repeat 之前，提供了更强的次序保证。
当可能时，我们推荐使用融合op：tf.contrib.data.shuffle_and_repeat 变换，这个变换在性能和更强的次序保证上都是最好的（good performance and strong ordering guarantees）。否则，我们推荐在 repeat 之前使用 shuffle。
4. 最优实现的总结（Summary of Best Practices）
下面是设计输入管道的最佳实践的总结：
- 使用
prefetch变换去重叠 数据读取器 和 数据消耗器的工作。我们尤其推荐在输入管道的末端添加prefetch(n)（n是batch size），以重叠 CPU 上的变换 及 GPU/TPU设备上的训练。详见【2.1】 - 通过设置
num_parallel_calls参数，来并行 map 变换。我们建议使用将该参数设置为 CPU 的核心数。详见【2.2】 - 如果你使用 batch 变换来将预处理好的元素 batching，我们建议使用融合op：
map_and_batch变换；尤其是你如果使用大的batch size。详见【2.2】 - 如果你的数据存在远程存储上，（且有时需要反序列化），我们建议使用
parallel_interleave来并行数据的读取和解析。详见【2.3】 - 将简单的用户自定义函数进行向量化，然后传递给 map 变换去分摊 用户自定义函数有关的调用、执行的负载。详见【3.1】
- 如果你的数据能够加载到内存，使用
cache变化去在训练的第一个 epoch 将数据集缓存到内存，所以能避免后来的 epoch 读取、解析、变换数据的负载。详见【3.2】 - 如果你的预处理会增加你数据的 size，我们建议你首先使用
interleave、prefetch、shuffle变换去减少内存使用量（如果可能）。详见【3.3】 - 我们建议在
repeat变换之前使用shuffle变换，最好使用融合op：shuffle_and_repeat变换。详见【3.4】
英文版本见：https://tensorflow.google.cn/performance/datasets_performance