Tensorflow并行計算：多核(multicore)，多線程(multi-thread)，圖分割(Graph Partition)

2019-11-08 03:00:22

字體：大中小

來源：轉載

供稿：網友

Github下載完整代碼

https://github.com/rockingdingo/tensorflow-tutorial/tree/master/mnist

簡介

利用tensorflow訓練深度神經網絡模型需要消耗很長時間，因為并行化計算就為提升運行速度提供了重要思路。Tensorflow提供了多種方法來使程序的并行運行，在使用這些方法時需要考慮的問題有：選取的計算設備是CPU還是GPU，每個CPU多少核的資源并行計算，構建圖Graph時消耗資源如何分配等等問題。下面我們以linux多核CPU的環境為例介紹幾種常見方法來提升你的tensorflow程序的運行速度。

一、多核并行：CPU多核的運算和資源調用

在Tensorflow程序中，我們會經常看到”with tf.device("/cpu:0"): “ 這個語句。單獨使用這個語句，而不做其他限制，實際上默認tensorflow程序占用所有可以使用的內存資源和CPU核，比如如果你的linux服務器是8核CPU，那么該程序會迅速占用可以使用的任意CPU，使用接近100%，最終結果就是影響整臺服務器的其他程序。因此我們會想到需要限制使用的CPU核的個數和資源。

在構建tf.session() 變量時，可以通過傳入tf.ConfigPRoto() 參數來改變一個tensorflow的session會話所使用的CPU核的個數以及線程數等等。

代碼1

config = tf.ConfigProto(device_count={"CPU": 4}, # limit to num_cpu_core CPU usage inter_op_parallelism_threads = 1, intra_op_parallelism_threads = 1, log_device_placement=True)with tf.Session(config = config) as sess: # To Do

例如，在上面代碼中我們通過 “device_count={"CPU": 4}” 參數來構建一個ConfigProto() 類，傳入tf.Session()來使每個會話分配相應的資源，這里我們給tensorflow程序共分配了4個CPU core。

實例

我們對tensorflow tutorial中的MNIST的CNN模型convolutional.py進行適當的修改，使得他限制在4個CPU core上運行，得到convolutional_multicore.py：運行結果，平均每個step計算一個batch的時間為611 ms左右：

圖1: 4 CPU core Single Thread

二、多線程，設置Multi-threads

在進行tf.ConfigProto()初始化時，我們也可以通過設置intra_op_parallelism_threads參數和inter_op_parallelism_threads參數，來控制每個操作符op并行計算的線程個數。二者的區別在于:

intra_op_parallelism_threads 控制運算符op內部的并行當運算符op為單一運算符，并且內部可以實現并行時，如矩陣乘法，reduce_sum之類的操作，可以通過設置intra_op_parallelism_threads 參數來并行, intra代表內部。inter_op_parallelism_threads 控制多個運算符op之間的并行計算當有多個運算符op，并且他們之間比較獨立，運算符和運算符之間沒有直接的路徑Path相連。Tensorflow會嘗試并行地計算他們，使用由inter_op_parallelism_threads參數來控制數量的一個線程池。

代碼2

我們還是修改上面的MNIST卷積的例子 convolutional_multicore.py，將參數如下:

config = tf.ConfigProto(device_count={"CPU": 4}, # limit to num_cpu_core CPU usage inter_op_parallelism_threads = 1, intra_op_parallelism_threads = 4, log_device_placement=True)with tf.Session(config = config) as sess: # To Do

實例

實例比較，下面我們比較線程數為2和4，平均每個batch的運行時間：

當參數為intra_op_parallelism_threads = 2時, 每個step的平均運行時間從610ms降低到380ms。

當參數為intra_op_parallelism_threads = 4時, 每個step的平均運行時間從610ms降低到230ms。

總結，在固定CPUcore的資源限制下，通過合理設置線程thread個數可以明顯提升tensorflow程序運行速度。

參考 tensorflow/core/protobuf/config.proto 實現

https://github.com/tensorflow/tensorflow/blob/26b4dfa65d360f2793ad75083c797d57f8661b93/tensorflow/core/protobuf/config.proto#L165

三、分割圖模型Graph：將Tensorflow的圖運算分配到不同計算單元

有時我們構建的深度網絡的結構十分復雜，會出現這種情況：多個CPU core同時運行時，有的核比較空閑，有的核使用率卻達到100%的情況。我們需要盡量避免這種運算符計算不均衡的情況。這時，如果我們將Graph拆分為多個部分，將每個部分（如每一層網絡結構）指定到不同的CPU 核上運算，優化計算量的分配，可以使運算速度得到提升。

一個很直觀的設計就是按照不同的層來劃分，把運算量大的Layer分配單獨的CPU，把運算量小的Layer合并分配到同一個CPU core上。

下面是我們做的一個測試，還是tensorflow官網上的 convolutional.py 例子改寫，將不同層分配到不同的CPU device上，優化了計算資源，使得程序的速度得以提升，例子為convolutional_graph_partitioned.py。

聲明了device_id全局變量記錄已經使用的CPU的ID；調用next_device() 函數返回下一個可用的CPU device id, 如果有可用的則分配并使全局變量device_id 加1, 最終獲得的可用的device_id 不會超過在 FLAGS.num_cpu_core中定義的核的總個數。在model()函數構建圖的過程中，通過with tf.device(next_device()): 語句，來將當成的Conv, Pool等運算符分配到單獨的CPU上。最終結果為每個batch 平均時間229 ms。

代碼3

device_id = -1 # Global Variable Counter for device_id used def next_device(use_cpu = True):    ''' See if there is available next device;        Args: use_cpu, global device_id        Return: new device id    '''  global device_id  if (use_cpu):    if ((device_id + 1) < FLAGS.num_cpu_core):      device_id += 1    device = '/cpu:%d' % device_id  else:    if ((device_id + 1) < FLAGS.num_gpu_core):      device_id += 1    device = '/gpu:%d' % device_id  return device   with tf.device(next_device()):  # To Do Insert Your Code  conv = ...  pool = ...
延伸閱讀
深語人工智能-技術博客: http://www.deepnlp.org/blog/tensorflow-parallelism/PyPI deepnlp: Deep Learning NLP Pipeline implemented on Tensorflowhttps://pypi.python.org/pypi/deepnlpTensorflow官網的卷積神經網絡MNIST的例子convolutional.py多線程convolutional_multithread.py分割Graph到不同Device上 convolutional_graph_partitioned.py