tensorflow高效读取数据之tfrecord

TFRecord输入数据格式

Tensorflow提供了一种统一的格式来存储数据，这个格式就是TFRecord，接下来介绍如何使用TFRecord来统一输入数据的格式。

• TFRcord格式介绍

TFRcord文件中的数据都是通过tf.train.Example Protocol Buffer的格式来存储的，以下代码给出了tf.train.Example的定义。

message example{
    Feature feature = 1;
}

message Feature {
    map < string,Feature > feature = 1;
}
message Feature {
    oneof kind {
        BytesList bytes_list = 1;
        FloatList float_list = 2;
        Int64List int64_list = 3;
    }
};

从以上代码可以看出tf.train.Example的数据结构是比较简洁的。tf.train.Example中包含了一个从属性名称到取值的字典。其中属性名称是一个字符串，属性数值可以为字符串（BytesList），实数列表（FloatList）或者整数列表（Int64List）。

---

TFRcord样例程序

• 本样例展示了把mnist数据集转化为TFRcord格式。

#coding；utf-8

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

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

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

mnist = input_data.read_data_sets('mnist/',dtype=tf.uint8,one_hot=True)

images = mnist.train.images
labels = mnist.train.labels

pixels = images.shape[1]

num_examples = mnist.train.num_examples

filename = 'output/train.tfrecords'

writer = tf.python_io.TFRecordWriter(filename)
for index in range(num_examples):
    image_raw = images[index].tostring()
    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)}))

    writer.write(example.SerializeToString())
writer.close()

以上程序可以将mnist数据集中所有的训练数据存储到一个TFRcord文件中，当数据量较大时，也可以将数据写入多个TFRcord文件。

---

• 本样例展示了读取TFRcord中的数据。

#coding；utf-8

import tensorflow as tf
import matplotlib as plt
import cv2

reader = tf.TFRecordReader()
#创建一个队列来维护输入的文件列表
filename_queue = tf.train.string_input_producer(['output/train.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),
                                   })
images = tf.decode_raw(features['image_raw'],tf.uint8)
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)
for i in range(10):
    image,label,pixel = sess.run([images,labels,pixels])
    cv2.imshow('image',image)

以上程序可以将之前生成的train.tfrecords文件转换为可读图片。