【TensorFlow】多线程输入数据处理框架

在tensorflow中，队列和变量类似，都是计算图上有状态的节点。其它的计算节点可以修改它们的状态。对于变量，可以通过赋值操作修改变量的取值。对于队列，修改队列状态的操作主要有Enqueue、EnqueueMany和Dequeue：

import tensorflow as tf

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

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

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

y = x+1

q_inc = q.enqueue([y])

with tf.Session() as sess:
	init.run()

	for _ in range(5):
		v, _ = sess.run([x, q_inc])
		print v
	
	# 0,10,1,11,2

tensorflow中提供了FIFOQueue(先进先出)和RandomShuffleQueue(将队列中的元素打乱，每次出队列操作得到的是从当前队列所有元素中随机选择的一个)两种队列。
在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

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:
			coord.request_stop()
		else:
			print worker_id
		time.sleep(1) #暂停1秒

# 声明一个tf.train.Coordinator类来协同多个线程
coord = tf.train.Coordinator()
# 声明创建5个线程
threads = [threading.Thread(target=MyLoop, args=(coord, i, )) for i in xrange(5)]

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

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

当所有线程启动之后，每个线程会打印各自的id，然后暂停1秒之后，所有线程又开始第二遍打印id。。。

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

import tensorflow as tf

queue = tf.FIFOQueue(100, "float")
enqueue_op = queue.enqueue([tf.random_noamal([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
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集合中所有的QueueRunner，因为这个函数只支持启动指定集合中的QueueRunner，所以一般来说tf.train.add_queue_runner函数和tf.train.start_quque_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)

tensorflow提供了tf.train.batch和tf.train.shuffle_batch函数来将单个的样例组织成batch的样式输出。这两个函数都会生成一个队列，队列的入队操作是生成单个样例的方法，而每次出对得到的是一个batch的样例。

import tensorflow as tf

example, label = features['i'], features['j']
batch_size = 3
# 组合样例的队列中最多可以存储的样例个数。队列如果太大，那么需要占用很多内存资源；如果太小，那么出队操作可能会因为没有数据而被阻碍，从而导致训练效率降低。一般来说这个队列的大小会和每一个batch的大小相关
capacity = 1000 + 3 * batch_size

# 使用tf.train.batch函数来组合样例。
example_batch, label_batch = tf.train.batch([example,label], batch_size=batch_size, capacity=capacity)
# 当队列长度等于容量时，tensorflow将暂停入队操作，而只等待元素出队。当元素数小于容量时，tensorflow将自动重新启动入队操作。

# 使用tf.train.shuffle_batch函数来组合样例。
# tf.train.shuffle_batch函数的参数大部分都和tf.train.batch函数相似，但是min_after_dequeue参数是tf.train.shuffle_batch特有的，min_atfer_dequeue参数限制了出队时队列中元素的最小个数。当队列中元素太少时。随机打乱样例的作用就不大了。当出队函数被调用但是队列中元素不够时，出队操作将等待更多的元素入队才会完成。
example_batch, label_batch = tf.train.shuffle_batch([example,label], batch_size=batch_size, capacity=capacity, min_after_dequeue=30)

with tf.Session() as sess:
	tf.initialize_all_variables().run()
	coord = tf.train.Coordinator()
	threads = tf.train.start_queue_runners(sess=sess, coord=coord)

	for i in range(2):
		cur_example_batch, cur_label_batch = sess.run([example_batch,label_batch])
	
	coord.request_stop()
	coord.join(threads)

tf.train.batch和tf.train.shuffle_batch函数除了可以将单个训练数据整理成输入batch，也提供了并行化处理输入数据的方法。tf.train.batch函数和tf.train.shuffle_batch函数并行化的方式一致，通过设置函数中的num_threads参数，可以指定多个线程同时执行入队操作，入队操作就是读取数据以及预处理的过程。当需要多个线程处理不同文件中的样例时，可以使用tf.train.shuffle_batch_join函数，此函数会从输入文件队列中获取不同的文件分配给不同的线程。

参考资料：
TensorFlow：实战Google深度学习框架-中国工信出版集团