TensorFlow2系列中TFrecord数据集建立与使用

TFrecord数据集形式介绍

TFrecord 是在TensorFlow中的一种数据集存储形式，可以使TensorFlow高效的读取和利用这些数据集。

TFrecord的结构可以有以下理解：

由若干个 tf.train.Feature 特征组成 tf.train.Example （字典形式），由若干个序列化的tf.train.Example字典元素组成的列表文件就是TFrecord数据集文件。

形式如下

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


# dataset.tfrecords
[
    {   # 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数据集

1，准备数据

此次作为测试的数据集是截取的MVSO图像情感数据集中的一小部分

文件结构:

文件结构

2, 建立TFrecord数据集

建立数据集时，我们要对数据文件中的每个元素进行下面的步骤：

读取该数据到内存
将该元素对应的特征（图像数据，类标等）建立Feature字典（由tf.train.Feature元素组成）。并将其转换为多个 tf.train.Example'对象
将该 tf.train.Example 对象序列化为字符串，并通过一个预先定义的 tf.io.TFRecordWriter 写入 TFRecord 文件。

代码示例

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28


import tensorflow as tf 
impport os
###  对原始数据文件名标签进行初始化
data_dir = './DataTes'
train_false = data_dir + '/Max-1/'
train_true = data_dir + '/Max1/'
tfrecord_file = data_dir + '/train.tfrecords'  ## TFrecord文件存储路径

## 遍历文件目录下所有文件路径并存储在列表中
train_false_filenames = [train_false + filename for filename in os.listdir(train_false)] 
train_true_filenames = [train_true + filename for filename in os.listdir(train_true)]
train_filenames = train_false_filenames + train_true_filenames
### 将max-1文件与max1文件夹标注标签
train_labels = [-1] * len(train_false_filenames) + [1] * len(train_true_filenames)

### 迭代读取每张图片，建立 tf.train.Feature 字典和 tf.train.Example 对象，序列化并写入 TFRecord 文件。
with tf.io.TFRecordWriter(tfrecord_file) as writer:
    for filename, label in zip(train_filenames, train_labels):
        image = open(filename, 'rb').read()     # 读取数据集图片到内存，image 为一个 Byte 类型的字符串
        feature = {                             # 建立 tf.train.Feature 字典
            # 图片是一个 Bytes 对象
            'image': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),  
            # 标签是一个 Int 对象
            'label': tf.train.Feature(int64_list=tf.train.Int64List(value=[label]))   
        }
        example = tf.train.Example(features=tf.train.Features(feature=feature)) # 通过字典建立 Example
        writer.write(example.SerializeToString())   # 将Example序列化并写入 TFRecord 文件
        

其中需要注意的点：

1，tf.io.TFRecordWriter`使用时是需要关闭的，上述代码中使用 with 调用函数代替执行了writer.close()。

2， 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 对象。

3，如果只希望保存一个元素而非数组，传入一个只有一个元素的数组即可。

文件执行完毕后，会在tfrecord_file 所指向的文件地址获得一个 train.tfrecords 文件。

3，读取TFrecord文件

首先解释下Dataset.map 方法，使用Dataset.map(转换函数) 方法可以将转换函数映射到数据集每一个元素。

因此，在读取TFrecord文件时可以通过此方法结合tf.io.parse_single_example 函数对数据集中的每一个序列化的 tf.train.Example 对象解码。

代码如下所示：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


raw_dataset = tf.data.TFRecordDataset(tfrecord_file)    # 读取 TFRecord 文件

feature_description = { # 定义Feature结构，告诉解码器每个Feature的类型是什么
    'image': tf.io.FixedLenFeature([], tf.string),
    'label': tf.io.FixedLenFeature([], tf.int64),
}

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'])    # 解码JPEG图片
    return feature_dict['image'], feature_dict['label']

## 利用dataset.map对其中所有元素进行解码
dataset = raw_dataset.map(_parse_example)  ## dataset 已经成为一个可用的数据集了

运行以上代码后，我们获得一个数据集对象 dataset ，这已经成为了一个可以用于训练的 tf.data.Dataset 对象

使用以下代码输出其中图片和对应类标

1
2
3
4
5


import matplotlib.pyplot as plt
for image, label in dataset:
    plt.title('-1' if label == -1 else '1')
    plt.imshow(image.numpy())
    plt.show()

部分结果如图所示

参考文档

本文属于学习过程手记，内容参考了下面两个教程文档

30天吃掉TensorFlow2，数据管道篇

简单粗暴TensorFlow2

文章目录