雷鋒網(wǎng) AI 研習(xí)社按，日前，谷歌和英偉達(dá)宣布將 NVIDIA TensorRT 集成到 TensorFlow 1.7 中。在谷歌開(kāi)發(fā)者博客中，他們介紹了此次合作的詳細(xì)信息以及整合之后的性能，雷鋒網(wǎng) AI 研習(xí)社編譯整理如下：

TensorRT 是一個(gè)可以用于優(yōu)化深度學(xué)習(xí)模型，以進(jìn)行推理，并為生產(chǎn)環(huán)境中的 GPU 創(chuàng)建運(yùn)行環(huán)境的庫(kù)。它能優(yōu)化 TensorFlow 中的 FP16 浮點(diǎn)數(shù)和 INT8 整型數(shù)，并能自動(dòng)選擇針對(duì)特定平臺(tái)的內(nèi)核，以最大化吞吐量，并最大限度的降低 GPU 推理期間的延遲。全新的集成工作流程簡(jiǎn)化了在 TensorFlow 中使用 TensorRT 的步驟，同時(shí)使得 TensorFlow 達(dá)到了世界一流的性能水平。

經(jīng)測(cè)試，在 NVIDIA Volta Tensor 核心上，集成了 TensorRT 的 TensorFlow 運(yùn)行 ResNet-50 比沒(méi)有集成 TensorRT 的 TensorFlow 執(zhí)行速度提高了 8 倍。

優(yōu)化 TensorFlow 中的子圖

在 TensorFlow 1.7 中，TensorRT 可以用于優(yōu)化子圖，而 TensorFlow 執(zhí)行其余未優(yōu)化的部分。這個(gè)方法使得開(kāi)發(fā)者既能夠使用 TensorFlow 的眾多功能來(lái)快速構(gòu)建模型，同時(shí)也可以在執(zhí)行推理時(shí)使用 TensorRT 獲得強(qiáng)大的優(yōu)化能力。如果你嘗試過(guò)在之前的 TensorFlow 模型中使用 TensorRT，你應(yīng)該知道，要想使用某些不受支持的 TensorFlow 層，必須手動(dòng)導(dǎo)入，這在某些情況下可能會(huì)耗費(fèi)大量時(shí)間。

從工作流程的角度來(lái)看，開(kāi)發(fā)者可以使用 TensorRT 來(lái)優(yōu)化 TensorFlow 的每個(gè)子圖。

在推斷過(guò)程中，TensorFlow 先將執(zhí)行所有支持區(qū)域的圖，之后調(diào)用 TensorRT 去執(zhí)行那些經(jīng)過(guò) TensorRT 優(yōu)化過(guò)的節(jié)點(diǎn)。舉個(gè)例子，如果你的圖包含 A，B，C 三段，其中 B 段被 TensorRT 優(yōu)化過(guò)，B 將被一個(gè)節(jié)點(diǎn)代替。那么在推理過(guò)程中，TensorFlow 將先執(zhí)行 A，之后調(diào)用 TensorRT 執(zhí)行 B，最后 TensorFlow 執(zhí)行 C。

這個(gè)用于優(yōu)化 TensorRT 的新加入的 TensorFlow API，以凍結(jié)的 TensorFlow 圖為輸入，針對(duì)該子圖進(jìn)行優(yōu)化，最后將優(yōu)化過(guò)的推理子圖發(fā)送回 TensorFlow 中。

下面為一段示例代碼：

# Reserve memory for TensorRT inference engine

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = number_between_0_and_1)

... 

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)  # Get optimized graph

per_process_gpu_memory_fraction 這個(gè)參數(shù)定義了 TensorFlow 允許使用的 GPU 顯存的比例，剩余的顯存將分配給 TensorRT。這個(gè)參數(shù)應(yīng)該在 TensorFlow-TensorRT 進(jìn)程第一次啟動(dòng)的時(shí)候設(shè)定好。比如，per_process_gpu_fraction=0.67，那么 67% 的顯存會(huì)被分配給 TensorFlow，其余的 33% 會(huì)被分配給 TensorRT 引擎。

Create_inference_graph 函數(shù)將凍結(jié)住的 TensorFlow 圖作為輸入，返回一個(gè)經(jīng)過(guò) TensorRT 節(jié)點(diǎn)優(yōu)化過(guò)的圖。我們看看這個(gè)函數(shù)的參數(shù)：

• Input_graph_def:凍結(jié)住的 TensorFlow 圖

• Outputs:輸出節(jié)點(diǎn)名字的字符串列表，比如：[“resnet_v1_50/predictions/Resape_1”]

• Max_batch_size:整數(shù)，輸入的 batch size，比如，16

• Max_workspace_size_bytes:整數(shù)，能分配給 TensorRT 的最大 GPU 顯存大小

• Precision_mode:字符串，可選的值為「FP32」, 「FP16」, 「INT8」

舉個(gè)例子，如果 GPU 有 12GB 顯存，想要給 TensorRT 引擎分配 4GB 顯存，那么應(yīng)該設(shè)置 per_process_gpu_memory_fraction 為（12-4）/12=0.67，max_workspace_size_bytes=4,000,000,000.

我們來(lái)試著將這個(gè)新的 API 應(yīng)用在 ResNet-50 上，看看經(jīng)過(guò)優(yōu)化后的模型在 TensorBoard 中看起來(lái)是什么樣的。左側(cè)的圖像是沒(méi)有經(jīng)過(guò) TensorRT 優(yōu)化的 ResNet-50，右側(cè)是經(jīng)過(guò)優(yōu)化的。在這個(gè)設(shè)定下，大部分圖被 TensorRT 優(yōu)化，并用一個(gè)單一節(jié)點(diǎn)代替了（圖中高亮部分）。

經(jīng)過(guò)優(yōu)化的INT8推理性能

TensorRT 兼容單精度（FP32）和半精度（FP16）訓(xùn)練的模型（也可以將它們量化為 INT8），同時(shí)能盡可能減少由精度降低而導(dǎo)致的準(zhǔn)確率降低。INT8 模型能夠更快的計(jì)算，同時(shí)對(duì)帶寬的需求也會(huì)降低，但是因?yàn)榭捎玫膭?dòng)態(tài)范圍降低了，這也對(duì)神經(jīng)網(wǎng)絡(luò)的權(quán)重和激活表示提出了很大的挑戰(zhàn)。

為了解決這個(gè)問(wèn)題，TensorRT 使用了一個(gè)校正過(guò)程，以盡可能減小將 FP32 網(wǎng)絡(luò)近似成 8-bit 整型表示時(shí)的信息損失。在使用 TensorRT 優(yōu)化 TensorFlow 圖之后，可以使用下面的命令將圖傳遞給 TensorRT 進(jìn)行校準(zhǔn)，如下：

trt_graph=trt.calib_graph_to_infer_graph(calibGraph)

除此之外的網(wǎng)絡(luò)推理流程都沒(méi)有變化。這一步的輸出為一個(gè)可以被 TensorFlow 執(zhí)行的凍結(jié)圖。

在NVIDIA Volta GPU上自動(dòng)使用Tensor核心

在 NVIDIA Volta GPU 的 Tensor 核心上通過(guò) TensorRT 進(jìn)行半精度 TensorFlow 模型推理，能夠提供相較于單精度模型八倍的吞吐量。相較于更高精度的 FP32 或者 FP64，半精度數(shù)據(jù)（FP16）能夠減少神經(jīng)網(wǎng)絡(luò)的顯存使用量，這使得開(kāi)發(fā)者能夠訓(xùn)練和部署更大規(guī)模的神經(jīng)網(wǎng)絡(luò)，同時(shí) FP16 相比 FP32 和 FP64 有更少的傳輸時(shí)間。

如果每個(gè) Tensor 核心執(zhí)行的是 D=A*B+C，其中 A 和 B 為半精度 4*4 矩陣，D 和 C 是單精度或者半精度 4*4 矩陣，那么 V100 上此時(shí) Tensor 核心的峰值性能是雙精度（FP64）性能的 10 倍，是單精度（FP32）性能的 4 倍。

Google 目前已經(jīng)發(fā)布了 TensorFlow 1.7，同時(shí)也將跟 NVIDIA 更緊密地合作。希望這個(gè)新的解決方案額能夠在提供最強(qiáng)性能的同時(shí)，保持 TensorFlow 的易用性和靈活性。隨著 TensorRT 支持越來(lái)越多的網(wǎng)絡(luò)結(jié)構(gòu)，大家只要更新就可以享受到這些好處，而無(wú)須改寫(xiě)代碼。

使用標(biāo)準(zhǔn)pip install即可更新到 TensorFlow 1.7:

pip install tensorflow-gpu r1.7

詳細(xì)的安裝說(shuō)明可在下面找到：

https://github.com/tensorflow/tensorflow/tree/r1.7/tensorflow/contrib/tensorrt

via: Google Develop Blog

雷鋒網(wǎng) AI 研習(xí)社編譯整理。

雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見(jiàn)轉(zhuǎn)載須知。