本文翻譯于博客Speed up TensorFlow Inference on GPUs with TensorRT，這篇博客介紹了如何使用TensorRT加速TensorFlow模型的推理速度，作者為：

• Siddharth Sharma?—?Technical Product Marketing Manager, NVidia
• Sami Kama?—?Deep Learning Developer Technologist, NVidia
• Julie Bernauer?—?Pursuit Engineering Solution Architect, NVidia
• Laurence Moroney?—?Developer Advocate, Google

概述

TensorFlow是當今最受歡迎的深度學習框架，在全球有著數(shù)十萬名用戶。NVIDIA?TensorRT?是一個深度學習平臺，用于優(yōu)化神經(jīng)網(wǎng)絡模型，加快數(shù)據(jù)中心，嵌入式芯片和汽車設備中運行的GPU加速平臺的推理速度。我們對TensorFlow與TensorRT的集成感到興奮，這似乎很自然，特別是NVIDIA提供的平臺非常適合用于加速TensorFlow。這使得TensorFlow用戶在使用TensorRT時具有極高的推理性能和近乎透明的工作流程。

圖1 tensor通過優(yōu)化訓練好的神經(jīng)網(wǎng)絡模型來生成部署就緒的推理引擎

TensorRT對神經(jīng)網(wǎng)絡圖執(zhí)行幾項重要的轉換和優(yōu)化（圖2）。首先，消除具有未使用的輸出層以避免不必要的計算。接下來，在可能的情況下，對卷積，偏置和ReLU層進行融合以形成單個層。另一種轉換是水平層融合或層聚合，以及聚合層到它們各自輸出的劃分。水平層融合通過對使用相同源tensor的層進行組合，并應用具有相似參數(shù)的相同操作來提高性能。請注意，這些圖優(yōu)化操作不會更改計算圖中的基礎計算：相反，它們會對計算圖進行重新構建，使其可以更快，更有效地進行推理。

圖2 （a）一個具有多個卷積層和激活層的卷積神經(jīng)網(wǎng)絡。（b）TensorRT通過融合垂直和水平層以及層消除操作簡化了GoogLeNet Inception模塊圖，減少了計算和內存開銷。

如果您已經(jīng)將TensorRT與TensorFlow模型一起使用，那么您應該知道，在過去的版本中，使用TensorRT優(yōu)化需要導出訓練好的TensorFlow計算圖。您還需要手動導入某些不受支持的TensorFlow圖層，然后在TensorRT中運行完整圖形。在大多數(shù)情況下，您再也不用這樣做了。在新的工作流中，您可以通過使用簡單的API在TensorFlow中應用TensorRT強大的FP16和INT8優(yōu)化?，F(xiàn)有的TensorFlow程序只需要幾行新代碼就可以使用這些優(yōu)化。

在ResNet-50的基準測試鐘，TensorRT將TensorFlow推斷速度提高了8倍。這些性能改進僅需幾行額外代碼，并可與TensorFlow 1.7及更高版本一起使用。在本文中，我們將介紹新的工作流程和API，以幫助您開始使用它。

在TensorFlow graphs中應用tensorRT優(yōu)化

如圖3所示，在TensorFlow推理工作流中加入tensorRT優(yōu)化需要一個額外操作，在這個額外操作（使用綠色進行高亮）中，TensorRT根據(jù)frozen TensorFlow graph構建優(yōu)化后的推理圖。

圖3 在TensorFlow中使用TensorRT時的工作流程圖

為了完成優(yōu)化操作，TensorRT對frozen TensorFlow graph進行解析，選出圖中可以進行優(yōu)化的子圖。之后，TensorRT對子圖進行優(yōu)化，并且在原TensorFlow graph中將需要優(yōu)化的子圖替換為TensorRT節(jié)點，圖中其余部分保持不變。在推理過程中，TensorFlow將調用TensorRT來執(zhí)行優(yōu)化后的TensorRT節(jié)點。通過這種方法，開發(fā)人員可以繼續(xù)使用靈活的TensorFlow功能集和TensorRT優(yōu)化。

舉個栗子??，假設有個計算圖由A,B,C三個子圖組成，TensorRT對子圖B進行了優(yōu)化，之后將子圖B替換為TensorRT節(jié)點。在推理過程中，TensorFlow執(zhí)行了子圖A，之后調用TensorFlow執(zhí)行了子圖B，最后使用TensorFlow執(zhí)行了子圖C。

TensorRT優(yōu)化了TensorFlow圖中可能的最大子圖。子圖中的計算越多，從TensorRT獲得的收益越大。您希望對大部分計算圖進行優(yōu)化，并使用最少數(shù)量的TensorRT節(jié)點替換，以獲得最佳性能。根據(jù)計算圖中的操作，優(yōu)化后的計算圖可能包含多個TensorRT節(jié)點。使用TensorFlow API，您可以指定子圖中節(jié)點的最小數(shù)量，以便將其轉換為TensorRT節(jié)點。任何小于指定設定節(jié)點數(shù)的子圖都不會轉換為TensorRT引擎，即使其與TensorRT兼容。這對于包含由不兼容節(jié)點分隔的小兼容子圖的模型非常有用，從而導致微小的TensorRT引擎。

讓我們看看如何更詳細地實現(xiàn)工作流。

使用新的 TensorFlow APIs

新的TensorFlow API支持使用幾行新代碼直接實現(xiàn)TensorRT優(yōu)化。首先，指定分配給TensorFlow的GPU內存比例，剩余內存將用于TensorRT引擎。這可以通過使用GPUOptions函數(shù)的新參數(shù)per_process_gpu_memory_fraction來完成。需要在TensorFlow-TensorRT進程第一次啟動時設置此參數(shù)。例如，將per_process_gpu_memory_fraction設置為0.67會為TensorFlow分配67％的GPU內存，為TensorRT引擎分配剩余三分之一的GPU內存。

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = 
                0 < memory_for_TensorFlow < 1)

下一步是使用tensorRT對TensorFlow計算圖進行分析，進行優(yōu)化，并將需要優(yōu)化的子圖替換為tensotRT節(jié)點。您可以使用新函數(shù)create_inference_graph來對frozen TensorFlow graph進行tensorRT優(yōu)化，這個函數(shù)的輸入為frozen TensorFlow graph，返回含有tensorRT節(jié)點的優(yōu)化圖，如下面的代碼片段所示：

trt_graph = trt.create_inference_graph(
   input_graph_def = frozen_graph_def,
   outputs = output_node_name,
   max_batch_size=batch_size,
   max_workspace_size_bytes=workspace_size,
   precision_mode=precision,
   minimum_segment_size=3)

讓我們看看這個函數(shù)的參數(shù)：

• input_graph_def：frozen TensorFlow graph
• outputs：輸出節(jié)點的tensor name，例如：[“resnet_v1_50/predictions/Reshape_1”]
• max_batch_size：整數(shù)，輸入batch的大小，例如：16
• max_workspace_size_bytes：整數(shù)，tensorRT可用的最大內存
• precision_mode：string，可選值為：“FP32”, “FP16” 或者 “INT8”
• minimum_segment_size：整數(shù)（默認為3），控制要創(chuàng)建的TensorRT引擎的子圖中的最小節(jié)點數(shù)，如果子圖中的節(jié)點數(shù)小于這個值，那么這個子圖將不會被替換為tensorRT節(jié)點

per_process_gpu_memory_fraction和per_process_gpu_memory_fraction兩個參數(shù)應該被一起使用，對TensorFlow和TensorRT的內存分配進行控制，進而獲得應用程序的最佳性能。為了最大化推理性能，您可以分配給tensorRT多一點內存，將剩余的內存分配給TensorFlow。例如：在一塊內存為12G的GPU中，如果您將per_process_gpu_memory_fraction設置為( 12–4 ) / 12 = 0.67，那么參數(shù)max_workspace_size_bytes應設置為4000000000，為TensorRT引擎收集大約4G的內存。同樣，最佳的內存分割方案是依賴于應用程序的，可能需要一些迭代來尋找最佳劃分。

使用TensorBoard對優(yōu)化后的計算圖進行可視化

在對 ResNet-50進行TensorRT優(yōu)化之后，TensorBoard可以使我們清楚地看到ResNet-50計算圖中的變化。如圖4所示，TensorRT幾乎優(yōu)化了整張計算圖，并使用名稱為“my_trt_op0”（使用紅色進行高亮）的單個節(jié)點進行替換。根據(jù)模型中的層和操作，由于優(yōu)化操作，TensorRT節(jié)點替換掉了圖中的一部分。命名為“ conv1”的操作其實不是真正的卷積層，只是將feature map的格式由NHWC轉換為NCHW（注：NCHW分別對應batch_size, channel, height, width）。

圖4 （a）經(jīng)過TensorBoard可視化之后的ResNet-50 （b）經(jīng)過TensorRT優(yōu)化之后的ResNet-50，其中的子圖被替換為TensorRT節(jié)點

在Volta GPU上使用Tensor Cores

與FP32或FP64相比，使用半精度（也被稱為FP16）可以降低神經(jīng)網(wǎng)絡的內存使用。FP16支持部署更大的網(wǎng)絡，同時比FP32或FP64花費更少的時間。NVIDIA的Volta架構硬件采用了稱為Tensor Cores的矩陣數(shù)學加速器。Tensor Cores提供4x4x4矩陣處理陣列，用于執(zhí)行操作D = A * B + C，其中A, B, C and D為4×4的矩陣。圖5展示了工作過程。矩陣乘法輸入A和B是FP16矩陣，而累加矩陣C和D可以是FP16或FP32矩陣。

圖5 Tensor Cores上的矩陣處理操作

當使用FP16 math檢測到推理時，TensorRT會自動使用硬件Tensor Cores。 Tensor Cores在NVIDIA Tesla V100上的峰值性能比雙精度（FP64）快一個數(shù)量級，而吞吐量比單精度（FP32）提高了4倍?？梢栽诤瘮?shù)create_inference_graph中，將precision_mode設置為“ FP16”啟用半精度計算，如下所示。getNetwork()用于從protobuf文件中讀出frozen network，返回數(shù)據(jù)結構為tf.GraphDef()的神經(jīng)網(wǎng)絡。

trt_graph = trt.create_inference_graph(
                getNetwork(network_file_name), 
                outputs,
                max_batch_size=batch_size, 
                max_workspace_size_bytes=workspace_size, 
                precision_mode=”FP16")

圖6所示，在相同的硬件環(huán)境，小于7ms延時的情況下，應用基于NVIDIA Volta Tensor Cores的TensorFlow-TensorRT集成后的ResNet-50比僅使用TensorFlow的推理速度快8倍。

圖6 ResNet-50推理吞吐量性能

使用INT8的精度進行推理

使用INT8精度執(zhí)行推理進一步提高了計算速度并降低了對帶寬的要求。減小的動態(tài)范圍給表示神經(jīng)網(wǎng)絡的權重和激活帶來挑戰(zhàn)。

各個數(shù)據(jù)類型的動態(tài)范圍

圖7 包含INT8推理的工作流程

trt_graph = trt.create_inference_graph(
                getNetwork(network_file_name), 
                outputs,
                max_batch_size=batch_size,
                max_workspace_size_bytes=workspace_size, 
                precision_mode=”INT8")

現(xiàn)在使用校準數(shù)據(jù)運行校準圖。 TensorRT使用節(jié)點數(shù)據(jù)的分布來量化節(jié)點的權重。您必須使用與實際問題數(shù)據(jù)集分布相同或者類似的校準數(shù)據(jù)。我們建議在首次使用INT8校準的模型時，檢查推理過程中的誤差累積。minimum_segment_size可以幫助調整優(yōu)化來最小化量化誤差。通過使用minimum_segment_size，您可以更改優(yōu)化的INT8引擎中的最小節(jié)點數(shù)，進而更改最終優(yōu)化計算圖以微調結果準確性。

在校驗數(shù)據(jù)上執(zhí)行完計算圖之后，使用函數(shù)calib_graph_to_infer_graph對檢驗之后的計算圖應用TensorRT優(yōu)化，這個函數(shù)同樣會將TensorFlow中的子圖替換為INT8優(yōu)化后的TensorRT節(jié)點。函數(shù)的輸出為可以用于推理的frozen TensorFlow graph。

trt_graph=trt.calib_graph_to_infer_graph(calibGraph)

在使用這兩個命令之后，就可以以INT8的精度對您的模型進行推理了。

如果您想查看此處顯示的示例，請查看運行這些示例所需的代碼：
https://developer.download.nvidia.com/devblogs/tftrt_sample.tar.xz

可用性

我們希望通過將TensorRT與TensorFlow集成，進而在使用NVIDIA GPU時可以獲得最高性能，同時保持TensorFlow的易用性和靈活性。 NVIDIA將繼續(xù)與Google TensorFlow團隊密切合作，進一步增強這些集成功能。隨著TensorRT支持更多網(wǎng)絡，開發(fā)人員將自動受益于更新，而無需更改現(xiàn)有代碼。

查找有關今天如何開始的說明：
https://www.tensorflow.org/install/install_linux

在不久的將來，我們希望標準的pip安裝過程也能正常運行。敬請關注！

我們相信，在使用GPU時，您會看到將TensorRT與TensorFlow集成的實質性好處。您可以在以下位置找到有關TensorFlow的更多信息：https://www.tensorflow.org/.

有關TensorRT的更多信息，請訪問NVIDIA的TensorRT頁面https://developer.nvidia.com/tensorrt.