Tensorflow深度学习之十七：队列与多线程

声明：本文参考《Tensorflow：实战Google深度学习框架》一书

在使用TensorFlow进行异步计算时，队列是一种强大的机制。

正如TensorFlow中的其他组件一样，队列就是TensorFlow图中的节点。这是一种有状态的节点，就像变量一样：其他节点可以修改它的内容。具体来说，其他节点可以把新元素插入到队列后端(rear)，也可以把队列前端(front)的元素删除。

为了感受一下队列，让我们来看一个简单的例子。我们先创建一个“先入先出”的队列（FIFOQueue），并将其内部所有元素初始化为零。然后，我们构建一个TensorFlow图，它从队列前端取走一个元素，加上1之后，放回队列的后端。慢慢地，队列的元素的值就会增加。

对于变量，可以通过赋值操作修改变量的取值。对于队列，修改队列的操作主要有enqueue，enqueue_many和dequeue。

import tensorflow as tf

# 创建一个先进先出队列，指定队列中最多可以保存两个元素，并指定数据类型为整形。
q = tf.FIFOQueue(capacity=2, dtypes=tf.int32, name="queue")

# 使用enqueue_many函数来初始化队列中的元素。
# 和变量初始化类似，在使用队列之前需要明确调用这个初始化过程。
init = q.enqueue_many(([0, 10],))

# 使用dequeue函数将队列中的第一个元素出队列。这个元素的值将被保存在变量x中。
x = q.dequeue()

# 将得到的值加1
y = x + 1

# 将加1后的值再重新加入到队列。
q_inc = q.enqueue(y)

# Tensorflow会话
with tf.Session() as sess:

    # 运行初始化队列的操作。
    init.run()
    for i in range(5):

        # 运行q_inc将执行数据出队列，出队的元素值加1，重新加入队列的整个过程。
        v, _ = sess.run([x, q_inc])

        # 打印出队元素的值。
        print(v)

程序运行结果如下：

队列开始有[0,10]两个元素，第一个出队的是0，加1后再次加入队列，得到的队列为[10,1]；
第二次出队的元素是10，加1之后变为11，11入队之后得到的队列为[1,11]；以此类推。

0
10
1
11
2

TensorFlow中提供了FIFOQueue和RandomShuffleQueue两种队列。在上面的程序中已经展示了如何使用FIFOQueue，它实现的是一个先进先出队列。RandomShuffleQueue会将队列中的元素打乱，每次出队操作得到的是从当前队列所有元素中随机选择一个。在训练神经网络是希望每次使用的训练数据尽量随机，RandomShuffleQueue就提供了这样的功能。

在TensorFlow中，队列不仅仅是一种数据结构，还是一步计算张量（tensor）取值的一个重要机制。比如多个线程可以同时向一个队列中写元素，或是同时读取一个队列中的元素。

TensorFlow提供了tf.Coordinator和tf.QueueRunner两个类来完成多线程协同的功能。

tf.Coordinator主要用于协同多个线程一起停止，并提供了should_stop、request_stop、join三个函数。在起订线程之前，需要先声明一个tf.Coordinator类，并将这个类传入每一个创建的线程中。启动的线程需要一直查询tf.Coordinator类中提供的should_stop函数，当这个函数的返回值为True时，则当前线程也需要退出。每一个启动的线程都可以通过调用request_stop函数来通知其他线程退出。当某一个线程调用request_stop函数之后should_stop函数的返回值将会被设置为True，这样其他的线程就可以同时终止了。

import tensorflow as tf
import numpy as np
import threading
import time

# 线程中需要运行的代码，每个1秒判断是否需要停止并打印自己的ID。
def MyLoop(coord, worker_id):

    # 使用tf.Coordinator类提供的协同工具判断当前线程是否需要停止
    while not coord.should_stop():

        # 随机停止所有的线程
        if np.random.rand() < 0.1:

            # 打印信息
            print("Stopping from id: %d"%(worker_id))

            # 调用request_stop()函数来通知其他的线程停止工作
            coord.request_stop()
        else:

            # 打印信息
            print("working on id: %d"%worker_id)

        # 暂停1秒钟
        time.sleep(1)

# 声明一个tf.train.Coordinator()类来协同多个线程
coord = tf.train.Coordinator()

# 声明创建5个线程
threads = [threading.Thread(target=MyLoop, args=(coord, i)) for i in range(5)]

# 启动所有的线程
for t in threads:
    t.start()

# 等待所有的线程退出
coord.join(threads)

程序运行结果如下：

working on id: 0
working on id: 1
working on id: 2
working on id: 3
working on id: 4
Stopping from id: 0
working on id: 1

也有可能如下：

working on id: 0
working on id: 1
working on id: 2
working on id: 3
working on id: 4
working on id: 0
working on id: 1
working on id: 4
working on id: 2
working on id: 3
Stopping from id: 0

当启动所有的线程之后，每个线程会打印各自的ID，于是前面打印出了它们的ID。然后在暂停1秒后，所有的线程又开始第二遍的打印自己的ID。在这个时候有个线程退出的条件达到了，于是调用了coord.request_stop函数来停止其他的线程。然而在第一次运行的时候，打印Stopping from id： 0 之后，可以看到仍然有线程在输出，这是因为这些线程已经执行完coord.should_stop的判断，于是仍然会输出自己的ID，但是在下一轮判断是否需要停止时将退出线程，于是在打印一次ID之后就不会再有输出了。

tf.QueueRunner主要用于启动多个线程来操作同一个队列，启动的这些线程可以通过上面介绍的tf.Coordinator类来统一管理。

import tensorflow as tf
import numpy as np
import threading
import time

# 声明一个先进先出的队列，队列中最多100个元素，类型为float32。
queue = tf.FIFOQueue(100, tf.float32)

# 定义队列的入队操作。
enqueue_op = queue.enqueue([tf.random_normal([1])])

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

# 将定义过的QueueRunner加入TensorFlow计算图上指定的集合。
# tf.train.add_queue_runner函数没有指定集合，则加入默认集合tf.GraphKeys.QUEUE_RUNNERS。
# 下面的函数就是将刚刚定义的qr加入默认的tf.GraphKeys.QUEUE_RUNNERS集合。
tf.train.add_queue_runner(qr)

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

# TensorFlow会话。
with tf.Session() as sess:

    # 使用tf.train.Coordinator来协同启动的线程。
    coord = tf.train.Coordinator()

    # 使用tf.train.QueueRunner时们需要明确调用tf.train.start_queue_runners来启动所有的线程。
    # 否则因为没有线程运行入队操作，当调用出队操作时，程序会一直等待入队操作被运行。
    # tf.train.start_queue_runners函数默认会启动tf.GraphKeys.QUEUE_RUNNERS集合中所有的QueueRunner。
    # 因为tf.train.start_queue_runners函数和tf.train.add_queue_runner函数会指定同一个集合。
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    # 获取队列中的值。
    for i in range(3):

        # 输出
        print(sess.run(out_tensor)[0])

        # 输出当前队列的数据元素
        print(queue.size().eval())

    # 使用tf.train.Coordinator来停止所有的线程。
    coord.request_stop()
    coord.join(threads)

    # 输出
    print(queue.size().eval())

程序运行结果如下所示：（不同机器运行结果或有不同）

-0.311776
2
-0.483378
3
1.68673
5
6

上面的程序将启动5个线程来执行队列的入队操作，其中每一个线程都是将随机数写入队列，于是在每次运行出队操作时，可以得到一个随机数。