tensorflow中的队列和多线程

在Tensorflow中，队列不仅仅是一种数据结构，还是异步张量取值的重要机制，比如多线程可以同时向一个队列写元素，或者从同一个队列读元素，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，这样其他的线程就可以同时终止了。下面程序展示了如何使用 tf.Coordinator。

#_*_coding:utf-8_*_
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\n" % worker_id)
            # 调用 coord.request_stop() 函数来通知其他线程停止
            coord.request_stop()
        else:
            # 打印当前线程的 ID
            print("Working on id: %d\n" % worker_id)
        # 暂停1 s
        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
Working on id: 0
Working on id: 1
Working on id: 3
Working on id: 4
Working on id: 2
Working on id: 0
Working on id: 1
Stopping from id: 4

首先启动五个线程，一旦有一个线程np.random.rand() < 0.1，会调用coord.request_stop()，其他线程通过coord.should_stop()获悉停止线程的通知后，停止线程。

tf.QueueRunner主要用于启动多个线程来操作同一个队列，启动的这些线程可以通过上面介绍的tf.Coordinator类来统一管理。以下代码展示了如何使用tf.QueueRunner和tf.Coordinator来管理多线程队列操作。

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个线程，每个线程中运行的是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()
 
with tf.Session() as sess:
    # 使用tf.train.Cooridator来协同启动的线程。
    sess.run(tf.local_variables_initializer())
    coord = tf.train.Coordinator()
    # 使用tf.train.QueueRunner时，需要说明调用tf.train.start_queue_runners
    # 来启动所有线程，否则因为没有线程运行入队操作，当调用出队操作时，程序会一直等待入队
    # 操作被运行。tf.train.start_queue_runners函数会默认启动
    # tf.GraphKeys.QUEUE_RUNNERS集合中所有的QueueRunner。因为这个函数只支持启动
    # 指定集合中的QueueRunner，所以一般来说tf.train.add_queue_runer函数和
    # 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)

以上程序启动5个线程来执行入队操作，由于入队列操作qr已通过tf.train.add_queue_runner(qr)加入tf.GraphKeys.QUEUE_RUNNERS集合中，所以执行start_queue_runners会启动线程进行入队列操作，故每次运行出队操作均会得到一个随机数。结果如下：

0.78807044
0.28835386
-1.2145209