Tensorflow中使用tfrecord方式读取数据

前言

本博客默认读者对神经网络与Tensorflow有一定了解，对其中的一些术语不再做具体解释。并且本博客主要以图片数据为例进行介绍，如有错误，敬请斧正。

使用Tensorflow训练神经网络时，我们可以用多种方式来读取自己的数据。如果数据集比较小，而且内存足够大，可以选择直接将所有数据读进内存，然后每次取一个batch的数据出来。如果数据较多，可以每次直接从硬盘中进行读取，不过这种方式的读取效率就比较低了。此篇博客就主要讲一下Tensorflow官方推荐的一种较为高效的数据读取方式——tfrecord。

从宏观来讲，tfrecord其实是一种数据存储形式。使用tfrecord时，实际上是先读取原生数据，然后转换成tfrecord格式，再存储在硬盘上。而使用时，再把数据从相应的tfrecord文件中解码读取出来。那么使用tfrecord和直接从硬盘读取原生数据相比到底有什么优势呢？其实，Tensorflow有和tfrecord配套的一些函数，可以加快数据的处理。实际读取tfrecord数据时，先以相应的tfrecord文件为参数，创建一个输入队列，这个队列有一定的容量（视具体硬件限制，用户可以设置不同的值），在一部分数据出队列时，tfrecord中的其他数据就可以通过预取进入队列，并且这个过程和网络的计算是独立进行的。也就是说，网络每一个iteration的训练不必等待数据队列准备好再开始，队列中的数据始终是充足的，而往队列中填充数据时，也可以使用多线程加速。

下面，本文将从以下4个方面对tfrecord进行介绍：

1. tfrecord格式简介
2. 利用自己的数据生成tfrecord文件
3. 从tfrecord文件读取数据
4. 实例测试

1. tfrecord格式简介

这部分主要参考了另一篇博文，Tensorflow 训练自己的数据集（二）（TFRecord）

tfecord文件中的数据是通过tf.train.Example Protocol Buffer的格式存储的，下面是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 的数据结构很简单。tf.train.Example中包含了一个从属性名称到取值的字典，其中属性名称为一个字符串，属性的取值可以为字符串（BytesList ），浮点数列表（FloatList ）或整数列表（Int64List ）。例如我们可以将图片转换为字符串进行存储，图像对应的类别标号作为整数存储，而用于回归任务的ground-truth可以作为浮点数存储。通过后面的代码我们会对tfrecord的这种字典形式有更直观的认识。

2. 利用自己的数据生成tfrecord文件

先上一段代码，然后我再针对代码进行相关介绍。

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
from scipy import misc
import scipy.io as sio


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

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


root_path = '/mount/temp/WZG/Multitask/Data/'
tfrecords_filename = root_path + 'tfrecords/train.tfrecords'
writer = tf.python_io.TFRecordWriter(tf