TensorFlow 深度学习框架（12）-- 多线程输入数据处理框架

（一）队列与多线程

TensorFlow 提供了FIFIQueue 和 RandomShuffleQueue 两种队列。FIFOQueue 顾名思义就是先进先出的意思，RandomShuffleQueu 会将队列中的元素打乱，每次出队列操作得到的是当前队列所有元素中随机的一个。（在训练神经网络时会希望每次使用的训练数据尽量随机）

在TensorFlow中，队列不仅仅是一种数据结构，还是异步计算张量取值的一个重要机制。比如多个线程可以同时向一个队列中写元素，或者同时读取一个队列中的元素。TensorFlow 提供了 tf.Coordinator 和 tf.QueueRunner 两个类 来完成多线程协同的功能。    tf.Coordinator 主要用于协同多个线程一起停止，并提供了 should_stop, request_stop, join 三个函数，在启动线程之前，需要声明一个 tf.Coordinator 类，并将这个类传入每一个创建的线程中。启动的线程需要一直查询 tf.Coordinator 类中提供的 should_stop 函数，当这个函数的值为 True 时，则当前线程也需要退出。每一个启动的线程都可以通过 request_stop 函数来通知其他线程退出。

import tensorflow as tf

# 声明一个先进先出的队列，队列中最多 100 个元素，类型为实数
queue = tf.FIFOQueue(100,"float")
# 定义队列的入队操作
enqueue_op = queue.enqueue([tf.random_normal([1])])

# 使用 tf.train.QueueRunner 来创建多个线程运行队列的入队操作
# tf.train.QueueRunner 的第一个参数给出了被操作的队列 [enqueue_op] * 5 表示启动了 5 个线程
qr = tf.train.QueueRunner(queue,[enqueue_op] * 5)

# 将定义过的 QueueRunner 加入TensorFlow 计算图上指定的集合
tf.train.add_queue_runner(qr)

# 定义出队操作
out_tensor = queue.dequeue()

with tf.Session() as sess:
    # 使用 tf.train.Coordinator 来协同启动的线程
    coord = tf.train.Coordinator()
    # 使用 tf.train.QueueRunner 时，需要明确调用 tf.train.start_queue_runners 
    # 来启动所有的线程。
    threads = tf.train.start_queue_runners(sess = sess,coord = coord)
    for _ in range(3):print sess.run(out_tensor)[0]
    
    # 使用 tf.train.Coordinator 来停止所有的线程
    coord.request_stop()
    coord.join(threads)

（二）多个 TFRecord 文件读取数据

假设我们将训练数据保存到了多个 TFRecords 文件，就以 /path/to/data.tfrecords-00000-of-00002 和 /path/to/data.tfrecords-00001-of-00002 为例。以下代码展示了 tf.train.match_filenames_once 和 tf.train.string_input_producer 函数的使用方法

import tensorflow as tf

# 使用 tf.train.match_filenames_once 函数获取文件列表
files = tf.train.match_filenames_once("/path/to/data.tfrecords-*")

# 通过 tf.train.string_input_producer 函数创建输入队列
# shuffle 设置为 True 可以打乱获取的文件列表
filename_queue = tf.train.string_input_producer(files,shuffle = False)

# 解析一个样本，在 TFRecord 文件存储介绍时提及过
reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)
features = tf.parse_single_example(
    serialized_example,
    features = {
            'i':tf.FixedLenFeature([],tf.int64),
            'j':tf.FixedLenFeature([],tf.int64),
                })

with tf.Session() as sess:
    # 初始化变量
    tf.initialize_all_variables().run()
    '''
    打印文件列表的结果：
        ["/path/to/data.tfrecords-00000-of-00002"
         "/path/to/data.tfrecords-00001-of-00002"]
    '''
    print sess.run(files)
    
    # 声明 tf.train.Coordinator 类来协同不同线程，并启动线程
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess = sess,coord = coord)
    
    # 依次读出 TFRecord 文件中的每一个样例。而且当所有的样例都被读完后，程序会重头开始。
    # 如果限制 num_epochs 为 1 ，那么程序会报错。因为每个文件中只有两个样例
    for i in range(6):
        print sess.run([featutes['i'],features['j']])
    
    # 读取完 6 个数据，告诉 threads 可以停止了
    coord.request_stop()
    coord.join(threads)

（三）组合成 batch 数据

上面只是展示了如何从多个 TFRecord 文件中多线程读取单个样例，但是我们在训练过程中往往是组织成一个 batch 数据。TensorFlow 提供了 tf.train.batch 和 tf.train.shuffle_batch 函数来将单个的样例数据组织成 batch 的形式输出。

import tensorflow as tf

# 假设 i 表示一个样例的特征向量，比如一张图像的像素矩阵。
# j 表示对应的标签
example,label = features['i'],features['j']

# 一个 batch 中的样例数
batch_size = 3

# 组合样例的队列中最多可以存储的样例个数。如果这个队列太大，会占用很多内存资源
# 如果太小，那么出队操作可能会因为没有数据而被阻碍（block），从而导致训练效率降低
# 合适的大小根据 batch_size 决定
capacity = 1000 + 3 * batch_size

# 使用 tf.train.batch 来组合样例。 [example,label] 参数给出了需要组合的元素，
# 一般 example 和 label 分别代表训练数据和正确的标签。
# 当队列长度等于容量 capacity 时，TensorFlow将暂停入队操作，而是只等待元素出队
example_batch, label_batch = tf.train.batch(
            [example,label],batch_size = batch_size,capacity = capacity)

with tf.Session() as sess:
    # 初始化变量
    tf.initialize_all_variables().run()
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess = sess,coord = coord)

    # 获取并打印组合之后的样例。这个输入一般会作为神经网络的输入
    for i in range(2):
        cur_example_batch,cur_label_batch = sess.run(
                [example_batch,label_batch])
        print cur_example_batch,cur_label_batch

    coord.request_stop()
    coord.join(threads)

tf.train.shuffle_batch 函数用法跟 tf.train.batch 函数差不多，但是多了一个 min_after_dequeue 参数来限制出队时队列中元素的最少个数，因为当队列中元素太少时，shuffle 的意义就不大了

example_batch,label_batch = tf.train.shuffle_batch(
    [example,label],batch_size = batch_size,
    capacity = capacity, min_after_dequeue = 30)

（四）完整的输入处理训练框架

以上，已经介绍完了利用队列 多线程从多个 TFRecord 文件中组织成训练数据 batch 的方法。这里给出一个完整的处理和训练框架，算是一个小结。

import tensorflow as tf
files = tf.train.match_filenames_once("/path/to/file_pattern-*")
filename_queue = tf.train.string_input_producer(files,shuffle = False)

reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)
features = tf.parse_single_example(
        serialized_example,
        features = {'image':tf.FixedLenFeature([],tf.string),
                    'label':tf.FixedLenFeature([],tf.int64),
                    'height':tf.FixedLenFeature([],tf.int64),
                    'width':tf.FixedLenFeature([],ft.int64),
                    'channels':tf.FixedLenFeature([],tf.int64),
                    })

image, label = features['image'],features['label']
height,width = features['height'],features['width']
channels = features['channels']

# 从原始图像中解析出像素矩阵，并根据图片尺寸还原图像
decoded_image = tf.decode_raw(image,tf.uint8)
decoded_image.set_shape([height,width,channels])

# 定义神经网络输入图片大小
image_size = 299
# 假设 preprocess_for_train 为图片随机处理函数，具体参考第 11 节
''' 由于这里采用了多线程输入框架，所以图像处理过程不会造成训练的瓶颈 '''
distorted_image = preprocess_for_train(decoded_image,image_size,image_size,None)

# 通过 tf.train.shuffle_batch 组织成 batch 训练数据
min_after_dequeue = 10000
batch_size = 100
capacity = min_after_dequeue + 3 * batch_size
image_batch,label_batch = tf.train.shuffle_batch(
        [distorted_image,label],batch_size = batch_size,
        capacity = capacity,min_after_dequeue = min_after_dequeue)

# 定义神经网络传输和优化过程
logit = inference(image_batch)
loss = calc_loss(logit,label_batch)
# 采用随机梯度下降算法优化损失函数
train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(loss)

with tf.Session() as sess:
    # 初始化变量
    tf.initialize_all_variables().run()
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess = sess,coord = coord)
    
    # 训练神经网络过程
    for i in range(TRAINING_STEPS):
        sess.run(train_step)

    coord.request_stop()
    coord.join(threads)