漫谈tf.data.Dataset以及Pipeline结构

基于TensorFlow 2.3.0。

PipeLine能够允许训练计算机的CPU执行IO和计算一个Batch数据的处理过程（activties），如从磁盘加载和转换特性（loading from disk and transforming features），而在此处理过程中，GPU能够操作前一个Batch的数据（in parallel with the GPU operating on the previous batch of data）。
TF2.3.0新版本在管道瓶颈与预处理优化方面的特性介绍。

一 TensorFlow中的PipeLine结构理解

大家无论学什么编程语言，都会接触到Pipeline的概念与结构，但是也不尽相同，TensorFlow中的PipeLine即是其中的一个。

① Pipeline结构的语言体系

1. ETL过程指Extract、Transform、Load三个过程；
2. Buffer代表缓冲区；Prefetch（数据预载入操作）和Shuffle（数据随机打乱操作）有各自的Buffer，而为了区分两者，因而称其各自的Buffer分别为Prefetch_Buffer和Shuffle_Buffer，并称其承载的元素Element数量的多少为：prefetch_buffer_size 和 shuffle_buffer_size；另，由于batch在Transform过程中的位置不固定，会引起各个buffer缓冲区中的元素既可能是一个batch，又可能是一个样本sample，因而，在本文将两者统称为element（元素）。
3. Preprocess预处理过程包含以下内容：IO读取、tf.map等tf.data.DataSet中的转换操作。其中前者是ETL过程中的Extract过程。后者则对应Transform过程。
4. 基本的Transform过程包含shuffle、prefetch、map、cache、batch操作。在本文中以以上五种操作构建的pipeline作为讲解的基础框架。如下图所示。

② PipeLine结构的设计思路介绍。

1. tf.data中的PipeLine结构的设计思路–Batch

如果是上述的基本的模型结构（从数据集到模型的单线程过程）来完成处理数据和模型训练的训练集的feed过程，可能会产生的问题是：
内存RAM不够大，全部数据集无法一次性地读入。
而相应地解决方法是：一次模型训练需要多少数据就将多少的数据读入内存，进行预处理操作以及转换操作，这样做对内存的压力比较小。
而将数据转换成tf.data.DataSet的数据格式即可实现这样想法：在执行的过程中，model.fit方法会使用类似next(iteration)的方法获取dataset中的element，当dataset并没有调用batch时，则调用一次next方法会获取到一个样本即一个element。

# 含有batch时
    features_labels_ds = dataSet.cache(
        "./keshan.tf-data").map(features_n_step, num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(
        batch_size=batch_size).shuffle(1000).prefetch(buffer_size=4)
# 不含有batch时
    features_labels_ds = dataSet.cache(
        "./keshan.tf-data").map(features_n_step, num_parallel_calls=tf.data.experimental.AUTOTUNE).shuffle(1000).prefetch(buffer_size=4)

Note: 注意这里的batch在指定在DataSet对象中后，每次程序从磁盘中读取数据的时候，只会读入指定大小(shuffle_buffer_size or prefetch_buffer_size)的数据量。

当加上batch后，模型从内存中获得的元素element就不再是sample而是包含了很多sample的batch。

2. tf.data中的PipeLine结构的设计思路–Shuffle

shuffle就是常说的打乱数据集中的样本elements的顺序，好处在于能有效地避免过拟合，基本上是模型训练的必选操作。

其中在tf.data中，shuffle的操作流程大致为：（假设是在batch操作后进行shuffle）定义一个shuffle_buffer_size，如果size为10，batch_size为32，则shuffle_Buffer中的的元素数目是10个batch，10*32个sample。进行shuffle操作时，Pipeline从磁盘中先获取到shuffle_buffer_size大小的samples，而后对于对这些样本再进行打乱操作，即打乱batch的次序，但不影响batch内部sample间的次序。如下述代码所示，batch内部的sample依旧保持原有的索引顺序，但batch之间的顺序变换了。

dataset = tf.data.Dataset.range(1, 14)  # ==> [ 1, 2, 3, 4, 5 ]
dataset = dataset.batch(3).shuffle(3) 
print(list(dataset.as_numpy_iterator()))
#######################============================
[array([1, 2, 3], dtype=int64), array([10, 11, 12], dtype=int64), array([4, 5, 6], dtype=int64), array([7, 8, 9], dtype=int64), array([13], dtype=int64)]

dataset = tf.data.Dataset.range(1, 14)  # ==> [ 1, 2, 3, 4, 5 ]
dataset = dataset.batch(3).shuffle(1)
print(list(dataset.as_numpy_iterator()))
#######################============================
[array([1, 2, 3], dtype=int64), array([4, 5, 6], dtype=int64), array([7, 8, 9], dtype=int64), array([10, 11, 12], dtype=int64), array([13], dtype=int64)]

当假设batch在shuffle操作之后，则是对打乱顺序后的sample进行batch分组。但是在batch分组时并不是按照shuffle后的顺序进行分组的，如下述代码所示：

dataset = tf.data.Dataset.range(1, 14) 
dataset = dataset.shuffle(10)
print(list(dataset.as_numpy_iterator()))
dataset = dataset.batch(5)
print(list(dataset.as_numpy_iterator()))
#######################=========================
[7, 5, 6, 1, 2, 9, 12, 10, 13, 4, 3, 11, 8]
[array([3, 4, 6, 8, 5], dtype=int64), array([ 2,  7, 11,  9, 13], dtype=int64), array([10, 12,  1], dtype=int64)]

另外，在从shuffle_Buffer中选择batch进入model，时的element的选择是一个个随机选择的，而每当shuffle_buffer中空缺一个element的内存位置时，会送一个call，提醒pipeline需要从“Prefetch缓冲区中”获取新的样本，处理后存储到缓冲区中。这样的一步操作其实也是下一小节的内容–Prefetch。

上述的内容的一个小例子：
当我们先进行shuffle，size为100，那么首先，100个sample会先进入buffer中，随机选择一个样本进入model，而后buffer中会留下一个位置由磁盘中原始数据集中的第101个元素填充。

3. tf.data中的PipeLine结构的设计思路–Prefetch

上述一节有简单提及，本节着重讲解为什么需要它。

当不存在prefetch时，程序的运行流程是串行的：首先向shuffle中输入100个batch，通过随机选择的方式，一次一次地向模型中送入batch，而当shuffle中地100个batch全部用完后，model就停止了训练，此时shuffle会再次从磁盘中读取shuffle_buffer_size的数据量。直到buffer中填充完毕，model会再次利用GPU资源进行模型的训练。

上述 情景可以正常执行，但是性能方面表现并不好，因为存在model训练空闲的时刻，model训练间歇，意味着GPU等计算资源被闲置，这是在极大的数据集场景中属于极大的浪费。
如果能够在每次从shuffle_Buffer中提取出一个batch时，给pipeline一个反馈信号，使得它能够预先载入一些batch，经过IO操作和预处理操作进入到某个缓冲区，等待shuffle缓冲区中缺少batch时，能够及时输送batch，以不耽误model的训练。而这样的一种行为就是通过prefetch来实现的。并且在单个batch数据集的预处理时间要远高于一个batch下model的训练时间时，这种优势会非常的明显。
下图展示的是当prefetch_buffer为1时和shuffle-buffer_size为1时的情况。

下图展示的是当prefetch_buffer为3时和shuffle-buffer_size为1时的情况。此时会使用CPU资源开启多线程来加速数据的预处理。如此即可实现GPU计算资源的一直使用。

最后，关于prefetch中的buffer_size通常设为tf.data.experimental.AUTOTUNE，他表示会根据计算机的性能选择合适的值。

tf_data会通过prefetch异步预载入下一个batch的数据，以提升模型的性能，使得GPU不需要等待数据，同时也可以选择并行预处理以及加载数据集。

③ PipeLine的数据流程图及硬件结构图介绍。

数据流程图部分的介绍见上节的设计思路，在此着重介绍设备结构图。
分为三个部分：磁盘承载原始数据集、内存和CPU承载IO读取及预处理Transform等过程、GPU和显存进行模型的训练。

假设场景：Transform中仅使用shuffle、prefetch和batch，且GPU仅进行模型的训练。（严格的讲，CPU会参与模型的一些基本计算。）

1. GPU、显存的空间占用来源：用于本次训练的一个Batchsize大小的数据量；模型计算过程中产生的可训练参数以及中间计算变量；模型训练中用于记录过程细节的变量或者对象。
2. CPU用于计算而不用于储存，储存通过内存RAM。
3. RAM的空间占用来源：shuffle和preFetch的各自缓冲区；计算过程中产生的中间变量。

构建高效数据管道的建议

关于tf.data还有很多方法，在此不一一介绍，详情看官网即可。另外，附加一些小的说明，仅供参考，因为具体的使用搭配要结合自身的项目特点进行。

1. 推荐使用prefetch 方法，它能让数据准备和参数迭代两个过程相互并行。
2. 使用 map 时推荐设置num_parallel_calls ，它能让数据转换过程多进行执行。
3. 推荐使用cache，它能够让数据在第一个epoch后缓存到内存中（即全部的训练数据），默认地，在不指定cache的filepath时，会将数据缓存到内存中，但是如果数据集过大，无法一次性为内存所承载，则可指定filepath，存储到disk中，同样能起到提高性能的作用。
   @Link
4. prefetch习惯上置于最后。
5. cache置于所有Transform之后，但在prefetch之前；用于将转换好的数据存储到内存中，第二次epoch时不必再进行相同的转换操作。
6. 习惯上先进行shuffle，再batch。但也可以batch后再shuffle。从shuffle缓冲区的内存占用角度考虑，如果对于以上两个顺序，使用相同的参数值，也就是说，前者的搭配是shuffle(1000).batch(20)；后者为batch(20).shuffle(1000)，则最终前者的shuffle缓冲区占用的内存空间是1000个element的大小；后者占用空间是1000个batch的大小，即1000*20个element。从避免过拟合效果的角度出发，前者更显无序。因为后者的batch内部的elements并未进行shuffle。

二 Tf.data.Dataset的Transform漫谈

tf.data.Dataset等方法可以创建数据集，并且可以在这个上面应用不同的变换操作Transform，但是如何消耗consume这些数据呢？Tensorflow提供了迭代器iterators来处理。

迭代器并不知道在数据集中元素的数目，但它存在一个 get_next方法，用来遍历迭代器上的数值，而一旦数据集被消耗尽，程序会返回一个tf.errors.OutOfRangeError 异常，

① flat_map

它能够将转换函数map_func映射到Dataset数据集对象中的每一个元素，并将嵌套的子元素（同属Dataset对象）压平。
文章中的一些表述对理解pipeline等有启发。

待完善。

参考文章：
Building efficient data pipelines using TensorFlow