TF卡里删掉文件后内存没变大_如何写一个高效的输入流 [TF 笔记 2 ]

前言

笔记 1 里讲到，TF 模型训练的第一步是构造数据流。模型在训练时会不断地消耗数据，计算损失函数，更新模型。模型训练时需要的数据是一个个的 batch ，而原始的数据往往是零散的单个样例。所以我们需要将原始的数据通过一系列的组合操作，拼装为一个个的 batch，送给模型来执行。这一系列的操作，就是数据流。

为什么我们需要数据流呢？为什么不直接把数据预先处理好，然后直接载到到内存里，然后传给模型呢？主要有以下几个原因：

• 当数据量大的时候，我们没有办法将全部数据 fit 到内存中；
• 模型训练的数据不是静态的，我们往往需要加入一些随机操作 （例如乱序，augmentation），来提升模型的泛化能力；
• 数据流帮助我们解耦了数据的预处理和数据执行的过程，能够帮我们更高效的应用硬件资源，例如当分布式训练的时候，一个好的数据流能够帮我们高效的分发数据到不同的硬件上，从而提高整体的训练效率。

一个合理的数据流，能够让你模型训练更加的高效。数据流的本质就是 ETL。一般来说，数据流由三部分组成，分别是初始化，预处理和遍历执行。具体如下：

• 从输入数据初始化数据流 (Extract)
• 添加各种操作处理数据流 (Transform)
• 遍历数据流，执行数据 (Load)

在 Tensorflow 里，我们使用 tf.data 来构建数据流。本文中我们会介绍构造数据流的基本操作，以及如何提高数据流的效率。如果觉得太长不想看，以下总结了本文的主要内容：

• 预先处理数据 （例如解码，resizing 等等）并将其序列化，尽量使用 TFrecords 来准备训练数据；
• 最常见的操作：shuffle -> repeat -> map(parse) -> batch -> prefetch ；
• 有些 map 操作放在 batch 前，有些 map 操作放在 batch 后；
• 在数据流早期进行采样和过滤，以免影响效率；
• 使用 AUTOTUNE 来设置并行执行的数量，不要去手动调节；
• 使用 cache / interleave / prefetch 这些空间换时间的操作时，注意不要消耗太多内存，以免其成为瓶颈。

1. 起点

数据流的起点是数据本身。根据数据的类型以及大小的不同，我们可以通过不同的方式初始化一个数据流。在实际中常见的有以下几种形式：

• 直接从内存中读取，tf.data.Dataset.from_tensor_slices()
• 使用一个 python 生成器 (generator) 初始化
• 从 TFrecords 初始化， tf.data.TFRecordDataset()
• 更复杂的线上应用，可以考虑使用 http://tensorlfow.io 初始化

当数据量不大且能够全部载入内存时，直接从内存中初始化数据流是最简单的方法：

# images: (K, H, W, C) numpy array
# labels: (K,) int list
dataset = tf.data.Dataset.from_tensor_slices((images, labels))

但实际上数据往往过大且不能完全载入内存。这时候我们有两个选择，第一是写一个 python 生成器，以此为出发点初始化数据流，第二是我们可以将数据转换为 TFrecords 形式。

python 生成器

生成器（generator）是 python 中一种节约内存的计算机制。生成器顺序的产生数据，我们并不需要在一开始存储所有的数据，而只有当我们需要用时，才进行计算。我们可以通过 tf.data.Dataset.from_generator 直接从生成器初始化一个数据对象。

deg gen_series(raw_files):
    # whatever python logical you want to write
    yield img, label

ds_series = tf.data.Dataset.from_generator(
    gen_series, 
    output_types=(tf.float32，tf.int32), 
    output_shapes=((), (None,)))

可以看到，这是一种非常灵活的方式，在 gen_series 函数中，我们可以随意的定义各种 pythonic 的逻辑。但是这种方法有两个问题：1. 低效，python 的方式写数据的逻辑很灵活，但是往往很低效， 我们当然可以使用很多工具来优化这个过程 （numba, dask），但是这样又牺牲了我们本来想要的灵活性；2. 可移植性差，这种方式要求生成器和训练代码本身工作在相同的python 进程里，当在更复杂的生产环境，例如多机器训练时，这种方法就不适用了。

TFrecords

一般来说在处理数据时，我们要先读入原始数据，然后解码原始数据，然后进行执行。这里可以看到，真正有用的是解码后的数据（具体来说，是内存里的二进制数据）。如果我们直接存储和读取内存对象，那样必然能够比从原始数据出发进行一系列的操作要更高效。我们常说的序列化 （serialization)，指的就是这种直接从硬盘上保存和读入内存对象的方法。

TFrecords 就是 TF 提供的进行数据序列化的工具。它基于 Proto Buffer 传输机制，每个样本可以看做是一个需要传输的消息。

类比真实的世界，当我们需要有大量的消息需要进行传输，我们可以将消息打包到多个固定大小的包裹里，然后传输包裹，收到包裹后我们先拆包，然后解读一个个的消息。

在 TF 的语境下，我们将每一个训练样本转化为一个 tf.train.Example 消息，然后分散打包到一个个的 TFrecord 文件里。

def _float_feature(value):
  return tf.train.Feature(float_list=tf.train.FloatList(value=value))

def _int64_feature(value):
  return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def serialize_example(example_feature, label_feature):
  # Creates a tf.Example message ready to be written to a file.
  feature = {
      'exmaple': _float_feature(example_feature),
      'label': _int64_feature(label_feature),
  }

  # Create a Features message using tf.train.Example.
  example_proto = tf.train.Example(features=tf.train.Features(feature=feature))
  return example_proto.SerializeToString()

serialized_example = serialize_example([0.666, 0.666], 0)

以上展示了我们如果将一个训练样本转化为一个 Example 消息，可以看到对于一个样本，我们首先生成 Feature，然后打包为一整个 Example 然后将其序列化。

filename = 'test.tfrecord'
writer = tf.data.experimental.TFRecordWriter(filename)
for example,label in example_label_list:
    serialized_feature = serialize_example(example,label)
    writer.write(serialized_feature)

然后我们就可以写入 tfrecord 文件中了。使用 TFrecord 时，有两点我们需要注意：

1. 每个 TFrecord 文件的大小应该适中，过小或者过大都会影响速度，TF 官方建议的是 100M-200M ；

2. 写入 TFrecord 文件时，应该对数据进行越多的预处理越好 (例如 resizing 操作)，避免在数据流中重复的进行预处理操作。在实际操作中，从 TFrecord 中读出的数据，应当直接输入到 augmentation 函数，因为 augmentation 往往是对单个数据进行随机化处理的第一步。

更复杂的线上应用 （可略过）

有些应用场景，例如增量训练等等，我们需要应对数据的不断动态变化。这种情况就需要我们自己来搭建框架来完成了。或者可以考虑使用 tensorflowio 搭配生产环境中的数据流工具完成，例如如下是 tfio + kafka 的一个示例：

import tensorflow_io.kafka as kafka_io

def decode_img(x):
  x = tf.io.decode_raw(x, out_type=tf.uint8)
  x = tf.reshape(x, [512, 512])
  x = tf.image.convert_image_dtype(x, tf.float32)
  return x

def decode_label(y):
  y = tf.io.decode_raw(y, out_type=tf.uint8)
  y = tf.reshape(y, [])
  return y

train_images = kafka_io.KafkaDataset(['img:0'], group='img', eof=True).map(decode_img)
train_labels = kafka_io.KafkaDataset(['label:0'], group='label', eof=True).map(decode_label)
train_kafka = tf.data.Dataset.zip((train_images, train_labels)).batch(K)

2. 基本操作

初始化数据流后，我们就得到了一个 data 对象。我们可以对该对象进行各种的操作，最后得到我们模型训练需要的数据 （一个包含样本和标注的batch）。如下是一个典型的操作，我们会对其中的每一个操作展开进行说明。

dataset = tf.data.TFRecordDataset(files)
dataset = dataset.shuffle(buffer_size=X).repeat(num_epoch)
dataset = dataset.map(lambda record: parse(record), num_parallel_calls=AUTOTUNE)
dataset = dataset.batch(batch_size=B)
dataset = dataset.prefetch(buffere_size=Y)

• 乱序，重复

dataset = dataset.shuffle(buffer_size=1000)
dataset = dataset.repeat(num_epoch)    

打乱和重复数据是常用的两个操作。乱序 (shuffle) 是随机梯度下降 (SGD) 能够随机 （Stochastic）的基础，重复 (repeat) 数据能使我们实现多个 epoch 的训练。

当我们使用 shuffle 方法的时候，dataset 对象会初始化一个固定大小的 buffer 空间。dataset 会同时维持 buffer 的大小：不断向该 buffer 中添加新数据；以及产生数据：随机的从该 buffer 抽取数据。

理论上说，如果我们想要得到最完美的随机化，我们应当把这个 buffer 的大小设置得和数据集大小一样。但是 buffer 空间越大，意味着更大的内存消耗，以及在初始化这个 buffer 时更多的时间消耗。在实际应用中，我们需要设定一个更合理的值。

repeat 方法能够帮助我们重复数据，从而实现多个 epoch 的训练。实际中我们往往看到 repeat 往往是接在 shuffle 后面的。那为何要这么做，而不是反过来，先 repeat 再 shuffle 呢？

原因是为了保证 epoch 的执行顺序。当我们训练模型的时候，在一个 epoch 里我们会遍历所有的数据，每个数据会被计算一次。当 shuffle 在 repeat 之前时，我们只有将一个 epoch 里所有的数据执行完毕之后，才会执行 repeat 进入下一个 epoch；当 shuffle 在 repeat 之后时，epoch 与 epoch 之间的边界就会模糊，就会出现未遍历完数据，已经计算过的数据又出现了情况。

但是当 shuffle 在 repeat 之前的时候，我们会遇到一个效率的问题。每当一个 epoch 结束的时候，我们需要再次初始化 buffer 空间，当 buffer 空间大或者数据预处理时间长的时候，这是一个很耗时的过程。

不过没关系，TF 已经帮你做好了这些优化。在TF 1 里，我们可以使用 tf.data.experimental.shuffle_and_repeat 函数， 该函数会将 shuffle 和 repeat 两个方法融合在一起。TF 2 里我们直接分开写即可，tf.data 会在执行时自动进行优化。

• 映射，打包 ，预支

载入了原始数据之后，我们往往需要对数据进行一些预处理，例如归一化 (normalization) 或者增广 (augmentation)。这时候我们就需要使用 map 函数了。

def parse(record):
    record = augmentation(record)
    record = preprocess(record)
    return record

dataset = dataset.map(lambda record: parse(record))

为了进一步提高效率，我们还能够多线程的调用 map 函数。可以看到这里我并没有有直接设置并行的线程数量，而是采用了自动设置。采用自动设置时，数据流会在训练的早期阶段通过启发式的方式搜索最合适的线程数量，往往最后的结果和手动调的结果没有太大差别。

AUTOTUNE = tf.data.experimental.AUTOTUNE
dataset = dataset.map(lambda record: parse(record)，num_parallel_calls=AUTOTUNE)

处理好数据后，我们就可以打包，创造模型训练需要的 batch 了：

dataset = dataset.batch(batch_size=K)

到此我们的 data pipeline 基本上完成了。当模型 (consumer) 需要执行训练时，数据流 (producer) 会为其产生相应的 batch。但这里依然存在一个效率问题，那就是生产者只在消费者需要时进行生产。我们能不能提前就生产一些，储备一些库存，从而进一步提高效率呢？例如下图所示，我们在 GPU/TPU 执行数据的同时就开始准备数据，不就可以更快么？

图1 预支数据，加快速度 Image Credit: Jiri Simsa https://www.youtube.com/watch?v=ZnukSLKEw34

我们可以使用 prefetch 方法实现这个目的。prefetch 方法会使用一个后台线程以及一个 buffer 来缓存 batch，提前为模型的执行程序准备好数据。一般来说，buffer 的大小应该至少和每一步训练消耗的 batch 数量一致，也就是 GPU/TPU 的数量。我们也可以使用 AUTOTUNE 来设置。

AUTOTUNE = tf.data.experimental.AUTOTUNE
dataset = dataset.map(lambda record: parse(record)，num_parallel_calls=AUTOTUNE)
dataset = dataset.batch(batch_size=K)
dataset = dataset.prefetch(buffer_size=AUTOTUNE)

能不能再快点？

在实际代码中，我们常看到的形式是先 map 再 batch，但这就是最好的方式吗？其实为了进一步的提高效率，我们应该将不同的预处理函数分离出来。有一些预处理函数适合向量化 (vectorization)，我们应当放到 batch 的后面，例如我们要对一个 batch 里所有的 tensor 进行同样的操作；而有一些预测函数，尤其是内部逻辑复杂且涉及到不定长的 tensor 计算的，我们应当放在 batch 前面，例如一些依赖 label 的 augmentation 操作，以及数据过滤/采样等等。

def parse(record):
    ...

def parse_batch(record_batch):
    ...

dataset = dataset.map(lambda record: parse(record))
dataset = dataset.batch(batch_size=K)
dataset = dataset.map(lambda records: parse_batch(records))

静态优化 （static optimization）
TF 会帮助数据流自动进行一些优化，例如我们之前提到的 shuffle+repeat。但有些优化需要我们手动的打开，例如融合 map 和 filter 的连续操作，向量化 map 操作等等。我们可以通过 tf.data.Options 来手动打开一些优化选项，如下所示。具体可参考文档。

options = tf.data.Options()
options.experimental_optimization.map_vectorization.enabled = True
options.experimental_optimization.map_and_filter_fusion = False
dataset = dataset.with_options(options)

• 过滤，采样

在大规模的场景下，模型训练往往会遇到目标不均衡的问题。比如一个二分类任务，可能负例样本的数量远远大于正例样本的数量。解决这个问题我们从目标函数出发，计算损失的时候对不同的类使用不同的权重，也可以对样本进行采样 (sample)，使得训练时使用的数据是均衡的。为了实现采样，我们往往需要使用过滤（filter）方法，对数据中较大的类别 (majority class) 进行下采样 (down sampling)。

除了采样，过滤方法还使用在别的一些场景中，例如 Curriculum Learning 。在模型不同的训练阶段我们需要对不同难度的数据进行学习，我们可以通过 filter 方法来选择数据。

neg_dataset = ori_dataset.filter(lambda features, label: label==0)

如上是一个简单的例子。具体来说，filter 方法接受一个函数作为其参数，该函数将一个 dataset 中一个元素映射为一个 boolean 值。所以以上的代码也可以转换为如下的更通用的形式：

def filter_func(feature, label):
    acceptance = tf.math.equal(label, 0)
    return acceptance

neg_dataset = ori_dataset.filter(filter_func)

数据采样最直接方法就是，对不同的类别构建不同的 dataset 对象，然后使用 tf.data.experimental.sample_from_datasets 方法进行采样。

neg_dataset = ori_dataset.filter(lambda features, label: label==0).repeat()
pos_dataset = ori_dataset.filter(lambda features, label: label==1).repeat()

balanced_dataset = tf.data.experimental.sample_from_datasets(
    [neg_dataset, pos_dataset], [0.5, 0.5])  

但是这种方法存在着一个问题：我们对每个class都要构建一个子数据集， N 个类对应着从原始数据中进行 N 次过滤，当 N 很大时效率会很低。为了提高效率，我们可以使用 tf.data.experimental.rejection_resample 方法从数据集整体的层面上进行采样。

def class_func(features, label):
  return label

sampled_dataset = tf.data.experimental.rejection_resample(
    class_func, target_dist=[0.5, 0.5])

balanced_dataset = sampled_dataset.map(lambda extra_label, 
    features_and_label: features_and_label)

可以看到这里有三个部分：

1. 首先我们定义一个函数 class_func , 该函数将一个 dataset 里的元素映射为类，这里例子 label 就是 class 所以我们直接返回 label ;
2. 然后我们进行采样，我们有两个类，且目标是均衡的分布，所以 target_dist 设置为 [0.5, 0.5] ;
3. 最后，rejection_resample 函数返回的是 (class, example)，所以我们需要丢掉 class 元素，只保留 example 元素。

顺序依然很重要

当我们需要采样或者过滤时候，从效率的角度考虑，我们需要在 pipeline 的越早进行这些操作越好。因为采样和过滤会丢弃掉一些数据，越晚进行这些操作，我们会进行越多的多余操作，从而浪费了时间。

• 缓存，间隔

除了上述的一些操作。还有一些可以进一步提高数据流效率的工具，例如缓存 （cache）， 间隔 （interleave）等等。

缓存操作帮助我们存储一些中间数据，尤其是消耗大量计算时间但是空间却比较少的数据。当我们进行缓存操作时，数据流只在第一个 epoch 进行这些操作，之后会直接读取缓存中的数据，从而提高整体的效率。

dataset.map(time_consuming_mapping).cache().map(memory_consuming_mapping)

但是在大多数情况下这并非完美的操作，因为如果我们如果知道有些操作只需要执行一次，那么在事先我们就应该对数据进行预操作然后序列化处理。当然，在某些无法接触到真实数据的场景下（例如联邦学习），这可能是提高效率的手段。

实际应用中，有时候训练数据不是存储在本地的。还有的时候你有一块慢得出奇的硬盘。如果文件读取速度是你的瓶颈，你可以考虑使用 interleave 。例如在下面的例子中，我们从一个 TFrecord 序列开始，生成一个 dataset，由于我们选择了 AUTOTUNE 的并行方式，TF 会找到最优的并行数量来读取文件。

dataset = tf.data.Dataset.from_tensor_slices(files).interleave(lambda x:
                               tf.data.TFRecordDataset(x),
                               num_parallel_calls=tf.data.experimental.AUTOTUNE))

3. 执行

构造好数据流，接下来就是执行了。TF 1 中常见的方式是构建 input 函数，然后使用 estimator 来调用，TF 2 中我们往往通过内置的方式（model.fit）来调用，或者更直接的自定义训练方式。接下我们来逐个的看一下这三种方法的大致框架。

Estimator

使用 Estimator 来训练模型， 核心在于定义两个函数，input_fn 和 model_fn。如下所示，当我们定义好生成 dataset 对象的函数后，直接传入给 estimator 的 train 方法即可。

def input_fn(tfrecords_list, batch_size, num_epochs):
     
     dataset = tf.data.TFRecordDataset(tfrecords_list, num_parallel_reads=AUTOTUNE)
     dataset = dataset.shuffle(buffer_size=1000)
     dataset = dataset.repeat(num_epoch)   
     dataset = dataset.map(parse_record, num_parallel_calls=AUTOTUNE)
     dataset = dataset.batch(batch_size)
     dataset = dataset.prefetch(buffer_size=batch_size)
 
     return dataset

# model_fn is the model function
clf = tf.estimator.Estimator(model_fn=model_fn, 
                             config=config)

train_input_fn = lambda：input_fn(tfrecords_list, num_gpus, batch_size_per_tower, num_epochs)

clf.train(input_fn=train_input_fn, max_steps=max_steps)

Keras

Keras 的训练方式和 Estimator 类似，但是区别在于 Estimator 中我们定义模型函数，而在 Keras 中我们定义模型对象。我们可以使用同样的方法定义数据流，模型对象在训练（fit）时直接调用数据流函数。

model = build_model()
model.compile(optimizer, loss)
model.fit(train_input_fn())     

自定义训练

自定义的训练方式下，我们可以直接将数据流当做一个生成器来使用。如下所示，我们从数据对象开始构建一个 iterator，然后遍历所有数据。

@tf.function
def train_step():
    ......

dataset = build_dataset()
dataset_iter = iter(dataset)

for epoch in range(num_epochs):    
    for i in range(steps_per_epoch):
        features, labels = dataset_iter.next()
        # train_step

参考文献：

1. inside tf.keras
2. http://tensorflow.org

-------------------------------------------------------------------------------

聊完了输入，下一章我会聊一聊输出。具体的会谈一谈 TF 是如何导出模型 （saved model, checkpoint），TF1 的 checkpoint 和 TF2 的 checkpoint 有什么不同，几种不同的恢复中断训练的方式等等。