文章目录
为什么要使用 TFRecord 文件
正常情况下我们用于训练的文件夹内部往往会存着成千上万的图片或文本等文件，这些文件通常被散列存放。这种存储方式有一些缺点：
- 占用磁盘空间；
- 在一个个读取的时候会非常耗时；
- 占用大量内存空间（有的大型数据不足以一次性加载）。
此时 TFRecord 格式的文件存储形式会很合理的帮我们存储数据。TFRecord 内部使用了 “Protocol Buffer” 二进制数据编码方案，它只占用一个内存块，只需要一次性加载一个二进制文件的方式即可，简单，快速，尤其对大型训练数据很友好。而且当我们的训练数据量比较大的时候，可以将数据分成多个 TFRecord 文件，来提高处理效率。
什么是 TFRecord 文件
TFRecord 是 TensorFlow 中的数据集存储格式。当我们将数据集整理成 TFRecord 格式后，TensorFlow 就可以高效地读取和处理这些数据集，从而帮助我们更高效地进行大规模的模型训练。
TFRecord 可以理解为一系列序列化的 tf.train.Example 元素所组成的列表文件，而每一个 tf.train.Example 又由若干个 tf.train.Feature 的字典组成。形式如下：
[
{ # example 1 (tf.train.Example)
'feature_1': tf.train.Feature,
...
'feature_k': tf.train.Feature
},
...
{ # example N (tf.train.Example)
'feature_1': tf.train.Feature,
...
'feature_k': tf.train.Feature
}
]
怎样写入 TFRecord 文件
步骤：
- 读取该数据元素到内存；
- 建立 tf.train.Feature 的字典；
- 将该元素转换为 tf.train.Example 对象（每一个 tf.train.Example 由若干个 tf.train.Feature 的字典组成）；
- 将该 tf.train.Example 对象序列化为字符串，并通过一个预先定义的 tf.io.TFRecordWriter 写入 TFRecord 文件。
1、导入需要的库
import tensorflow as tf
import os
2、导入图片
data_dir = './faces/' # 图片所在文件夹
tfrecord_file = './tfrecord_files/train.tfrecords' # 要保存的 TFRecord 文件
train_filenames = [data_dir + filename for filename in os.listdir(data_dir)] # 将所有图片的名称写入一个列表中
3、写入 TFRecord 文件
with tf.io.TFRecordWriter(tfrecord_file) as writer:
for filename in train_filenames:
image = open(filename, 'rb').read() # 读取数据集图片到内存，image 为一个 Byte 类型的字符串
feature = { # 建立 tf.train.Feature 字典
'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])) # 图片是一个 Bytes 对象
}
# tf.train.Example 在 tf.train.Features 外面又多了一层封装
example = tf.train.Example(features=tf.train.Features(feature=feature)) # 通过字典建立 Example
writer.write(example.SerializeToString()) # 将 Example 序列化并写入 TFRecord 文件
值得注意的是， tf.train.Feature 支持三种数据格式：
-
tf.train.BytesList ：字符串或原始 Byte 文件（如图片），通过 bytes_list 参数传入一个由字符串数组初始化的 tf.train.BytesList 对象；
-
tf.train.FloatList ：浮点数，通过 float_list 参数传入一个由浮点数数组初始化的 tf.train.FloatList 对象；
-
tf.train.Int64List ：整数，通过 int64_list 参数传入一个由整数数组初始化的 tf.train.Int64List 对象。
对这三种数据格式分别举例说明：
# 将字符串列表转化为utf-8编码
favorite_books = [name.encode('utf-8')
for name in ['machine learning', 'cc150']]
# 生成bytes_list
favorite_books_bytelist = tf.train.BytesList(value = favorite_books)
print(favorite_books_bytelist)
# 生成float_list
hours_floatlist = tf.train.FloatList(value = [15.5, 9.5, 7.0, 8.0])
print(hours_floatlist)
# 生成int64_list
age_int64list = tf.train.Int64List(value=[42]) # 如果只希望保存一个元素而非数组，传入一个只有一个元素的数组即可
print(age_int64list)
value: "machine learning"
value: "cc150"
value: 15.5
value: 9.5
value: 7.0
value: 8.0
value: 42
运行以上代码即可在 tfrecord_file 所指向的文件地址获得一个 train.tfrecords 文件。
怎样读取 TFRecord 文件
-
通过 tf.data.TFRecordDataset 读入原始的 TFRecord 文件（此时文件中的 tf.train.Example 对象尚未被反序列化），获得一个 tf.data.Dataset 数据集对象；
-
通过 Dataset.map 方法，对该数据集对象中的每一个序列化的 tf.train.Example 字符串执行 tf.io.parse_single_example 函数，从而实现反序列化。
我们可以通过以下代码，读取之前建立的 train.tfrecords 文件。
1、初步读取 TFRecord 文件
raw_dataset = tf.data.TFRecordDataset(tfrecord_file)
这里得到的 raw_dataset 仍是字符串类型。
2、生成描述文件
feature_description = {
'image': tf.io.FixedLenFeature([], tf.string)
}
这里的 feature_description 类似于一个数据集的 “描述文件”，通过一个由键值对组成的字典，告知解码器每个 Feature 的类型是什么。
tf.io.FixedLenFeature 的三个输入参数 shape 、 dtype 和 default_value （可省略）为每个 Feature 的形状、类型和默认值。这里我们的数据项都是单个的字符串，所以 shape 为空数组。
3、定义解码器
def _parse_example(example_string): # 将 TFRecord 文件中的每一个序列化的 tf.train.Example 解码
feature_dict = tf.io.parse_single_example(example_string, feature_description)
feature_dict['image'] = tf.io.decode_jpeg(feature_dict['image']) # 解码 JEPG 图片
return feature_dict['image']
dataset = raw_dataset.map(_parse_example)
得到的数据集对象 dataset 是一个可以用于训练的 tf.data.Dataset 对象。
4、展示图片
import matplotlib.pyplot as plt
%matplotlib inline
for image in dataset.take(1):
plt.title('face')
plt.imshow(image.numpy())
plt.show()
