如何讓TensorFlow模型運行提速36.8%

在訓練TensorFlow模型的時候，我們傳統的做法是在每個Epoch將數據通過feed_dict導入到session中，即不斷地從Python到C++之間來回切換，這種做法十分不高效。而且，訓練操作與導入數據操作都是屬於同一個主線程，它們屬於同步操作，訓練操作必須等待導入數據操作完成以後纔開始執行。這種做法帶來的就是運行效率底下，我們應該趕緊拋棄，取而代之的是TensorFlow中的線程和隊列。下面，我將通過實例代碼來說明如何使用線程和隊列以及其如何爲程序效率提高36%。

爲了比較，我寫了一段傳統的基於TensorFlow訓練神經網絡的代碼。構建了一個隱含層神經元個數爲128的雙向動態LSTM網絡（BLSTM），輸入數據是400個樣本，每個樣本的維度是[128, 512]，即時間長度爲128，特徵長度爲512，而對應標籤是將這個序列進行二分類，得到0或1。

x和y是標準的輸入輸出，網絡的預測輸出則是在BLSTM頂端加了一個線性MLP，最後通過sigmoid層並輸出交叉熵誤差。在這個程序裏，數據都是通過feed_dict導入的，代碼如下：

#coding:utf-8

import time

import tensorflow as tf
from tensorflow.contrib.rnn import LSTMCell

'''
構建了一個隱含層神經元個數爲128的雙向動態LSTM網絡（BLSTM），
輸入數據是400個樣本，每個樣本的維度是[128, 512]，即時間長度爲128，特徵長度爲512，
而對應標籤是將這個序列進行二分類，得到0或1。
'''

time_length = 128

batch_size = 400

feature_size = 512

hidden_size = 128

# 隨機產生以均值爲0 方差爲1 的[time_length, batch_size, feature_size]數據
x = tf.random_normal([time_length,
    batch_size, feature_size], mean=0, stddev=1)

y = tf.reduce_mean(tf.reduce_sum(x, axis=0), axis=1, keep_dims=True)
y = tf.cast(tf.greater(y, 0), tf.int32)

inputs = tf.placeholder(tf.float32,
     shape=[time_length, batch_size, feature_size])
labels = tf.placeholder(tf.int32, shape=[batch_size, 1])

sequence_length = tf.Variable([time_length]*batch_size, dtype=tf.int32)

cell_fw = LSTMCell(num_units=hidden_size)
cell_bw = LSTMCell(num_units=hidden_size)
outputs, state = tf.nn.bidirectional_dynamic_rnn(
      cell_fw=cell_fw,
      cell_bw=cell_bw,
      inputs=inputs,
      sequence_length=sequence_length,
      dtype=tf.float32,
      time_major=True)
outputs_fw, outputs_bw = outputs
outputs = tf.concat([outputs_fw, outputs_bw], axis=2)
outputs = tf.reduce_mean(outputs, axis=0)
outputs = tf.contrib.layers.fully_connected(
            inputs=outputs,
            num_outputs=1,
            activation_fn=None)


losses_op = tf.nn.sigmoid_cross_entropy_with_logits(None, tf.cast(labels, tf.float32), outputs)

losses_op = tf.reduce_mean(losses_op)

y_pred = tf.cast(tf.greater(outputs, 0), tf.int32)

accuracy = tf.reduce_mean(tf.cast(tf.equal(y_pred, y), tf.float32))

# adam = tf.train.AdamOptimizer(0.001)

train_op = tf.train.AdamOptimizer(0.001).minimize(losses_op, name="train_op")

t1 = time.time()
with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  for i in range(50):
    data_x, data_y = sess.run([x, y])
    _, losses, acc = sess.run([train_op, losses_op, accuracy],
                              feed_dict={inputs: data_x, labels: data_y})
    print('epoch:%d, loss: %f,   accuracy: %f' % (i, losses, acc))

print('time:', (time.time()-t1))

將這段代碼在我電腦的單一GPU上執行（GPU爲TITAN X系列），運行結果如下所示，可以看到，整個程序花費了58.94秒。

 

接下來我使用TensorFlow的queue來優化這個程序，和TensorFlow中的其他對象一樣，queue也只是TensorFlow中的一個節點，其可以將feed_dict操作轉化爲異步計算。從tensorflow.python.ops.data_flow_ops的源碼中可以看到，所有的Queue類都是繼承自QueueBase父類，繼承自QueueBase類的子類有：

RandomShuffleQueue 隨機隊列
FIFOQueue 先入先出
PaddingFIFOQueue 可以自帶padding
PriorityQueue 優先隊列

並且每個類都有以下幾個常用的方法：

enqueue() 即向queue中壓入數據
dequeue() 即從queue中彈出數據
enqueue_many() 即向queue中壓入多個數據
dequeue_many() 即從queue中彈出多個數據

當然，單個這個Queue類的方法好像並沒有什麼吸引人之處，因爲就它自身而言，如果我們啓動了一個queue，它是生產者消費者模型，如果只是在單一線程下面工作，那仍然是無濟於事的，就好比很多個廚師一起做菜，然而卻只有一個竈臺可以利用。因此，要想提高運行效率，我們必須要讓enqueue和dequeue分別處在不同線程上，這個時候就需要用到QueueRunner類和Coordinator類了。

QueueRunner類和Coordinator類的作用是處理隊列的操作、保持同步，並且這些操作都是在不同的線程上面。根據tensorflow.python.training中的源碼可知，QueueRunner需要一個隊列和入隊的操作（可以是很多個操作），然後根據session即可創造出很多個執行入隊操作的線程，然後調用tf.train.add_queue_runner方法即可將queue_runner添加到TensorFlow QUEUE_RUNNERS集合中去，再調用tf.train.start_queue_runner方法即可啓動所有線程。這樣，就可以調用上層API Coordinate來執行線程了。

Coordinator類可以讓多個線程停止，它主要有三個方法：

tf.train.Coordinator.should_stop 確認線程是否應該停止
tf.train.Coordinator.request_stop 要求線程停止
tf.train.Coordinator.join 要求等待線程結束

代碼如下：

import time
import tensorflow as tf
from tensorflow.contrib.rnn import LSTMCell

time_length = 128
batch_size = 400
feature_size = 512

hidden_size = 128

## prepare data
x = tf.random_normal([time_length, batch_size, feature_size], mean=0, stddev=1)

q = tf.FIFOQueue(capacity=4, dtypes=tf.float32)
enqueue_op = q.enqueue(x)
num_threads = 1
qr = tf.train.QueueRunner(q, [enqueue_op] * num_threads)
tf.train.add_queue_runner(qr)
inputs = q.dequeue()
inputs.set_shape(x.get_shape())
y = tf.reduce_mean(tf.reduce_sum(inputs, axis=0), axis=1, keep_dims=True)
labels = tf.cast(tf.greater(y, 0), tf.int32)

## build model
sequence_length = tf.Variable([time_length]*batch_size, dtype=tf.int32)
cell_fw = LSTMCell(num_units=hidden_size)
cell_bw = LSTMCell(num_units=hidden_size)
outputs, state = tf.nn.bidirectional_dynamic_rnn(
      cell_fw=cell_fw,
      cell_bw=cell_bw,
      inputs=inputs,
      sequence_length=sequence_length,
      dtype=tf.float32,
      time_major=True)

outputs_fw, outputs_bw = outputs
outputs = tf.concat([outputs_fw, outputs_bw], axis=2)
outputs = tf.reduce_mean(outputs, axis=0)
outputs = tf.contrib.layers.fully_connected(
            inputs=outputs,
            num_outputs=1,
            activation_fn=None)

losses_op = tf.nn.sigmoid_cross_entropy_with_logits(None, tf.cast(labels, tf.float32), outputs)
losses_op = tf.reduce_mean(losses_op)

y_pred = tf.cast(tf.greater(outputs, 0), tf.int32)
accuracy = tf.reduce_mean(tf.cast(tf.equal(y_pred, labels), tf.float32))
# adam = tf.train.AdamOptimizer(0.001)
train_op = tf.train.AdamOptimizer(0.001).minimize(losses_op, name="train_op")

t1 = time.time()
with tf.Session() as sess:
  sess.run(tf.global_variables_initializer())
  coord = tf.train.Coordinator()
  threads = tf.train.start_queue_runners(coord=coord)
  for i in range(50):
    _, losses, acc = sess.run([train_op, losses_op, accuracy])
    print('epoch:%d, loss: %f' % (i, losses))

  coord.request_stop()
  coord.join(threads)
  print("Time taken: %f" % (time.time() - t1))

運行結果如下圖所示：

可以看到這裏我們同樣的程序，同樣的數據，同樣的GPU，只花費了37.25秒，由此可計算出該方法竟然比傳統的feed_dict做法提速了36.8%!