Tensorflow中的TFRecord、Queue和多线程

Queues, Threads, and Reading Data

输入管线

如果训练数据量较小，Tesnsorflow会把数据一次性加载到内存当中。如果数据过于庞大，Tesorflow需要把存储到硬盘中的数据进行有关处理，分批次加载数据块到内存当中。这些硬盘中的数据可能需要格式化、乱序等的操作。当数据很大的时候，训练过程与处理数据的过程需要用多线程操作，以节约时间。为了使输入输出更加高效，Tensorflow有自己专用的数据格式，处理到大批次数据时，最好转化成专有格式。

TFRecord

TFRecord是一种二进制数据文件，它忽略掉原来的输入数据格式（图片、语音、文字等），统一它们的格式，并把数据进行序列化处理。序列化操作是基于协议缓冲区（protobuf）的，这与平台、编程语言等无关。以下内容主要来自这篇博客，需要使用Google才能访问，在这里就是翻译一下。
第一部分讲述了怎样使用Numpy获取原生数据；第二部分讲述了在没有图模型的情况下，仅仅借助于Tensorflow的内置函数，把数据集转换成TFRecord格式；第三部分讲述了定义一个用于读取二进制文件的模型，并且以随机的方式进行输入。

使用Numpy获取原始数据

在这里我们把图片转化成原始数据，并恢复数据，并检查恢复后的图片是否与原图片一致

# 注意下面的代码在Tensorflow可能无法运行，我这里无法通过
# 就当是一个举例吧
import numpy as np
import skimage.io as io

cat_img = io.imread('cat.jpg')
io.imshow(cat_img)

图片显示：

# 使用 ndarray.tostring() 函数， 把图片转化成string编码
cat_string = cat_img.tostring()

# 把string转化回原来的图片，注意：dtype需要显式地声明，否则报错
# 重构图片的是一维类型，我们需要整个图片的尺寸来恢复原始图片
reconstructed_cat_1d = np.fromstring(cat_string, dtype=np.uint8)

# 在这里重新塑造图片形状
reconstructed_cat_img = reconstructed_cat_1d.reshape(cat_img.shape)

# 检测两个图片是否一致
print(np.allclose(cat_img, reconstructed_cat_img))
# 结果输出True

创建一个.tfrecord文件，并且在不用图模型的情况下读取

在使用之前，需要了解一下protocol buffer的定义，至少要有一个了解，否则有很多细节的东西不能很好的理解。给出官网：https://developers.google.com/protocol-buffers/ (需要Google)
tf.train.Example的定义：

syntax = "proto3";

import "tensorflow/core/example/feature.proto";
option cc_enable_arenas = true;
option java_outer_classname = "ExampleProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.example";

package tensorflow;

message Example {
  Features features = 1;
};

message SequenceExample {
  Features context = 1;
  FeatureLists feature_lists = 2;
};

tf.train.Features的定义：

syntax = "proto3";
option cc_enable_arenas = true;
option java_outer_classname = "FeatureProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.example";

package tensorflow;

// Containers to hold repeated fundamental values.
message BytesList {
  repeated bytes value = 1;
}
message FloatList {
  repeated float value = 1 [packed = true];
}
message Int64List {
  repeated int64 value = 1 [packed = true];
}

// Containers for non-sequential data.
message Feature {
  // Each feature can be exactly one kind.
  oneof kind {
    BytesList bytes_list = 1;
    FloatList float_list = 2;
    Int64List int64_list = 3;
  }
};

message Features {
  // Map from feature name to feature.
  map<string, Feature> feature = 1;
};
message FeatureList {
  repeated Feature feature = 1;
};

message FeatureLists {
  // Map from feature name to feature list.
  map<string, FeatureList> feature_list = 1;
};

以上是基于protocol buffer协议的一个数据格式定义， 这是使用TFRecord数据格式的基础。
下面，我们将展示在没有图模型参与的情况下的小批量数据写入和读取操作。

# 定义数据的路径
filename_pairs = [
('/home/dpakhom1/tf_projects/segmentation/VOCdevkit/VOCdevkit/VOC2012/JPEGImages/2007_000032.jpg',
'/home/dpakhom1/tf_projects/segmentation/VOCdevkit/VOCdevkit/VOC2012/SegmentationClass/2007_000032.png'),
('/home/dpakhom1/tf_projects/segmentation/VOCdevkit/VOCdevkit/VOC2012/JPEGImages/2007_000039.jpg',
'/home/dpakhom1/tf_projects/segmentation/VOCdevkit/VOCdevkit/VOC2012/SegmentationClass/2007_000039.png'),
('/home/dpakhom1/tf_projects/segmentation/VOCdevkit/VOCdevkit/VOC2012/JPEGImages/2007_000063.jpg',
'/home/dpakhom1/tf_projects/segmentation/VOCdevkit/VOCdevkit/VOC2012/SegmentationClass/2007_000063.png')
                 ]

from PIL import Image
import numpy as np
import skimage.io as io
import tensorflow as tf


# 定义的两个特征写入函数
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]))

# 写入的文件名 
tfrecords_filename = 'pascal_voc_segmentation.tfrecords'

# 定义的TFRecordWriter
writer = tf.python_io.TFRecordWriter(tfrecords_filename)

# 用于保存图片的初始信息
original_images = []

# 对于路径中的每一个图片，img_path存储路径，annotation_path存储图片编号
for img_path, annotation_path in filename_pairs:
    # 转换成Numpy数组的格式
    img = np.array(Image.open(img_path))
    annotation = np.array(Image.open(annotation_path))

    # 保存原始图片的高度和宽度
    height = img.shape[0]
    width = img.shape[1]

    # 把数据追加到原始数组中
    original_images.append((img, annotation))

    # 转换成string类型
    img_raw = img.tostring()
    annotation_raw = annotation.tostring()

    # 确定protocol buffer的格式，用于写入`.tfrecord`文件
    example = tf.train.Example(features=tf.train.Features(feature={
        'height': _int64_feature(height),
        'width': _int64_feature(width),
        'image_raw': _bytes_feature(img_raw),
        'mask_raw': _bytes_feature(annotation_raw)}))
    # 数据序列化之后写入
    writer.write(example.SerializeToString())
writer.close()
# 用于存储重构图片的数据
reconstructed_images = []
# 确定存储的路径
record_iterator = tf.python_io.tf_record_iterator(path=tfrecords_filename)

for string_record in record_iterator:
    # 获取数据，
    example = tf.train.Example()
    example.ParseFromString(string_record)

    height = int(example.features.feature['height']
                                 .int64_list
                                 .value[0])

    width = int(example.features.feature['width']
                                .int64_list
                                .value[0])

    img_string = (example.features.feature['image_raw']
                                  .bytes_list
                                  .value[0])

    annotation_string = (example.features.feature['mask_raw']
                                .bytes_list
                                .value[0])
    # 恢复数据成一维的
    img_1d = np.fromstring(img_string, dtype=np.uint8)
    二维展开图像
    reconstructed_img = img_1d.reshape((height, width, -1))
    # 获取路径
    annotation_1d = np.fromstring(annotation_string, dtype=np.uint8)

    # Annotations don't have depth (3rd dimension)
    reconstructed_annotation = annotation_1d.reshape((height, width))

    reconstructed_images.append((reconstructed_img, reconstructed_annotation))

最后以MNIST数据集为例子：

import tensorflow as tf
from tensorflow.contrib.learn.python.learn.datasets import mnist
import os

save_dir = "/home/erick/Desktop/mnist"

# Download data to save_dir
data_sets = mnist.read_data_sets(save_dir,
                                 dtype=tf.uint8,
                                 reshape=False,
                                 validation_size=1000)

data_splits = ["train", "test", "validation"]
for d in range(len(data_splits)):
    print("saving" + data_splits[d])
    data_set = data_sets[d]

    # 实例化TFRecordWriter对象，并确定写入的目录
    filename = os.path.join(save_dir, data_splits[d] + '.tfrecords')
    writer = tf.python_io.TFRecordWriter(filename)

    # 遍历每个图片
    for index in range(data_set.images.shape[0]):
        # 把图片从Numpy数组转化成二进制字符串
        image = data_set.images[index].tostring()
        # 这是一个存储协议
        example = tf.train.Example(features=tf.train.Features(features={
            'height': tf.train.Feature(int64_list=
                                       tf.train.Int64List(value=
                                                          [data_set.images.shape[1]])),
            'width': tf.train.Feature(int64_list=
                                      tf.train.Int64List(value=
                                                         [data_set.images.shape[2]])),
            'depth': tf.train.Feature(int64_list=
                                      tf.train.Int64List(value=
                                                         [data_set.images.shape[3]])),
            'label': tf.train.Feature(int64_list=
                                      tf.train.Int64List(value=
                                                         [int(data_set.labels[index])])),
            'image_raw': tf.train.Feature(bytes_list=
                                      tf.train.BytesList(value=
                                                         [image]))
        }))

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

Queues

这里的队列指的是计算图中的队列，与一般队列不同，这里的队列需要有计算图的介入才有效。

Enqueuing and Dequeuing 入队和出队

入队和出队的操作需要提前声明，不过声明的完后不会运行，只有当图模型执行该操作后，定义的操作才会有效！不能对一个空队列执行出队操作，否则会导致主线程一直处于挂起状态。

import tensorflow as tf

sess = tf.InteractiveSession()
# 建立容量是10 的string存储string的队列
queue1 = tf.FIFOQueue(capacity=10, dtypes=[tf.string])
# 入队操作
enque_op = queue1.enqueue(["F"])
print(sess.run(queue1.size()))  # 输出0，因为此时没有计算图执行该操作
enque_op.run()  # 计算图介入操作
print(sess.run(queue1.size()))  # 输出1

# 继续执行入队的操作
enque_op = queue1.enqueue(["I"])
enque_op.run()
enque_op = queue1.enqueue(["F"])
enque_op.run()
enque_op = queue1.enqueue(["O"])
enque_op.run()

print(sess.run(queue1.size()))  # 输出4

x = queue1.dequeue()  # 出队操作
# 下面的出队操作最多执行4次，对空队列执行出队操作，会让主线程永远挂起！！
print(x.eval())  # 输出 F
print(x.eval())  # 输出 I
print(x.eval())  # 输出 F
print(x.eval())  # 输出 O

多个元素同时出队

可以使用dequeue_many()一次使多个元素出队，不过出队的个数不能超过队列的元素的个数，否则主线程会挂起。

import tensorflow as tf

sess = tf.InteractiveSession()
# 建立容量是10的string存储string的队列，需要提前说明队列的类型
queue1 = tf.FIFOQueue(capacity=10, dtypes=[tf.string], shapes=[()])
# 入队操作
enque_op = queue1.enqueue(["F"])
print(sess.run(queue1.size()))  # 输出0
enque_op.run()  # 计算图介入操作
print(sess.run(queue1.size()))  # 输出1

# 继续执行入队的操作
enque_op = queue1.enqueue(["I"])
enque_op.run()
enque_op = queue1.enqueue(["F"])
enque_op.run()
enque_op = queue1.enqueue(["O"])
enque_op.run()

print(sess.run(queue1.size()))  # 输出4

# 执行多次出队操作
inputs = queue1.dequeue_many(4)
print(inputs.eval())  # 输出 [b'F' b'I' b'F' b'O']

Threading and Queues 线程和队列

一个Tensorflow的session可以是多线程的，一个session可以同时并行执行多个操作。但是，如果一次调用sess.run()不能充分利用资源，可以利用多个并行调用增加吞吐量。一个典型的应用场景：执行图像识别算法训练时，一方面要把图像转换成标准的输入流格式，另一方面要用输入流对神经网络进行训练，很明显输入和训练的过程是可以并行执行的。
关于python多线程的知识，可以参考一下这几篇博客
但是，多线程使用的时候，必须同时停止；如果有异常，必须在所有线程结束后，捕获所有的异常；队列必须要在线程结束的时候适当的停止。Tensorflow提供了tf.train.Coordinator和tf.train.QueueRunner这两个类在协助使用多线程技术。前者用来同时停止多个线程，或者等待所有线程结束后报告异常；后者用来创建一定数量的线程，这些线程写作让tensor加入同一个队列。

Coordinator

• tf.train.Coordinator.should_stop： 返回True，如果线程需要停止
• tf.train.Coordinator.request_stop：要求线程停止
• tf.train.Coordinator.join： 一直等到特定的线程终止

我们首先建立一个Coordinator对象，之后创建一定数量的可以使用线程协调器的线程。当should_stop()返回True时，所有的线程循环停止。任何一个线程都可以决定计算图模型的停止时间，只要该线程调用request_stop()函数，那么其他的线程就会停止，且它们的shoule_stop()函数会返回True。
官方文档给出的一般格式：

# Thread body: loop until the coordinator indicates a stop was requested.
# If some condition becomes true, ask the coordinator to stop.
def MyLoop(coord):
  while not coord.should_stop():
    ...do something...
    if ...some condition...:
      coord.request_stop()

# Main thread: create a coordinator.
coord = tf.train.Coordinator()

# Create 10 threads that run 'MyLoop()'
threads = [threading.Thread(target=MyLoop, args=(coord,)) for i in xrange(10)]

# Start the threads and wait for all of them to stop.
for t in threads:
  t.start()
coord.join(threads)

实例代码：

import tensorflow as tf
import threading, time


def loop(coord, i):
    while not coord.should_stop():
        # 输出线程编号和当前系统时间
        print(i, time.strftime('%H:%M:%S', time.localtime(time.time())))
        # 睡眠一秒，防止运行过快
        time.sleep(1)
        # 5号线程申请终止
        if i == 5:
            coord.request_stop()


# 主线程
coord = tf.train.Coordinator()
# 申请10个线程
threads = [threading.Thread(target=loop, args=(coord, i)) for i in range(10)]

# 启动所有线程，并等待线程结束
for t in threads:
    t.start()
coord.join(threads)

上述代码中，5号线程仅出现一次，并且在休眠一秒后，终止了所有的线程。

QueueRunner

QueueRunner创建了一定数量的线程，来重复的执行入队操作。
给出一般结构：

import tensorflow as tf

gen_random_normal = tf.random_normal(shape=())
queue = tf.RandomShuffleQueue(capacity=100, dtypes=[tf.float32], min_after_dequeue=1)
enqueue_op = queue.enqueue(gen_random_normal)
# Create a queue runner that will run 4 threads in parallel to enqueue examples.
qr = tf.train.QueueRunner(queue, [enqueue_op] * 4)

# Launch the graph.
sess = tf.Session()
# Create a coordinator, launch the queue runner threads.
coord = tf.train.Coordinator()
enqueue_threads = qr.create_threads(sess, coord=coord, start=True)

# Run the training loop, controlling termination with the coordinator.
# for step in range(1000000):
#     if coord.should_stop():
#         break
#     sess.run(train_op)

# When done, ask the threads to stop.
coord.request_stop()
# And wait for them to actually do it.
coord.join(enqueue_threads)