【技術】TensorFlow官方解讀：如何在多系統和網絡拓撲中構建高性能模型
這個文檔和附帶的腳本詳細介紹了如何構建針對各種系統和網絡拓撲的高性能可拓展模型。這個技術在本文檔中用了一些低級的 Tensorflow Python 基元。在未來，這些技術將被并入高級 API。
輸入管道
性能指南闡述了如何診斷輸入管道可能存在的問題及其***解決方法。在使用大量輸入和每秒更高的采樣處理中我們發現 tf.FIFOQueue 和 tf.train.queue_runner 無法使用當前多個 GPU 生成飽和，例如在使用 AlexNet 訓練 ImageNet 時。這是因為使用了 Python 線程作為底層實現，而 Python 線程的開銷太大了。
我們在腳本中采用的另一種方法是通過 Tensorflow 中的本機并行構建輸入管道。我們的方法主要由如下 3 個階段組成：
- I/O 讀?。簭拇疟P中選擇和讀取圖像文件。
- 圖像處理：將圖像記錄解碼為像素、預處理并生成最小批量。
- CPU 到 GPU 的數據傳輸：將圖像從 CPU 傳輸至 GPU。
通過利用 data_flow_ops.StagingArea，每個階段的主要部分與其他階段并行執行。StagingArea 是一個像隊列(queue)一樣且類似于 tf.FIFOQueue 的運算符。不同之處在于 StagingArea 提供了更簡單的功能且可在 CPU 和 GPU 中與其他階段并行執行。將輸入管道拆分為 3 個獨立并行操作的階段，并且這是可擴展的，充分利用大型的多核環境。本節的余下部分將詳細介紹每個階段以及 data_flow_ops.StagingArea 的使用細節。
并行 I/O 讀取
data_flow_ops.RecordInput 用于磁盤的并行讀取。給定一個代表 TFRecords 的輸入文件列表，RecordInput 可使用后臺線程連續讀取記錄。這些記錄被放置在大型的內部池中，當這個池加載量達到其容量的一半時，會有相應的張量輸出。這個操作有其內部線程，線程由占用最少的 CPU 資源的 I/O 時間主導，這就允許它可與模型的其余部分并行運行。
并行圖像處理
從 RecordInput 讀取圖像后，它們作為張量被傳遞至圖像處理管道。為了更方便解釋圖像處理管道，假設輸入管道的目標是 8 個批量大小為 256(每個 GPU 32 個)GPU。256 個圖像記錄的讀取和處理是獨立并行的。從圖中 256 個 RecordInput 讀操作開始，每個讀取操作后都有一個與之相匹配的圖像預處理操作，這些操作是彼此獨立和并行執行的。這些圖像預處理操作包括諸如圖像解碼、失真和調整大小。
當圖像通過預處理器后，它們被聯接成 8 個大小為 32 的張量。為了達到這一目的，使用了 tf.parallel_stack，而不是 tf.concat ，目的作為單一操作被實現，且在將它們聯結在一起之前需要所有輸入準備就緒。tf.parallel_stack 將未初始化的張量作為輸出，并且在有張量輸入時，每個輸入的張量被寫入輸出張量的指定部分。
當所有的張量完成輸入時，輸出張量在圖中傳遞。這有效隱藏了由于產生所有輸入張量的長尾(long tail)而導致的內存延遲。
并行從 CPU 到 GPU 的數據傳輸
繼續假設目標是批量大小為 256(每個 GPU 32 個)8 個 GPU，一旦輸入圖像被處理完并被 CPU 聯接后，我們將得到 8 個批量大小為 32 的張量。Tensorflow 可以使一個設備的張量直接用在任何其他設備上。為使張量在任何設備中可用，Tensorflow 插入了隱式副本。在張量被實際使用之前，會在設備之間調度副本運行。一旦副本無法按時完成運行，需要這些張量的計算將會停止并且導致性能下降。
在此實現中，data_flow_ops.StagingArea 用于明確排定并行副本。最終的結果是當 GPU 上的計算開始時，所有張量已可用。
軟件管道
由于所有的階段都可以在不同的處理器下運行，在它們之間使用 data_flow_ops.StagingArea 可使其并行運行。StagingArea 是一個與 tf.FIFOQueue 相似且像隊列(queue)一樣的運算符，tf.FIFOQueue 提供更簡單的功能可在 CPU 和 GPU 中被執行。
在模型開始運行所有的階段之前，輸入管道階段將被預熱，以將其間的分段緩存區置于一組數據之間。在每個運行階段中，開始時從分段緩沖區中讀取一組數據，并在***將該組數據推送。
例如有 A、B、C 三個階段，這之間就有兩個分段區域 S1 和 S2。在預熱時，我們運行：
- Warm up:
- Step 1: A0
- Step 2: A1 B0
- Actual execution:
- Step 3: A2 B1 C0
- Step 4: A3 B2 C1
- Step 5: A4 B3 C2
預熱結束之后，S1 和 S2 各有一組數據。對于實際執行的每個步驟，會計算一組來自分段區域的數據，同時分段區域會添加一組新數據。
此方案的好處是：
- 所有的階段都是非阻塞的，因為預熱后分段區域總會有一組數據存在。
- 每個階段都可以并行處理，因為它們可以立即啟動。
- 分段緩存區具有固定的內存開銷，并至多有一組額外的數據。
- 運行一個步驟的所有階段只需要調用 singlesession.run()，這使得分析和調試更加容易。
構建高性能模型的***實踐
以下收集的是一些額外的***實踐，可以改善模型性能，增加模型靈活性。
使用 NHWC 和 NCHW 建模
CNN 使用的絕大多數 Tensorflow 操作都支持 NHWC 和 NCHW 數據格式。在 GPU 中，NCHW 更快;但是在 CPU 中，NHWC 只是偶爾更快。
構建一個支持日期格式的模型可增加其靈活性，能夠在任何平臺上良好運行?；鶞誓_本是為了支持 NCHW 和 NHWC 而編寫的。使用 GPU 訓練模型時會經常用到 NCHW。NHWC 在 CPU 中有時速度更快。在 GPU 中可以使用 NCHW 對一個靈活的模型進行訓練，在 CPU 中使用 NHWC 進行推理，并從訓練中獲得合適的權重參數。
使用融合的批處理歸一化
Tensorflow 中默認的批處理歸一化被實現為復合操作，這是很通用的做法，但是其性能不好。融合的批處理歸一化是一種替代選擇，其在 GPU 中能取得更好的性能。如下是用 tf.contrib.layers.batch_norm 實現融合批處理歸一化的一個實例：
- bn = tf.contrib.layers.batch_norm(
- input_layer, fused=True, data_format='NCHW'
- scope=scope)
變量分布和梯度聚合
訓練期間，訓練的變量值通過聚合的梯度和增量進行更新。在基準腳本中，展示了通過使用靈活和通用的 Tensorflow 原語，我們可以構建各種各樣的高性能分布和聚合方案。
在基準腳本中包括 3 個變量分布和聚合的例子：
- 參數服務器，訓練模型的每個副本都從參數服務器中讀取變量并獨立更新變量。當每個模型需要變量時，它們將被復制到由 Tensorflow 運行時添加的標準隱式副本中。示例腳本介紹了使用此方法如何進行本地訓練、分布式同步訓練和分布式異步訓練。
- 拷貝，在每個 GPU 上放置每個訓練變量相同的副本，在變量數據立即可用時，正向計算和反向計算立即開始。所有 GPU 中的梯度都會被累加，累加的總和應用于每個 GPU 變量副本，以使其保持同步。
- 分布式復制，將每個 GPU 中的訓練參數副本與參數服務器上的主副本放置在一起，在變量數據可用時，正向計算和反向計算立即開始。一臺服務器上每個 GPU 的梯度會被累加，然后每個服務器中聚合的梯度會被應用到主副本中。當所有的模塊都執行此操作后，每個模塊都將從主副本中更新變量副本。
以下是有關每種方法的其他細節。
參數服務器變量
在 Tensorflow 模型中管理變量的最常見方式是參數服務器模式。
在分布式系統中，每個工作器(worker)進程運行相同的模型，參數服務器處理其自有的變量主副本。當一個工作器需要一個來自參數服務器的變量時，它可從其中直接引用。Tensorflow 在運行時會將隱式副本添加到圖形中，這使得在需要它的計算設備上變量值可用。當在工作器上計算梯度時，這個梯度會被傳輸到擁有特定變量的參數服務器中，而相應的優化器被用于更新變量。
以下是一些提高吞吐量的技術：
- 為了使負載平衡，這些變量根據其大小在參數服務器之間傳輸。
- 當每個工作器有多個 GPU 時，累加每個 GPU 的梯度，并把這個單一的聚合梯度發送到參數服務器。這將降低網絡帶寬，減少參數服務器的工作量。
為了協調工作器，常常采用異步更新模式，其中每個工作器更新變量的主副本，而不與其他工作器同步。在我們的模型中，我們展示了在工作器中引入同步機制是非常容易的，所以在下一步開始之前所有的工作器必須完成更新。
這個參數服務器方法同樣可以應用在本地訓練中，在這種情況下，它們不是在參數服務器之間傳播變量的主副本，而是在 CPU 上或分布在可用的 GPU 上。
由于該設置的簡單性，這種架構在社區中獲得廣泛的推廣。
通過傳遞參數 variable_update=parameter_server，也可以在腳本中使用此模式。
變量復制
在這種設計中，服務器中的每個 GPU 都有自己的變量副本。通過將完全聚合的梯度應用于變量的每個 GPU 副本，使得這些值在 GPU 之間保持同步。
因為變量和數據在訓練的初始階段就準備好了，所以訓練的前向計算可以立即開始。聚合各個設備的梯度以得到一個完全聚合的梯度，并將該梯度應用到每個本地副本中。
服務器間的梯度聚合可通過不同的方法實現：
- 使用 Tensorflow 標準操作在單個設備上(CPU 或 GPU)累加整和，然后將其拷貝回所有的 GPU。
- 使用英偉達 NCCL，這個將在下面的 NCCL 章節闡述。
分布式訓練中的變量復制
上述變量復制的方法可擴展到分布式訓練中。一種類似的方法是：完全地聚合集群中的梯度，并將它們應用于每個本地副本。這種方法在未來版本的腳本中可能會出現，但是當前的腳本采用不同的方法。描述如下。
在這一模式中，除了變量的每一個 GPU 副本之外，主副本被存儲在參數服務器之中。借助這一復制模式，可使用變量的本地副本立刻開始訓練。
隨著權重的梯度可用，它們會被送回至參數服務器，并所有的本地副本都會被更新：
- 同一個工作器中把 GPU 所有的梯度聚合在一起。
- 將來自各個工作器的聚合梯度發送至自帶變量的參數服務器中，其中使用特殊的優化器來更新變量的主副本。
- 每個工作器從主副本中更新變量的本地副本。在示例模型中，這是在一個擁有交叉副本的負載中在等待所有的模塊完成變量更新后進行的，并且只有在負載被所有副本釋放以后才能獲取新的變量。一旦所有的變量完成復制，這就標志著一個訓練步驟的完成，和下一個訓練步驟的開始。
盡管這些聽起來與參數服務器的標準用法很相似，但是其性能在很多案例中表現更佳。這很大程度因為計算沒有任何延遲，早期梯度的大部分復制延遲可被稍后的計算層隱藏。
通過傳遞參數 variable_update=distributed_replicated 可以在腳本中使用該模式。
NCCL
為了在同一臺主機的不同 GPU 上傳播變量和聚合梯度，我們可以使用 Tensorflow 默認的隱式復制機制。
然而，我們也可以選擇 NCCL(tf.contrib.nccl)。NCCL 是英偉達的一個庫，可以跨不同的 GPU 實現數據的高效傳輸和聚合。它在每個 GPU 上分配一個協作內核，這個內核知道如何***地利用底層硬件拓撲結構，并使用單個 SM 的 GPU。
通過實驗證明，盡管 NCCL 通常會加速數據的聚合，但并不一定會加速訓練。我們的假設是：隱式副本基本是不耗時的，因為它們本在 GPU 上復制引擎，只要它的延遲可以被主計算本身隱藏起來，那么。雖然 NCCL 可以更快地傳輸數據，但是它需要一個 SM，并且給底層的 L2 緩存增加了更多的壓力。我們的研究結果表明，在 8 個 GPU 的條件下，NCCL 表現出了更優異的性能;但是如果 GPU 更少的情況下，隱式副本通常會有更好的表現。
分段變量
我們進一步介紹一種分段變量模式，我們使用分段區域來進行變量讀取和更新。與輸入管道中的軟件流水線類似，這可以隱藏數據拷貝的延遲。如果計算所花的時間比復制和聚合的時間更長，那么可以認為復制本身是不耗時的。
這種方法的缺點是所有的權重都來自之前的訓練步驟，所以這是一個不同于 SGD 的算法，但是通過調整學習率和其他超參數，還是有可能提高收斂性。
腳本的執行
這一節將列出執行主腳本的核心命令行參數和一些基本示例(tf_cnn_benchmarks.py)
注意：tf_cnn_benchmarks.py 使用的配置文件 force_gpu_compatible 是在 Tensorflow 1.1 版本之后引入的，直到 1.2 版本發布才建議從源頭建立。
主要的命令行參數
- model：使用的模型有 resnet50、inception3、vgg16 和 alexnet。
- num_gpus：這里指所用 GPU 的數量。
- data_dir：數據處理的路徑，如果沒有被設置，那么將會使用合成數據。為了使用 Imagenet 數據，可把這些指示 (https://github.com/tensorflow/tensorflow/blob/master/tensorflow_models/inception#getting-started) 作為起點。
- batch_size：每個 GPU 的批量大小。
- variable_update：管理變量的方法：parameter_server 、replicated、distributed_replicated、independent。
- local_parameter_device：作為參數服務器使用的設備：CPU 或者 GPU。
單個實例
- # VGG16 training ImageNet with 8 GPUs using arguments that optimize for
- # Google Compute Engine.
- python tf_cnn_benchmarks.py --local_parameter_device=cpu --num_gpus=8 \
- --batch_size=32 --model=vgg16 --data_dir=/home/ubuntu/imagenet/train \
- --variable_update=parameter_server --nodistortions
- # VGG16 training synthetic ImageNet data with 8 GPUs using arguments that
- # optimize for the NVIDIA DGX-1.
- python tf_cnn_benchmarks.py --local_parameter_device=gpu --num_gpus=8 \
- --batch_size=64 --model=vgg16 --variable_update=replicated --use_nccl=True
- # VGG16 training ImageNet data with 8 GPUs using arguments that optimize for
- # Amazon EC2.
- python tf_cnn_benchmarks.py --local_parameter_device=gpu --num_gpus=8 \
- --batch_size=64 --model=vgg16 --variable_update=parameter_server
- # ResNet-50 training ImageNet data with 8 GPUs using arguments that optimize for
- # Amazon EC2.
- python tf_cnn_benchmarks.py --local_parameter_device=gpu --num_gpus=8 \
- --batch_size=64 --model=resnet50 --variable_update=replicated --use_nccl=False
分布式命令行參數
1)ps_hosts：在<host>:port 的格式中(比如 10.0.0.2:50000)，逗號分隔的主機列表用做參數服務器。
2)worker_hosts：(比如 10.0.0.2:50001)，逗號分隔的主機列表用作工作器，在<host>:port 的格式中。
3)task_index：正在啟動的 ps_host 或 worker_hosts 列表中的主機索引。
4)job_name：工作的類別，例如 ps 或者 worker。
分布式實例
如下是在兩個主機(host_0 (10.0.0.1) 和 host_1 (10.0.0.2))上訓練 ResNet-50 的實例，這個例子使用的是合成數據，如果要使用真實數據請傳遞 data_dir 參數。# Run the following commands on host_0 (10.0.0.1):
- python tf_cnn_benchmarks.py --local_parameter_device=gpu --num_gpus=8 \
- --batch_size=64 --model=resnet50 --variable_update=distributed_replicated \
- --job_name=worker --ps_hosts=10.0.0.1:50000,10.0.0.2:50000 \
- --worker_hosts=10.0.0.1:50001,10.0.0.2:50001 --task_index=0
- python tf_cnn_benchmarks.py --local_parameter_device=gpu --num_gpus=8 \
- --batch_size=64 --model=resnet50 --variable_update=distributed_replicated \
- --job_name=ps --ps_hosts=10.0.0.1:50000,10.0.0.2:50000 \
- --worker_hosts=10.0.0.1:50001,10.0.0.2:50001 --task_index=0
- # Run the following commands on host_1 (10.0.0.2):
- python tf_cnn_benchmarks.py --local_parameter_device=gpu --num_gpus=8 \
- --batch_size=64 --model=resnet50 --variable_update=distributed_replicated \
- --job_name=worker --ps_hosts=10.0.0.1:50000,10.0.0.2:50000 \
- --worker_hosts=10.0.0.1:50001,10.0.0.2:50001 --task_index=1
- python tf_cnn_benchmarks.py --local_parameter_device=gpu --num_gpus=8 \
- --batch_size=64 --model=resnet50 --variable_update=distributed_replicated \
- --job_name=ps --ps_hosts=10.0.0.1:50000,10.0.0.2:50000 \
- --worker_hosts=10.0.0.1:50001,10.0.0.2:50001 --task_index=1
