TensorFlow2.0 Guide官方教程学习笔记15 -‘Better performance with tf.data‘

最新推荐文章于 2023-04-23 16:52:05 发布

黄水生

最新推荐文章于 2023-04-23 16:52:05 发布

阅读量980

点赞数 1

分类专栏： TensorFlow 2.0 学习笔记文章标签： TensorFlow2.0

本文链接：https://blog.csdn.net/jackhh1/article/details/102763999

版权

TensorFlow 2.0 学习笔记专栏收录该内容

24 篇文章 4 订阅

订阅专栏

本笔记参照TensorFlow官方教程，主要是对‘Better performance with tf.data’教程内容翻译和内容结构编排，原文链接：Better performance with tf.data

一、概览

GPU和TPUs可以从根本上减少执行单个训练步骤所需的时间。要达到最佳性能，需要一个高效的输入管道，在当前步骤完成之前为下一个步骤交付数据。tf.data API有助于建立灵活和有效的输入流水线。本文档解释了tf.data API的特性和最佳实践，用于跨各种模型和加速器构建高性能的TensorFlow输入流水线。
本指南主要阐述以下内容：
- 说明TensorFlow输入管道流水线上是一个ETL过程。（ETL：Extracting-Transforming-Loading）
- 描述设计performant TensorFlow输入流水线的推荐实践。
- 讨论应用转换的顺序对性能的影响。

二、输入流水线架构

一个典型的TensorFlow训练输入流水线框架可以视为是ETL过程：
1.提取（EXtract）：从内存中读取数据（Numpy）或者永久性存储-本地（HDD或SSD）或远程存储（如GCS或HDFS）
2.转换（Transform）：使用CPU对数据进行解析和执行预处理操作，如洗牌、批处理和域特定转换，如图像解压缩和增强，文本向量化或视频时间采用。
3.加载（Load）：将转换过的数据加载到用来执行机器学习模型的加速装置（如GPU或TPU）。
这个模式有效地利用了CPU，同时保留了加速器来进行繁重的模型训练。此外，将输入管道视为ETL流程提供了一个框架，可以简化性能优化的应用。
下面的示例表示输入流水线的一个简单实现，它读取包含标记图像的TFRecord文件，并将它们转换为适合训练的成批图像-标签对。输入流水线表示为‘tf.data.Dataset’，它可以传递给高级的TensorFlow API（如tf.keras）。

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.io.image.decode_image(parsed["image"])
  image = _augment_helper(image)  # augments image using slice, reshape, resize_bilinear
  return image, parsed["label"]

def make_dataset():
  dataset = tf.data.TFRecordDataset("/path/to/dataset/train-*.tfrecord")
  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

下一节将以这个输入流水线为基础，说明设计performant TensorFlow输入流水线的最佳实践。

三、性能优化

新的计算设备（如GPU和TPU）使得训练神经网络的速度越来越快，CPU处理速度成为瓶颈。tf.data API为用户提供了构建模块来设计有效利用CPU的输入流水线，优化ETL流程的每个步骤。

3.1 流水线（Pipelining）

要执行训练步骤，必须首先提取和转换训练数据，然后将其提供给在加速器上运行的模型。然而，在一个简单的同步实现中，当CPU准备数据时，加速器处于空闲状态。相反，当加速器在训练模型时，CPU处于空闲状态。因此，训练步长是CPU预处理时间和加速器训练时间的总和。
流水线操作与单个训练步骤的预处理和模型执行重叠。当加速器执行训练步骤N时，CPU正在为步骤N+1准备数据。这样做可以将步骤时间减少到最大（而不是总和），并减少提取和转换数据所需的时间。
没有流水线操作，CPU和GPU/TPU大部分时间处于空闲：
在这里插入图片描述
执行流水线操作，空闲时间将会大大压缩：

通过‘’tf.data.Dataset.prefetch‘转换，tf.data’ API提供了一个软件流水线操作机制，可以用来解耦数据产生的时间和数据消耗的时间。特别地，转换使用一个后台线程和一个内部缓冲区，以便在请求输入数据集的元素之前预取它们。预取元素的数量应该等于(或者可能大于)单个训练步骤所消耗的批数。您可以手动调整这个值，或者将其设置为tf.data.experimental.AUTOTUNE,它将提示tf.data runtime在运行时动态地调整值。
在上面运行的例子中插入：

dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

当作我们输入流水线的最后一道转换。

	注意，只要有机会将“生产者”的工作与“消费者”的工作重叠，prefetch转换就会带来好处。

3.2 并行数据转换（Parallelize data transformation）

当准备一个批时，输入元素可能需要被预处理。在教程末尾，tf.data API提供‘tf.data.Dataset.map’转换，它可以为输入数据集的每个元素应用用户自定义函数（如运行示例中的‘parse_fn’）。因为输入元素相互独立，所以预处理可以通过多核CPU并行进行。要想实现并行预处理，可以使用‘map’转换提供的‘num_parallel_calls’参数来制定并行处理的级别。例如，下面这张图解释了设置‘num_parallel_calls=2’后的影响：
在这里插入图片描述
‘num_parallel_calls’参数的最佳值取决于我们的硬件、训练数据特性（如大小size和形状shape）、map函数代价和预处理时CPU正在处理的事情。简单的做法是直接去CPU的核数（比如6核，就设置‘num_parallel_calls’），如果这个值设置的远远大于实际CPU核数，将导致调度效率低下，从而导致速度下降。
与‘prefetch’转换类似，‘map’转换也支持‘tf.data.experimental.AUTOTUNE’,将委托tf.data运行时决定使用何种并行度。
要应用刚刚的变化，我们需要将

dataset = dataset.map(map_func=parse_fn)

替换成

dataset = dataset.map(map_func=parse_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)

3.3 并行数据提取（Parallelize data extraction）

在真实设置环境里，输入数据是远程存储的（比如GCS或HDFS），或者输入数据在本地不适合，又或是由于是分布式训练，为每个机器复杂输入数据没有意义。在本地读取数据时工作良好的数据集流水线在远程读取数据时可能成为I/O瓶颈，原因如下:
- Time-to-first-byte:从远程存储读取文件的第一个字节要比从本地存储读取长几个数量级。
- 读取吞吐量:虽然远程存储通常提供大的聚合带宽，但读取单个文件可能只能利用该带宽的一小部分。
另外，此外，将原始字节读入内存后，可能还需要反序列化和/或解密数据(例如protobuf)，这需要额外的计算。无论数据是存储在本地还是远程，都存在这种开销，但是如果没有有效地预取数据，则在远程情况下情况会更糟。
为了减轻各种数据提取开销的影响，可以使用tf.data.Dataset.interleave转换并行化数据提取步骤，交错地处理其他数据集的内容(比如数据文件读取器)。
要重叠的数据集的数量可以由cycle_length参数指定，而并行度的级别可以由num_parallel_calls参数指定。与预取和映射转换类似，interleave转换支持tf.data.experimental.AUTOTUNE，它将委托tf.data运行时决定使用何种并行度。
下面这张图解释设置‘cycle_length=2’和‘num_parallel_calls=2’后对‘interleave’转换的影响：
在这里插入图片描述
要应用上述变化，需将：

dataset = tf.data.TFRecordDataset("/path/to/dataset/train-*.tfrecord")

替换成：

files = tf.data.Dataset.list_files("/path/to/dataset/train-*.tfrecord")
dataset = files.interleave(
    tf.data.TFRecordDataset, cycle_length=FLAGS.num_parallel_reads,
    num_parallel_calls=tf.data.experimental.AUTOTUNE)

四、性能考虑（Performance considerations）

tf.data API是围绕可组合转换设计的，为用户提供了灵活性。虽然许多转换是可交换的，但是某些转换的顺序对性能有影响。

4.1 映射和批（map and batch）

调用传递给映射转换的用户定义函数会产生与调度和执行用户定义函数相关的开销。通常，与函数执行的计算量相比，这种开销很小。但是，如果map做的工作很少，那么这种开销就会控制总成本。在这种情况下，我们建议向量化用户定义函数(即，让它一次操作一批输入)，并在映射转换(map transformation)之前应用批转换(batch transformation)。

4.2 映射和缓存（map and cache）

tf.data.Dataset.cache转换可以缓存数据集，无论是在内存中还是在本地存储中。如果传递给映射转换的用户定义函数开销很大，那么在映射转换之后应用缓存转换，只要结果数据集仍然可以放入内存或本地存储中。如果用户定义的函数增加了存储数据集所需的空间，超出了缓存容量，请考虑在训练作业之前对数据进行预处理，以减少资源使用。

4.3 映射和交织/预取/洗牌（Map and interleave / prefetch / shuffle）

许多转换，包括交织、预取和洗牌，维护元素的内部缓冲区。如果传递给映射转换的用户定义函数改变了元素的大小，那么映射转换的顺序和缓冲元素的转换将影响内存使用。通常，我们建议选择导致更低内存占用的顺序，除非不同的顺序对性能有好处(例如，启用映射和批处理转换的融合)。

4.4 重复和洗牌（repeat and shuffle）

tf.data.Dataset.repeat转换将输入数据重复有限(或无限)次;数据的每次重复通常被称为一次纪元（epoch）。tf.data.Dataset.shuffle随机转换数据集示例的顺序。
如果‘重复’转换在‘洗牌’转换之前，那‘纪元’的边界会变得模糊。也就是说，某些元素可以在其他元素出现之前重复出现一次。另一方面，如果在重复变换之前应用了shuffle变换，那么在与shuffle变换内部状态初始化相关的每个epoch开始时，性能可能会下降。换句话说，前者(repeat before shuffle)提供了更好的性能，而后者(shuffle befor repeat)提供了更强的排序保证。

五、最佳实践总结

下面是设计performant TensorFlow输入流水线的最佳实践总结:

使用预取转换来重叠生产者和消费者的工作。特别是，我们建议将预取添加到输入管道的末端，以使在CPU上执行的转换与在加速器上完成的培训重叠。要么手动调整缓冲区大小，要么使用tf.data.experimental.AUTOTUNE将决策委托给tf.data。
通过设置num_parallel_calls参数来并行化映射转换。要么手动调整并行度，要么使用tf.data.experimental.AUTOTUNE将决策委托给tf.data。
如果我们正在处理远程存储的数据和/或需要反序列化，谷歌建议使用interleave转换并行读取(和反序列化)来自不同文件的数据。
将传入映射转换的廉价用户定义函数向量化，以摊销与调度和执行函数相关的开销。
如果你的数据可以放入内存，那么可以在第一个epoch期间使用缓存转换将其缓存在内存中，以便后续的epoch可以避免与读取、解析和转换相关的开销。
如果预处理增加了数据的大小，我们建议首先应用交错（interleave）、预取（prefetch）和洗牌(shuffle)(如果可能的话)来减少内存使用。
谷歌建议我们在‘repeat’转换前使用‘shuffle’转换。