本篇主要介绍怎么使用 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
- 1
- 2
- 3
- 4
- 5
- 6
- 7
- 8
- 9
- 10
- 11
- 12
- 13
- 14
- 15
- 16
- 17
- 18
下一部分建立在这个输入管道上。
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
- 1
- 2
改为:
dataset = dataset.batch(batch_size=FLAGS.batch_size)
dataset = dataset.prefetch(buffer_size=FLAGS.prefetch_buffer_size)
return dataset
- 1
- 2
- 3
注意:如果 数据产生器 和 数据消耗器 的工作可能重合,那么 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)
- 1
改为:
dataset = dataset.map(map_func=parse_fn, num_parallel_calls=FLAGS.num_parallel_calls)
- 1
进一步,如果你的 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)
- 1
- 2
改为:
dataset = dataset.apply(tf.contrib.data.map_and_batch(
map_func=parse_fn, batch_size=FLAGS.batch_size))
- 1
- 2
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)
- 1
改为:
dataset = files.apply(tf.contrib.data.parallel_interleave(
tf.data.TFRecordDataset, cycle_length=FLAGS.num_parallel_readers))
- 1
- 2
由于负载或网络事件,远程存储系统的吞吐量会随时间而变化。
为了缓解这种变化,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