TensorFlow之input pipeline性能指南

  GPU和TPU可以从根本上减少执行单个training step所需的时间。为了达到高性能，我们需要一个高效的input pipeline，它可以在当前step完成后为下一个step高效分发数据。tf.data的API可以帮助构建灵活和高效的input pipeline。该文档解释了tf.data的API特性以及最佳实践，来跨多种模型和加速器构建高性能tensorflow input pipelines。

Input Pipeline结构

  一个典型的tensorflow training input pipeline可以看成是一个ETL process问题：

• Extract(E)：从持久化存储中读取数据，可以是本地(比如HDD或SSD)或远程(比如GCS或HDFS)。
• Transform(T)：利用CPU cores来解析和执行在数据上的预处理操作，例如图片解压、数据转换(随机裁减、翻转、颜色扭曲)、重排(shuffling)以及batching。
• Load(L)：在加速设备上(例如GPU或TPU)加载转换后的数据，执行机器学习模型。

这种模式(pattern)可以有效地利用CPU，同时保留加速器让它去处理模型训练部分的重任。另外，将input pipeline看成是一个ETL process，这种结构更有利于性能优化。

  当使用tf.estimator.Estimator的API时，前两个phases(Extract和Transform)可以在input_fn中被捕获，然后传给tf.estimator.Estimator.train。在代码中，看起来实现如下：

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

性能优化

  新计算设备(比如CPU或TPU)可以以一个越来越快的速率来训练神经网络，而CPU处理速度被证明是个瓶颈。tf.data的API提供给用户相应的构建块来设计input pipeline，可以有效利用CPU，优化在ETL过程中的每一步。

Pipeling

  为了执行一个training step，你必须首先extract和transform训练数据，接着将它feed给一个运行在加速器上的模型。然而，在一个原始的同步实现(naive synchronous implementation)中，CPU会准备数据，加速器会保持空闲状态。相反地，当加速器在训练模型时，CPU会保持空闲状态。训练过程的时间会是CPU处理时间和加速器训练时间的总和。
  Pipelining会将一个training step的预处理和模型执行在时序上重叠。当加速器执行了N个training step时，CPU会为第N + 1个step准备数据，这样做是为了减少总时间。如果没有做pipeling，CPU和GPU/TPU的空闲时间会有很多：

  而有了pipeling，空闲时间会急剧减少：

  tf.data的API通过tf.data.Dataset.prefetch转换，提供了一个软件实现的管道机制(software pipeling)，该转换可以被用于将数据生产时间和数据消费时间相解耦。特别的，该转换会使用一个后台线程，以及一个内部缓存(internal buffer)，来prefetch来自input dataset的元素(在它们被请求之前)。这样，为了达到上述的pipeling效果，你可以添加prefetch(1)到你的dataset pipeline中(或者prefetch(n)，如果单个training 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转换在任何时间都会有好处，有机会将一个producer的工作量和一个consumer的工作量在时序上重合。

并行化数据转换

  当准备一个batch时，input elements需要被预处理。为了这个目的，tf.data的API提供了tf.data.Dataset.map转换，它会将一个用户定义的函数(例如运行示例中的parse_fn)应用到input dataset中的每个元素上。由于input elements相互独立，预处理可以跨多个CPU core并行进行。为了达到该目的，map转换提供了num_parallel_calls参数来指定并行度。例如，下图展示了在map转换中设置num_parallel_calls=2的效果：

  选择num_parallel_calls的最佳值取决于你的硬件、训练数据的特性(比如size和shape)、map函数的开销、以及在CPU上同时发生的其它处理过程，一个简单的启发法(heuristic)是使用可提供的CPU cores的数目。例如，如果机器具有4个cores，那么设置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是几百或几千，你的pipeline将可能从batch创建并行化中获得额外收益。为了达到该目的，tf.data的API提供了tf.contrib.data.map_and_batch转换，它可以用效地将map转换和batch转换相混合(fuse)。
  可以将下面代码：

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

数据抽取并行化

  在现实中，input data可以在远端存储(例如GCS或HDFS)，可能因为input data太大在本地存不下，也可能因为训练是分布式的，不希望在每台机器上复制input data。当在本地读取数据时，一个dataset pipeline可以很好地工作，而当远端读取时可能会有I/O瓶颈。这是由于以下在本地存储和远端存储的不同之处：

• 首字节的时间(Time-to-first-byte)：从远端存储上读取一个文件的首字节，可以会比本地存储花费多个数量级。
• 读取吞吐量(Read throughput)：远端存储通常提供了更大的聚集带宽，读取单个文件可能只利用了该带宽的一小部分。

另一方面，一旦原始字节被读到内存中，有必要做并行化和数据压缩(比如protobuf)，这会增加额外开销。这种开销不用管数据是否在本地或远端存储，但在远端存储的情况下，如果数据没有有效地做prefetch，开销会很大。
  为了减轻多种数据抽取开销的影响，tf.data提供了tf.contrib.data.parallel_interleave转换。使用该转换来并行化其它datasets内容(比如data file readers)交错执行。datasets的重合数目可以通过cycle_length参数进行指定。

  为了在运行示例中应用该变化，将：

dataset = files.interleave(tf.data.TFRecordDataset)

变更为：

dataset = files.apply(
    tf.contrib.data.parallel_interleave(
        tf.data.TFRecordDataset, cycle_length=FLAGS.num_parallel_readers
    )
)

远端存储系统的吞吐量可以随时间变化，这是由于负载和网络事件。