tfrecord文件

在学习TensorFlow的过程中，有很多小伙伴反映读取数据这一块很难理解。确实这一块官方的教程比较简略，网上也找不到什么合适的学习材料。今天这篇文章就以图片的形式，用最简单的语言，为大家详细解释一下TensorFlow的数据读取机制，文章的最后还会给出实战代码以供参考。

TensorFlow读取机制图解

首先需要思考的一个问题是，什么是数据读取？以图像数据为例，读取数据的过程可以用下图来表示：

假设我们的硬盘中有一个图片数据集0001.jpg，0002.jpg，0003.jpg……我们只需要把它们读取到内存中，然后提供给GPU或是CPU进行计算就可以了。这听起来很容易，但事实远没有那么简单。事实上，我们必须要把数据先读入后才能进行计算，假设读入用时0.1s，计算用时0.9s，那么就意味着每过1s，GPU都会有0.1s无事可做，这就大大降低了运算的效率。

如何解决这个问题？方法就是将读入数据和计算分别放在两个线程中，将数据读入内存的一个队列，如下图所示：

读取线程源源不断地将文件系统中的图片读入到一个内存的队列中，而负责计算的是另一个线程，计算需要数据时，直接从内存队列中取就可以了。这样就可以解决GPU因为IO而空闲的问题！

而在TensorFlow中，为了方便管理，在内存队列前又添加了一层所谓的“文件名队列”。

为什么要添加这一层文件名队列？我们首先得了解机器学习中的一个概念：epoch。对于一个数据集来讲，运行一个epoch就是将这个数据集中的图片全部计算一遍。如一个数据集中有三张图片A.jpg、B.jpg、C.jpg，那么跑一个epoch就是指对A、B、C三张图片都计算了一遍。两个epoch就是指先对A、B、C各计算一遍，然后再全部计算一遍，也就是说每张图片都计算了两遍。

TensorFlow使用文件名队列+内存队列双队列的形式读入文件，可以很好地管理epoch。下面我们用图片的形式来说明这个机制的运行方式。如下图，还是以数据集A.jpg, B.jpg, C.jpg为例，假定我们要跑一个epoch，那么我们就在文件名队列中把A、B、C各放入一次，并在之后标注队列结束。

程序运行后，内存队列首先读入A（此时A从文件名队列中出队）：

再依次读入B和C：

此时，如果再尝试读入，系统由于检测到了“结束”，就会自动抛出一个异常（OutOfRange）。外部捕捉到这个异常后就可以结束程序了。这就是TensorFlow中读取数据的基本机制。如果我们要跑2个epoch而不是1个epoch，那只要在文件名队列中将A、B、C依次放入两次再标记结束就可以了。

TensorFlow读取数据机制的对应函数

如何在TensorFlow中创建上述的两个队列呢？

对于文件名队列，我们使用tf.train.string_input_producer函数。这个函数需要传入一个文件名list，系统会自动将它转为一个文件名队列。

此外tf.train.string_input_producer还有两个重要的参数，一个是num_epochs，它就是我们上文中提到的epoch数。另外一个就是shuffle，shuffle是指在一个epoch内文件的顺序是否被打乱。若设置shuffle=False，如下图，每个epoch内，数据还是按照A、B、C的顺序进入文件名队列，这个顺序不会改变：

如果设置shuffle=True，那么在一个epoch内，数据的前后顺序就会被打乱，如下图所示：

在TensorFlow中，内存队列不需要我们自己建立，我们只需要使用reader对象从文件名队列中读取数据就可以了，具体实现可以参考下面的实战代码。

除了tf.train.string_input_producer外，我们还要额外介绍一个函数：tf.train.start_queue_runners。初学者会经常在代码中看到这个函数，但往往很难理解它的用处，在这里，有了上面的铺垫后，我们就可以解释这个函数的作用了。

在我们使用tf.train.string_input_producer创建文件名队列后，整个系统其实还是处于“停滞状态”的，也就是说，我们文件名并没有真正被加入到队列中（如下图所示）。此时如果我们开始计算，因为内存队列中什么也没有，计算单元就会一直等待，导致整个系统被阻塞。

而使用tf.train.start_queue_runners之后，才会启动填充队列的线程，这时系统就不再“停滞”。此后计算单元就可以拿到数据并进行计算，整个程序也就跑起来了，这就是函数tf.train.start_queue_runners的用处。

 

#写tfrecord
def write(input_file, output_file):
    writer = tf.python_io.TFRecordWriter(output_file)  # 定义writer，传入目标文件路径
    path = input_file
    file_names = [f for f in os.listdir(path) if f.endswith('.jpg')]  # 获取待存文件路径
    for file_name in file_names:
        img = cv2.imread(path + file_name)
        shape=img.shape
        raw_img = img.tobytes()  # 需要把图片文件转化成bytes形式（二进制比特流）
        file_name_byte=file_name.encode('utf-8')
        # 把数据合并成feature，注意这里的"value="后面一定要是一个"[]"形式的列表，否则读取的时候会出现can't parse的情况
        features = tf.train.Features(
            feature={'img_name': tf.train.Feature(bytes_list=tf.train.BytesList(value=[file_name_byte])),
                     'img_shape_a': tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[0]])),
                     'img_shape_b': tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[1]])),
                     'img_shape_c': tf.train.Feature(int64_list=tf.train.Int64List(value=[shape[2]])),
                     'raw_img': tf.train.Feature(bytes_list=tf.train.BytesList(value=[raw_img]))})
        # 把features存入example
        example = tf.train.Example(features=features)
        # example序列化，并写入文件
        writer.write(example.SerializeToString())
    writer.close()


# input_file = '../samples/'
# output_file = 'samples.tfrecords'
# write(input_file, output_file)
# print ('Write tfrecords: %s done' % output_file)


#读tfrecord

def read_and_decode(file_name):
    filename_queue = tf.train.string_input_producer([file_name])
    #tf.train.slice_input_producer
    reader = tf.TFRecordReader()
    _, serialized_example = reader.read(filename_queue)
    features = tf.parse_single_example(serialized_example,
                                       features={"img_name": tf.FixedLenFeature([], tf.string),
                                                 "img_shape_a": tf.FixedLenFeature([], tf.int64),
                                                 "img_shape_b": tf.FixedLenFeature([], tf.int64),
                                                 "img_shape_c": tf.FixedLenFeature([], tf.int64),
                                                 "raw_img": tf.FixedLenFeature([], tf.string)})
    img_name = features["img_name"]
    shape_a = features["img_shape_a"]
    shape_b = features["img_shape_b"]
    shape_c = features["img_shape_c"]

    image = tf.decode_raw(features['raw_img'], tf.uint8)

    image = tf.reshape(image, [shape_a,shape_b,shape_c])
    return img_name, image


path = 'samples.tfrecords'
with tf.Session() as sess:
    img_name, img = read_and_decode(path)
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)
    try:
        for i in range(2):
            name, image = sess.run([img_name, img])
            print(name)
            plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            plt.show()
    except tf.errors.OutOfRangeError:
        print
        'Done training -- epoch limit reached'
    finally:
        coord.request_stop()

    coord.join(threads)

 

 

TensorFlow的Session对象是支持多线程的，可以在同一个会话（Session）中创建多个线程，并行执行。在Session中的所有线程都必须能被同步终止，异常必须能被正确捕获并报告，会话终止的时候， 队列必须能被正确地关闭。

TensorFlow提供了两个类来实现对Session中多线程的管理：tf.Coordinator和 tf.QueueRunner，这两个类往往一起使用。

Coordinator类用来管理在Session中的多个线程，可以用来同时停止多个工作线程并且向那个在等待所有工作线程终止的程序报告异常，该线程捕获到这个异常之后就会终止所有线程。使用 tf.train.Coordinator()来创建一个线程管理器（协调器）对象。

QueueRunner类用来启动tensor的入队线程，可以用来启动多个工作线程同时将多个tensor（训练数据）推送入文件名称队列中，具体执行函数是 tf.train.start_queue_runners ， 只有调用 tf.train.start_queue_runners 之后，才会真正把tensor推入内存序列中，供计算单元调用，否则会由于内存序列为空，数据流图会处于一直等待状态。

tf中的数据读取机制如下图：

 

调用 tf.train.slice_input_producer，从 本地文件里抽取tensor，准备放入Filename Queue（文件名队列）中;
调用 tf.train.batch，从文件名队列中提取tensor，使用单个或多个线程，准备放入文件队列;
调用 tf.train.Coordinator() 来创建一个线程协调器，用来管理之后在Session中启动的所有线程;
调用tf.train.start_queue_runners, 启动入队线程，由多个或单个线程，按照设定规则，把文件读入Filename Queue中。函数返回线程ID的列表，一般情况下，系统有多少个核，就会启动多少个入队线程（入队具体使用多少个线程在tf.train.batch中定义）;
文件从 Filename Queue中读入内存队列的操作不用手动执行，由tf自动完成;
调用sess.run 来启动数据出列和执行计算;
使用 coord.should_stop()来查询是否应该终止所有线程，当文件队列（queue）中的所有文件都已经读取出列的时候，会抛出一个 OutofRangeError 的异常，这时候就应该停止Sesson中的所有线程了;
使用coord.request_stop()来发出终止所有线程的命令，使用coord.join(threads)把线程加入主线程，等待threads结束。