TensorFlow 中层 API TFRecordDataset

本文作者YJango，本文首发于知乎专栏【超智能体】，AI 研习社获其授权转载。

半年没有更新了， 由于抑郁，我把 gitbook 上的《超智能体》电子书删掉了，所有以 gitbook 作为资源所显示的图片以及所有引向 gitbook 的链接全部失效。CSDN 上是不能看了。

遇到图片和链接失效的朋友到我知乎的专栏里找相应的文章，如果没有了只能说声抱歉。

很欣慰还有人喜欢我写的文章，以及对超智能体专栏（http://t.cn/RJvw4iB）的支持。

有很多想说的，却又不知道说什么。只是，谢谢。以往 YJango 的文章都是以教学为主，并不覆盖高效的实际应用。

这篇文章是专门写给那些支持过我的读者们，感谢你们。

完整代码可以从下面的 github（http://t.cn/R8yKP99） 上找到。

目录

前言

优势

Dataset API

TFRecord

概念

数据说明

数据存储

常用存储

TFRecord 存储

实现

生成数据

写入 TFRecord file

存储类型

如何存储张量 feature

使用 Dataset

创建 dataset

操作 dataset

解析函数

迭代样本

Shuffle

Batch

Batch padding

Epoch

帮助函数

前言

半年没有更新了， 由于抑郁，我把 gitbook 上的《超智能体》电子书删掉了，所有以 gitbook 作为资源所显示的图片以及所有引向 gitbook 的链接全部失效。CSDN 上是不能看了。

遇到图片和链接失效的朋友到我知乎的专栏里找相应的文章，如果没有了只能说声抱歉。

很欣慰还有人喜欢我写的文章，以及对超智能体专栏（http://t.cn/RJvw4iB）的支持。

有很多想说的，却又不知道说什么。只是，谢谢。以往 YJango 的文章都是以教学为主，并不覆盖高效的实际应用。

这篇文章是专门写给那些支持过我的读者们，感谢你们。

完整代码可以从下面的 github（http://t.cn/R8yKP99） 上找到。

优势

一、为什么用 Dataset API？

1. 简洁性：

常规方式：用 python 代码来进行 batch，shuffle，padding 等 numpy 类型的数据处理，再用 placeholder + feed_dict 来将其导入到 graph 中变成 tensor 类型。因此在网络的训练过程中，不得不在 tensorflow 的代码中穿插 python 代码来实现控制。

Dataset API：将数据直接放在 graph 中进行处理，整体对数据集进行上述数据操作，使代码更加简洁。

2. 对接性：TensorFlow 中也加入了高级 API （Estimator、Experiment，Dataset）帮助建立网络，和 Keras 等库不一样的是：这些 API 并不注重网络结构的搭建，而是将不同类型的操作分开，帮助周边操作。可以在保证网络结构控制权的基础上，节省工作量。若使用 Dataset API 导入数据，后续还可选择与 Estimator 对接。

二、为什么用 TFRecord？

在数据集较小时，我们会把数据全部加载到内存里方便快速导入，但当数据量超过内存大小时，就只能放在硬盘上来一点点读取，这时就不得不考虑数据的移动、读取、处理等速度。使用 TFRecord 就是为了提速和节约空间的。

概念

在进行代码功能讲解之前，先明确一下想要存储和读取的数据是什么样子（老手跳过）。

一、数据说明：

假设要学习判断个人收入的模型。我们会事先搜集反映个人信息的输入，用这些信息作为判断个人收入的依据。同时也会把拥有的人的实际收入也搜集。这样搜集n个人的后形成我们的数据集。

1. 训练：在每一步训练中，神经网络会把输入和 正确的输出送入中来更新一次神经网络f()中的参数 θ。用很多个不同的不断更新 θ，最终希望当遇到新的时，可以用判断出正确的。

2. 专有名词：结合下图说明名称

样本 (example)：：输入和正确的输出一起叫做样本。给网络展示了什么输入该产生什么样的输出。这里每个是五维向量，每个是一维向量。

表征 (representation)：：集合了代表个人的全部特征。

特征 (feature)：中的某个维：如年龄，职业。是某人的一个特点。

标签 (label)：：正确的输出。

一个样本 (an example)

二、数据存储

为达成上述的训练，我们需要把所有的样本存储成合适的类型以供随后的训练。

1. 常用存储：

输入和标签是分开存储，若有 100 个样本，所有的输入存储成一个 100×5 的 numpy 矩阵；所有的输出则是100×1。

2. TFRecord 存储：

TFRecord 是以字典的方式一次写一个样本，字典的 keys 可以不以输入和标签，而以不同的特征（如学历，年龄，职业，收入）区分，在随后的读取中再选择哪些特征形成输入，哪些形成标签。这样的好处是，后续可以根据需要只挑选特定的特征；也可以方便应对例如多任务学习这样有多个输入和标签的机器学习任务。

注：一般而言，单数的 feature 是一个维度，即标量。所有的 features 组成 representation。但在 TFRecord 的存储中，字典中 feature 的 value 可以不是标量。如：key 为学历的 value 就可以是：[初中，高中，大学]，3 个 features 所形成的向量。亦可是任何维度的张量。

实现

一、生成数据

除了标量和向量外，feature 有时会是矩阵（如段落），有时会还会是三维张量（如图片）。

所以这里展示如何写入三个样本，每个样本有四个 feature，分别是标量，向量，矩阵，三维张量（图片）。

1. 导入库包

2. 生成数据

三个样本的数值是递增的，方便认清顺序

显示结果

二、写入 TFRecord file

1. 打开 TFRecord file

2. 创建样本写入字典

这里准备一个样本一个样本的写入 TFRecord file 中。

先把每个样本中所有 feature 的信息和值存到字典中，key 为 feature 名，value 为 feature 值。

feature 值需要转变成 tensorflow 指定的 feature 类型中的一个：

2.1. 存储类型

int64：

float32：

string：

注：输入必须是 list(向量)

2.2. 如何处理类型是张量的 feature

tensorflow feature 类型只接受 list 数据，但如果数据类型是矩阵或者张量该如何处理？

两种方式：

转成 list 类型：将张量 fatten 成 list(也就是向量)，再用写入 list 的方式写入。

转成 string 类型：将张量用. tostring() 转换成 string 类型，再用 tf.train.Feature(bytes_list=tf.train.BytesList(value=[input.tostring()])) 来存储。

形状信息：不管那种方式都会使数据丢失形状信息，所以在向该样本中写入 feature 时应该额外加入 shape 信息作为额外 feature。shape 信息是 int 类型，这里我是用原 feature 名字 +'_shape'来指定 shape 信息的 feature 名。

3. 转成 tf_features

4. 转成 tf_example

5. 序列化样本

6. 写入样本

7. 关闭 TFRecord file

三、使用 Dataset

1. 创建 dataset

Dataset 是你的数据集，包含了某次将要使用的所有样本，且所有样本的结构需相同（在 tensorflow 官网介绍中，样本 example 也被称作 element）。样本需从 source 导入到 dataset 中，导入的方式有很多中。随后也可从已有的 dataset 中构建出新的 dataset。

1.1. 直接导入（非本文重点，随后不再提）

1.2. 从 TFRecord 文件导入

2. 操作 dataset：

如优势中所提到的，我们希望对 dataset 中的所有样本进行统一的操作（batch，shuffle，padding 等）。接下来就是对 dataset 的操作。

2.1. dataset.map(func)

由于从 tfrecord 文件中导入的样本是刚才写入的 tf_serialized 序列化样本，所以我们需要对每一个样本进行解析。这里就用 dataset.map(parse_function) 来对 dataset 里的每个样本进行相同的解析操作。

注：dataset.map(输入) 中的输入是一个函数。

2.1.1. feature 信息

解析基本就是写入时的逆过程，所以会需要写入时的信息，这里先列出刚才写入时，所有 feature 的各项信息。

注：用到了 pandas，没有的请 pip install pandas。

显示结果

有 6 个信息，name, type, shape, isbyte, length_type, default。前 3 个好懂，这里额外说明后 3 个：

isbyte：是用于记录该 feature 是否字符化了。

default：是当所读的样本中该 feature 值缺失用什么填补，这里并没有使用，所以全部都是 np.NaN

length_type：是指示读取向量的方式是否定长，之后详细说明。

注：这里的信息都是在写入时数据的原始信息。但是为了展示某些特性，这里做了改动：

把 vector 的 shape 从 (3,) 改动成了 (1,3)

把 matrix 的 length_type 改成了 var（不定长）

2.1.2. 创建解析函数

接下就创建 parse function。

Step 1. 创建样本解析字典

该字典存放着所有 feature 的解析方式，key 为 feature 名，value 为 feature 的解析方式。

解析方式有两种：

定长特征解析：tf.FixedLenFeature(shape, dtype, default_value)

shape：可当 reshape 来用，如 vector 的 shape 从 (3,) 改动成了 (1,3)。

注：如果写入的 feature 使用了. tostring() 其 shape 就是 ()

dtype：必须是 tf.float32， tf.int64， tf.string 中的一种。

default_value：feature 值缺失时所指定的值。

不定长特征解析：tf.VarLenFeature(dtype)

注：可以不明确指定 shape，但得到的 tensor 是 SparseTensor。

Step 2. 解析样本

Step 3. 转变特征

得到的 parsed_example 也是一个字典，其中每个 key 是对应 feature 的名字，value 是相应的 feature 解析值。如果使用了下面两种情况，则还需要对这些值进行转变。其他情况则不用。

string 类型：tf.decode_raw(parsed_feature, type) 来解码

注：这里 type 必须要和当初. tostring() 化前的一致。如 tensor 转变前是 tf.uint8，这里就需是 tf.uint8；转变前是 tf.float32，则 tf.float32

VarLen 解析：由于得到的是 SparseTensor，所以视情况需要用 tf.sparse_tensor_to_dense(SparseTensor) 来转变成 DenseTensor

Step 4. 改变形状

到此为止得到的特征都是向量，需要根据之前存储的 shape 信息对每个 feature 进行 reshape。

Step 5. 返回样本

现在样本中的所有 feature 都被正确设定了。可以根据需求将不同的 feature 进行拆分合并等处理，得到想要的输入 和标签 ，最终在 parse_function 末尾返回。这里为了展示，我直接返回存有 4 个特征的字典。

2.1.3. 执行解析函数

创建好解析函数后，将创建的 parse_function 送入 dataset.map() 得到新的数据集

2.2. 创建迭代器

有了解析过的数据集后，接下来就是获取当中的样本。

2.3. 获取样本

显示结果，还会显示先前保存的头像

我们写进 test.tfrecord 文件中了 3 个样本，用

导入了两次，所以有 6 个样本。scalar 的值，也符合所写入的数据。

2.4. Shuffle

可以轻松使用来打乱顺序。设置成一个大于你数据集中样本数量的值来确保其充分打乱。

注：对于数据集特别巨大的情况，请参考 YJango：tensorflow 中读取大规模 tfrecord 如何充分 shuffle？（http://t.cn/R8y0jmQ）

顺序打乱了，但 1,2,3 都出现过 2 次

2.5. Batch

再从乱序后的数据集上进行 batch。

6 个样本，以 4 个进行 batch，第一个得到 4 个，第二个得到余下的 2 个

2.6. Batch_padding

也可以在每个 batch 内进行 padding

padded_shapes 指定了内部数据是如何 pad 的。

rank 数要与元数据对应

rank 中的任何一维被设定成 None 或 - 1 时都表示将 pad 到该 batch 下的最大长度

显示结果

2.7. Epoch

使用来指定要遍历几遍整个数据集

显示结果

帮助函数

如果你不想设定，直接送入数据就可以读写。可以用下面的方法

复制 dataset_helper.py（http://t.cn/R8yOJFQ）

把所有样本写成一个 list

list 中的每个元素都是一个字典

字典的 key 是 feature 名

字典的 value 是 feature 值

feature 值的形状和类型都指定好 (int64, float32)

导入帮助类

数据写成 tfrecord

从 tfrecord 导入数据到 dataset

使用迭代器获取样本

YJango/TFRecord-Dataset-APIgithub.com（http://t.cn/R8yOmmX）

时间允许会结合剩下的两个高层 API：Estimator、Experiment 作为下一篇介绍内容。

参考资料：

TensorFlow Importing Data（http://t.cn/Rl2azT8）

Tfrecords Guide（http://t.cn/R8yWPJo）

NLP 工程师入门实践班：基于深度学习的自然语言处理

三大模块，五大应用，手把手快速入门 NLP

海外博士讲师，丰富项目经验

算法 + 实践，搭配典型行业应用

随到随学，专业社群，讲师在线答疑

AI 科技评论年度特辑