【Tensorflow】超大规模数据集解决方案：通过线程来预取（上）

环境Tensorflow1.2（这是最新的一个版本），python2.7

这是我重点要讲的解决方案，我怕篇幅过长，分成了两篇，上篇介绍一下预备的东西，下篇来进行实验

一.Tensorflow中的队列机制

队列和线程是Temsorflow中实现异步的重要工具。为什么要异步？用一个形象的例子来解释这个问题。

可以把数据导入的过程看作io操作，在数据规模极大的情况下，io请求需要大量时间执行。同步意味着我们一次处理完io请求，然后再执行程序后面的操作，所以之前在【Tensorflow】怎样为你的网络预加工和打包训练数据？（二）：小数据集的处理方案和【Tensorflow】怎样为你的网络预加工和打包训练数据？（一）中的处理都可以看作同步的，因为都是一次处理完所有的数据，然后在feed给我们的网络。而这里，我们需要扩展一个异步的概念，所谓异步，也就是开辟一个线程来单独处理这个io请求，处理完成，就通知给主程序告诉它“我已经完成了”，在此期间，主程序可以去做其他的事，也就是io请求并不耽误程序的执行。所以异步方式可以显著提高效率，那是不是说异步一定比同步好呢，当然不是，异步只适用于等待时间很长的情况，如果处理小数据集，就不如同步方式了。

Tensorflow是怎样实现异步的？这里需要用到队列这个数据结构，先让我们看看Tensorflow是如何实现队列的。

我们来看一个简单的例子。创建一个“先进先出”队列（FIFOQueue）并填充零。然后，创建一个将元素从队列中取出的图，将该元素加一，并将其放回队列的末尾。缓慢地，队列上的数字增加。用如下的图来表示

和如上图类似，tensorflow中使用RandomShuffleQueue作为输入框架，来准备训练数据。这种机制的运行流程大致如下：

1.使用多线程来准备训练样本，并把他们放入队列

2.一个训练线程执行优化，并吧mini-batch数据提取出队列

扫描二维码关注公众号，回复： 3295788 查看本文章

TensorFlow Session对象是多线程的，所以多个线程可以轻松地使用相同的Session并并行运行op。但是，实现一个驱动线程的Python程序并不容易。所有线程必须能够一起停止，异常必须被捕获和报告，队列必须在停止时正确关闭。TensorFlow提供了两个类来帮助线程驱动：tf.train.Coordinator和tf.train.QueueRunner。这两个类被设计为一起使用，协调器类帮助多个线程一起停止，并向等待其停止的程序报告异常。QueueRunner类用于创建多个线程，用于协调在同一队列中放入张量。

二.线程协调器

先看一下一些关键方法：

tf.train.Coordinator.should_stop ： returns True if the threads should stop

tf.train.Coordinator.request_stop ：requests that threads should stop

tf.train.Coordinator.join ：waits until the specified threads have stopped

首先创建一个Coordinator对象，然后创建一些使用协调器的线程。线程执行运行循环，当should_stop（）返回True时停止。任何线程都可以决定计算是否停止。它只需要调用request_stop（），其他线程将停止，因为should_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)

三.队列运行器

QueueRunner类创建了一些重复执行入队操作的线程。这些线程可以使用协调器一起停止，此外，队列运行器运行一个队列关闭的线程，一旦向协调器报告异常，则会自动关闭队列。

创建TensorFlow队列（例如tf.RandomShuffleQueue）用作样本输入的过程如下：

example = ...ops to create one example...
# Create a queue, and an op that enqueues examples one at a time in the queue.
queue = tf.RandomShuffleQueue(...)
enqueue_op = queue.enqueue(example)
# Create a training graph that starts by dequeuing a batch of examples.
inputs = queue.dequeue_many(batch_size)
train_op = ...use 'inputs' to build the training part of the graph...

然后我们要为这个队列创建一个队列运行器，使用多线程来执行入队操作：

# 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 xrange(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)

四.具体的使用方法

在训练当中如何结合以上的工具呢？tf.train.QueueRunner对象要求您在运行任何训练步骤之前调用tf.train.start_queue_runners，否则线程将永久挂起。所以在训练之前要调用tf.train.start_queue_runners启动输入流水线的线程，填充样本队列，以使得样本的出队操作能成功执行。Tensorflow提供了一个参考代码模板：

# Create the graph, etc.
init_op = tf.global_variables_initializer()

# Create a session for running operations in the Graph.
sess = tf.Session()

# Initialize the variables (like the epoch counter).
sess.run(init_op)

# Start input enqueue threads.
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)

try:
    while not coord.should_stop():
        # Run training steps or whatever
        sess.run(train_op)

except tf.errors.OutOfRangeError:
    print('Done training -- epoch limit reached')
finally:
    # When done, ask the threads to stop.
    coord.request_stop()

# Wait for threads to finish.
coord.join(threads)
sess.close()

ok，到这里让我们详细解释一下上面的代码可以做什么，首先我们创建计算图，第一阶段获得文件名路径并将其排入文件名队列。第二阶段使用文件名，生成样本，并将它们排列在样本队列中。一旦启动运行这些入队操作的线程，我们的训练循环就可以从样本队列中取出训练样本用以训练。用一张图来表示：

但是在我们实际使用中，要定义一个以上所描述的pipeline是非常复杂的，所以一般就只使用一个文件名队列，而把训练数据的预处理放在主程序当中，因为预处理一个batch的训练数据通常不需要花费太多时间。我们可以使用tf.train.add_queue_runner来添加一个QueueRunner，创建完图后，tf.train.start_queue_runners函数将在图中查询每个QueueRunner，以启动其运行入队操作的线程。此外，我们还要设定epoch的限制，保证在epoch到达时，程序会正确停止，并抛出一个tf.errors.OutOfRangeError异常。

五.一些小的示例程序

首先我们来看第一个先进先出队列例子的实现代码

import tensorflow as tf


q = tf.FIFOQueue(3, "float")
init = q.enqueue_many(vals=[[0., 0., 0.],])

x = q.dequeue()
y = x+1
q_inc = q.enqueue([y])

with tf.Session() as sess:
    init.run()
    q_inc.run()
    q_inc.run()
    q_inc.run()
    q_inc.run()
    print sess.run(x)
    print sess.run(x)
    print sess.run(x)

需要注意的是，enqueue_many方法在实现的时候还是有比较多的坑的，参数vals必须二维，而且要写成[[0,0,0]]的形式而不是，[[0],[0],[0]]，这个要重点记一下。

整个流程的动图上面贴过了，这里再贴一遍吧

然后就是RandomShuffleQueue的例子：

import tensorflow as tf


q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")

sess = tf.Session()
for i in range(10):
    sess.run(q.enqueue('File:'+str(i)))

for i in range(10):
    print(sess.run(q.dequeue()))

因为是 RandomShuffleQueue，所以出队顺序是随机的，一直到这里，入队都是用主线程来做的，以上简单版程序没有用多线程。

最后是加上多线程的终极版，让我们来看看：

import tensorflow as tf

q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")
enqueue_op = q.enqueue('File:')
qr = tf.train.QueueRunner(q, enqueue_ops=[enqueue_op] * 1)
sess = tf.Session()
enqueue_threads = qr.create_threads(sess, start=True)
for i in range(100):
    print(sess.run(q.dequeue()))

好了，运行这个程序，我们会发现输出了100个File，说明在出队的时候，是又有一个线程在做入队操作不停的补充队列的。但是发现一个问题，貌似这个程序自动终止不了？因为我们创建的线程还在运行。

所以要加上协调器，来使得程序得以终止：

import tensorflow as tf

q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")
enqueue_op = q.enqueue('File:')
qr = tf.train.QueueRunner(q, enqueue_ops=[enqueue_op] * 1)
sess = tf.Session()
coord = tf.train.Coordinator()
enqueue_threads = qr.create_threads(sess, start=True,coord=coord)
for i in range(100):
    print(sess.run(q.dequeue()))
coord.request_stop()   
coord.join(enqueue_threads)

这样程序就得以终止了，然后，如果你想再设置一些epoch limit的异常：

import tensorflow as tf

q = tf.RandomShuffleQueue(capacity=10,min_after_dequeue=0, dtypes="string")
enqueue_op = q.enqueue('File:')
qr = tf.train.QueueRunner(q, enqueue_ops=[enqueue_op] * 1)
sess = tf.Session()
coord = tf.train.Coordinator()
enqueue_threads = qr.create_threads(sess, start=True,coord=coord)
try:
    for i in range(100):
        print(sess.run(q.dequeue()))
        if i>=50:
            coord.request_stop()
            coord.join(enqueue_threads)
except tf.errors.OutOfRangeError:
    print('Done training -- epoch limit reached')
finally:
    coord.request_stop()
    coord.join(enqueue_threads)

我们可以在i>50以后强行终止线程，这样就可以获得队列为空且出队的异常了。当然这个程序仅供实验学习用，本身没有应用价值。

下一篇，我们会付诸实践，上手用数据来实现以上所述的处理流程。