TensorFlow 2.0 常用模块3：tf.data 流水线加速

在上一篇文章中，我们介绍了 高效的数据流水线模块 tf.data 的基本使用方式。

本篇文章我们将介绍如何通过 prefetch 和 map 的并行化参数，让 tf.data 的性能得到明显提升。

使用 tf.data 的并行化策略提高训练流程效率

当训练模型时，我们希望充分利用计算资源，减少 CPU/GPU 的空载时间。然而有时，数据集的准备处理非常耗时，使得我们在每进行一次训练前都需要花费大量的时间准备待训练的数据，而此时 GPU 只能空载而等待数据，造成了计算资源的浪费，如下图所示：

常规训练流程，在准备数据时，GPU 只能空载

此时， tf.data 的数据集对象为我们提供了 Dataset.prefetch() 方法，使得我们可以让数据集对象 Dataset 在训练时预取出若干个元素，使得在 GPU 训练的同时 CPU 可以准备数据，从而提升训练流程的效率，如下图所示：

使用 Dataset.prefetch() 方法进行数据预加载后的训练流程，在 GPU 进行训练的同时 CPU 进行数据预加载，提高了训练效率

Dataset.prefetch() 的使用方法和前节的 Dataset.batch() 、 Dataset.shuffle() 等非常类似。继续以前节的 MNIST 数据集为例，若希望开启预加载数据，使用如下代码即可：

1mnist_dataset = mnist_dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)

此处参数 buffer_size 既可手工设置，也可设置为 tf.data.experimental.AUTOTUNE 从而由 TensorFlow 自动选择合适的数值。

与此类似，Dataset.map() 也可以利用多 GPU 资源，并行化地对数据项进行变换，从而提高效率。以前节的 MNIST 数据集为例，假设用于训练的计算机具有 2 核的 CPU，我们希望充分利用多核心的优势对数据进行并行化变换（比如 前节 的旋转 90 度函数 rot90 ），可以使用以下代码：

1mnist_dataset = mnist_dataset.map(map_func=rot90, num_parallel_calls=2)

其运行过程如下图所示：

通过设置 Dataset.map() 的 num_parallel_calls 参数实现数据转换的并行化，上部分是未并行化的图示，下部分是 2 核并行的图示

当然，这里同样可以将 num_parallel_calls 设置为 tf.data.experimental.AUTOTUNE 以让 TensorFlow 自动选择合适的数值。

除此以外，还有很多提升数据集处理性能的方式，可参考 TensorFlow 文档 进一步了解。
注：TensorFlow 文档

https://tensorflow.google.cn/guide/data_performance

实例：cats_vs_dogs 图像分类

为了加速 tf.data ，将上一篇文章中的 cats_vs_dogs 图像分类实例进行以下修改：
为训练集的 map 部分增加 num_parallel_calls 参数，即：

1 train_dataset = train_dataset.map(
2        map_func=_decode_and_resize, 
3        num_parallel_calls=tf.data.experimental.AUTOTUNE)

在训练集处理的最后加入：

1 train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

经过修改后的代码如下：

1import tensorflow as tf
2import os
3
4num_epochs = 10
5batch_size = 32
6learning_rate = 0.001
7data_dir = 'C:/datasets/cats_vs_dogs'
8train_cats_dir = data_dir + '/train/cats/'
9train_dogs_dir = data_dir + '/train/dogs/'
10test_cats_dir = data_dir + '/valid/cats/'
11test_dogs_dir = data_dir + '/valid/dogs/'
12
13def _decode_and_resize(filename, label):
14    image_string = tf.io.read_file(filename)            # 读取原始文件
15    image_decoded = tf.image.decode_jpeg(image_string)  # 解码JPEG图片
16    image_resized = tf.image.resize(image_decoded, [256, 256]) / 255.0
17    return image_resized, label
18
19if __name__ == '__main__':
20    # 构建训练数据集
21    train_cat_filenames = tf.constant([train_cats_dir + filename for filename in os.listdir(train_cats_dir)])
22    train_dog_filenames = tf.constant([train_dogs_dir + filename for filename in os.listdir(train_dogs_dir)])
23    train_filenames = tf.concat([train_cat_filenames, train_dog_filenames], axis=-1)
24    train_labels = tf.concat([
25        tf.zeros(train_cat_filenames.shape, dtype=tf.int32), 
26        tf.ones(train_dog_filenames.shape, dtype=tf.int32)], 
27        axis=-1)
28
29    train_dataset = tf.data.Dataset.from_tensor_slices((train_filenames, train_labels))
30    train_dataset = train_dataset.map(
31        map_func=_decode_and_resize, 
32        num_parallel_calls=tf.data.experimental.AUTOTUNE)
33    # 取出前buffer_size个数据放入buffer，并从其中随机采样，采样后的数据用后续数据替换
34    train_dataset = train_dataset.shuffle(buffer_size=23000)    
35    train_dataset = train_dataset.batch(batch_size)
36    train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)
37
38    model = tf.keras.Sequential([
39        tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(256, 256, 3)),
40        tf.keras.layers.MaxPooling2D(),
41        tf.keras.layers.Conv2D(32, 5, activation='relu'),
42        tf.keras.layers.MaxPooling2D(),
43        tf.keras.layers.Flatten(),
44        tf.keras.layers.Dense(64, activation='relu'),
45        tf.keras.layers.Dense(2, activation='softmax')
46    ])
47
48    model.compile(
49        optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
50        loss=tf.keras.losses.sparse_categorical_crossentropy,
51        metrics=[tf.keras.metrics.sparse_categorical_accuracy]
52    )
53
54    model.fit(train_dataset, epochs=num_epochs)

通过对以上示例进行性能测试，我们可以感受到 tf.data 的强大并行化性能。通过 prefetch() 的使用和在 map() 过程中加入 num_parallel_calls 参数，模型训练的时间可缩减至原来的一半甚至更低。测试结果如下：

tf.data 的并行化策略性能测试（纵轴为每 epoch 训练所需时间，单位：秒）

如果帮到了，希望你们能收藏评论加点赞，一键三连走一波，感谢支持。

原文来自微信公众号，侵删。