理解TensorFlow的Queue

巾梵關注

2017.05.17 21:37* 字數 1067 閱讀 9817評論 1喜歡 14讚賞 2

這篇文章來說說TensorFlow裏與Queue有關的概念和用法。

其實概念只有三個：

Queue是TF隊列和緩存機制的實現
QueueRunner是TF中對操作Queue的線程的封裝
Coordinator是TF中用來協調線程運行的工具

雖然它們經常同時出現，但這三樣東西在TensorFlow裏面是可以單獨使用的，不妨先分開來看待。

1. Queue

根據實現的方式不同，分成具體的幾種類型，例如：

tf.FIFOQueue 按入列順序出列的隊列
tf.RandomShuffleQueue 隨機順序出列的隊列
tf.PaddingFIFOQueue 以固定長度批量出列的隊列
tf.PriorityQueue 帶優先級出列的隊列
… …

這些類型的Queue除了自身的性質不太一樣外，創建、使用的方法基本是相同的。

創建函數的參數：

tf.FIFOQueue(capacity, dtypes, shapes=None, names=None ...)

Queue主要包含入列（enqueue）和出列（dequeue）兩個操作。enqueue操作返回計算圖中的一個Operation節點，dequeue操作返回一個Tensor值。Tensor在創建時同樣只是一個定義（或稱爲“聲明”），需要放在Session中運行才能獲得真正的數值。下面是一個單獨使用Queue的例子：

import tensorflow as tf
tf.InteractiveSession()

q = tf.FIFOQueue(2, "float")
init = q.enqueue_many(([0,0],))

x = q.dequeue()
y = x+1
q_inc = q.enqueue([y])

init.run()
q_inc.run()
q_inc.run()
q_inc.run()
x.eval()  # 返回1
x.eval()  # 返回2
x.eval()  # 卡住

注意，如果一次性入列超過Queue Size的數據，enqueue操作會卡住，直到有數據（被其他線程）從隊列取出。對一個已經取空的隊列使用dequeue操作也會卡住，直到有新的數據（從其他線程）寫入。

2. QueueRunner

Tensorflow的計算主要在使用CPU/GPU和內存，而數據讀取涉及磁盤操作，速度遠低於前者操作。因此通常會使用多個線程讀取數據，然後使用一個線程消費數據。QueueRunner就是來管理這些讀寫隊列的線程的。

QueueRunner需要與Queue一起使用（這名字已經註定了它和Queue脫不開干係），但並不一定必須使用Coordinator。看下面這個例子：

import tensorflow as tf  
import sys  
q = tf.FIFOQueue(10, "float")  
counter = tf.Variable(0.0)  #計數器
# 給計數器加一
increment_op = tf.assign_add(counter, 1.0)
# 將計數器加入隊列
enqueue_op = q.enqueue(counter)

# 創建QueueRunner
# 用多個線程向隊列添加數據
# 這裏實際創建了4個線程，兩個增加計數，兩個執行入隊
qr = tf.train.QueueRunner(q, enqueue_ops=[increment_op, enqueue_op] * 2)

# 主線程
sess = tf.InteractiveSession()
tf.global_variables_initializer().run()
# 啓動入隊線程
qr.create_threads(sess, start=True)
for i in range(20):
    print (sess.run(q.dequeue()))

增加計數的進程會不停的後臺運行，執行入隊的進程會先執行10次（因爲隊列長度只有10），然後主線程開始消費數據，當一部分數據消費被後，入隊的進程又會開始執行。最終主線程消費完20個數據後停止，但其他線程繼續運行，程序不會結束。

3. Coordinator

Coordinator是個用來保存線程組運行狀態的協調器對象，它和TensorFlow的Queue沒有必然關係，是可以單獨和Python線程使用的。例如：

import tensorflow as tf
import threading, time

# 子線程函數
def loop(coord, id):
    t = 0
    while not coord.should_stop():
        print(id)
        time.sleep(1)
        t += 1
        # 只有1號線程調用request_stop方法
        if (t >= 2 and id == 1):
            coord.request_stop()

# 主線程
coord = tf.train.Coordinator()
# 使用Python API創建10個線程
threads = [threading.Thread(target=loop, args=(coord, i)) for i in range(10)]

# 啓動所有線程，並等待線程結束
for t in threads: t.start()
coord.join(threads)

將這個程序運行起來，會發現所有的子線程執行完兩個週期後都會停止，主線程會等待所有子線程都停止後結束，從而使整個程序結束。由此可見，只要有任何一個線程調用了Coordinator的request_stop方法，所有的線程都可以通過should_stop方法感知並停止當前線程。

將QueueRunner和Coordinator一起使用，實際上就是封裝了這個判斷操作，從而使任何一個現成出現異常時，能夠正常結束整個程序，同時主線程也可以直接調用request_stop方法來停止所有子線程的執行。

4. 在一起

在TensorFlow中用Queue的經典模式有兩種，都是配合了QueueRunner和Coordinator一起使用的。

第一種，顯式的創建QueueRunner，然後調用它的create_threads方法啓動線程。例如下面這段代碼：

import tensorflow as tf

# 1000個4維輸入向量，每個數取值爲1-10之間的隨機數
data = 10 * np.random.randn(1000, 4) + 1
# 1000個隨機的目標值，值爲0或1
target = np.random.randint(0, 2, size=1000)

# 創建Queue，隊列中每一項包含一個輸入數據和相應的目標值
queue = tf.FIFOQueue(capacity=50, dtypes=[tf.float32, tf.int32], shapes=[[4], []])

# 批量入列數據（這是一個Operation）
enqueue_op = queue.enqueue_many([data, target])
# 出列數據（這是一個Tensor定義）
data_sample, label_sample = queue.dequeue()

# 創建包含4個線程的QueueRunner
qr = tf.train.QueueRunner(queue, [enqueue_op] * 4)

with tf.Session() as sess:
    # 創建Coordinator
    coord = tf.train.Coordinator()
    # 啓動QueueRunner管理的線程
    enqueue_threads = qr.create_threads(sess, coord=coord, start=True)
    # 主線程，消費100個數據
    for step in range(100):
        if coord.should_stop():
            break
        data_batch, label_batch = sess.run([data_sample, label_sample])
    # 主線程計算完成，停止所有采集數據的進程
    coord.request_stop()
    coord.join(enqueue_threads)

第二種，使用全局的start_queue_runners方法啓動線程。

import tensorflow as tf

# 同時打開多個文件，顯示創建Queue，同時隱含了QueueRunner的創建
filename_queue = tf.train.string_input_producer(["data1.csv","data2.csv"])
reader = tf.TextLineReader(skip_header_lines=1)
# Tensorflow的Reader對象可以直接接受一個Queue作爲輸入
key, value = reader.read(filename_queue)

with tf.Session() as sess:
    coord = tf.train.Coordinator()
    # 啓動計算圖中所有的隊列線程
    threads = tf.train.start_queue_runners(coord=coord)
    # 主線程，消費100個數據
    for _ in range(100):
        features, labels = sess.run([data_batch, label_batch])
    # 主線程計算完成，停止所有采集數據的進程
    coord.request_stop()
    coord.join(threads)