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)