TensorFlow TFRecord格式

一、TFRecord格式介绍

TFRecord 文件中的数据都是通过 tf.train.Example 以 Protocol Buffer(以下简称PB) 的格式存储。

这里讲一下PB， PB是Google的一种数据交换的格式，他独立于语言，独立于平台，以二进制的形式存在。它能更好的利用内存，方便复制和移动。

下面我们看一下 tf.train.Example 的定义：

message Example {
 Features features = 1;
};

message Features{
 map<string,Feature> featrue = 1;
};

message Feature{
    oneof kind{
        BytesList bytes_list = 1;
        FloatList float_list = 2;
        Int64List int64_list = 3;
    }
};

从代码中我们可以看到， tf.train.Example 包含了一个字典，key是字符串，value为Feature，Feature可以取值为字符串（BytesList ）、浮点数列表（FloatList ）、整型数列表（Int64List ）。

二、TFRecord 文件的写入

实现TFRecord 文件的写入分为以下几步：

1、首先要获取我们需要转化的数据

2、将数据填入到Example PB，并且将Example PB 转化为一个字符串

3、通过 tf.python_io.TFRecordWriter 将字符串写入TFRecord 文件

下面的程序（有详细注释）给出了如何将MNIST输入数据转化为TFRecord格式：

#coding:utf-8
#将MNIST输入数据转化为TFRecord的格式

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import numpy as np

#把传入的value转化为整数型的属性，int64_list对应着 tf.train.Example 的定义
def _int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

#把传入的value转化为字符串型的属性，bytes_list对应着 tf.train.Example 的定义
def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
#读取MNIST数据
mnist = input_data.read_data_sets("/home/sun/AI/CNN/handWrite1/data", dtype=tf.uint8, one_hot=True)
#训练数据的图像，可以作为一个属性来存储
images = mnist.train.images
#训练数据所对应的正确答案，可以作为一个属性来存储
labels = mnist.train.labels
#训练数据的图像分辨率，可以作为一个属性来存储
pixels = images.shape[0]
#训练数据的个数
num_examples = mnist.train.num_examples
#指定要写入TFRecord文件的地址
filename = "/home/sun/AI/grammar/TFRecord/output.tfrecords"
#创建一个write来写TFRecord文件
writer = tf.python_io.TFRecordWriter(filename)
for index in range(num_examples):
    #把图像矩阵转化为字符串
    image_raw = images[index].tostring()
    #将一个样例转化为Example Protocol Buffer，并将所有的信息写入这个数据结构
    example = tf.train.Example(features=tf.train.Features(feature={
        'pixels': _int64_feature(pixels),
        'label': _int64_feature(np.argmax(labels[index])),
        'image_raw': _bytes_feature(image_raw)}))
    #将 Example 写入TFRecord文件
    writer.write(example.SerializeToString())
    
writer.close()

三、TFRecord文件的读取

下面的程序（有详细注释）给出了如何读取上面程序生成的 TFRecord 文件：

#coding:utf-8
#读取TFRecord文件中的数据

import tensorflow as tf

#创建一个reader来读取TFRecord文件中的样例
reader = tf.TFRecordReader()
#通过 tf.train.string_input_producer 创建输入队列
filename_queue = tf.train.string_input_producer(["/home/sun/AI/grammar/TFRecord/output.tfrecords"])
#从文件中读取一个样例
_, serialized_example = reader.read(filename_queue)
#解析读入的一个样例
features = tf.parse_single_example(
	serialized_example,
	features={
		#这里解析数据的格式需要和上面程序写入数据的格式一致
		'image_raw': tf.FixedLenFeature([], tf.string),
		'pixels': tf.FixedLenFeature([], tf.int64),
		'label': tf.FixedLenFeature([], tf.int64),
	})
#tf.decode_raw可以将字符串解析成图像对应的像素数组
images = tf.decode_raw(features['image_raw'], tf.uint8)
#tf.cast可以将传入的数据转化为想要改成的数据类型
labels = tf.cast(features['label'], tf.int32)
pixels = tf.cast(features['pixels'], tf.int32)

sess = tf.Session()
#启动多线程处理输入数据
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)

#每次运行可以读取TFRecord文件中的一个样例。当所有样例都读完之后，再次样例中的程序会重头读取
for i in range(10):
	image, label, pixel = sess.run([images, labels, pixels])
	print label

在上面第三小节读取TFRecord文件中，用到了队列和多线程，对于队列和多线程详细的用法，可以参考我这一篇博文：点击打开链接

TensorFlow TFRecord格式

猜你喜欢