TensorFlow队列与多线程

1  tf.Coordinator     

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：

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

#线程中运行的程序，这个程序每隔1秒判断是否需要停止并打印自己的ID
def MyLoop(coord, worker_id):
    #使用tf.Coordinator类提供的协同工具判断当前线程是否需要停止
    while not coord.should_stop():
        #随即停止所有线程
        if np.random.rand()<0.1:
            print ('Stoping from id: %d\n' % (worker_id))
            #调用coord.request_stop()函数来通知其他线程停止
            coord.request_stop()
        else:
            #打印当前线程id
            print ('Working on id: %d\n' % (worker_id))
            #暂停一秒
        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: 1

Working on id: 0
Working on id: 2

Working on id: 3


Working on id: 4

Stoping from id: 1
Working on id: 0

2 tf.QueueRunner

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

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