0
雷鋒網(wǎng) AI 研習社按,日前,谷歌和英偉達宣布將 NVIDIA TensorRT 集成到 TensorFlow 1.7 中。在谷歌開發(fā)者博客中,他們介紹了此次合作的詳細信息以及整合之后的性能,雷鋒網(wǎng) AI 研習社編譯整理如下:
TensorRT 是一個可以用于優(yōu)化深度學習模型,以進行推理,并為生產(chǎn)環(huán)境中的 GPU 創(chuàng)建運行環(huán)境的庫。它能優(yōu)化 TensorFlow 中的 FP16 浮點數(shù)和 INT8 整型數(shù),并能自動選擇針對特定平臺的內核,以最大化吞吐量,并最大限度的降低 GPU 推理期間的延遲。全新的集成工作流程簡化了在 TensorFlow 中使用 TensorRT 的步驟,同時使得 TensorFlow 達到了世界一流的性能水平。
經(jīng)測試,在 NVIDIA Volta Tensor 核心上,集成了 TensorRT 的 TensorFlow 運行 ResNet-50 比沒有集成 TensorRT 的 TensorFlow 執(zhí)行速度提高了 8 倍。
優(yōu)化 TensorFlow 中的子圖
在 TensorFlow 1.7 中,TensorRT 可以用于優(yōu)化子圖,而 TensorFlow 執(zhí)行其余未優(yōu)化的部分。這個方法使得開發(fā)者既能夠使用 TensorFlow 的眾多功能來快速構建模型,同時也可以在執(zhí)行推理時使用 TensorRT 獲得強大的優(yōu)化能力。如果你嘗試過在之前的 TensorFlow 模型中使用 TensorRT,你應該知道,要想使用某些不受支持的 TensorFlow 層,必須手動導入,這在某些情況下可能會耗費大量時間。
從工作流程的角度來看,開發(fā)者可以使用 TensorRT 來優(yōu)化 TensorFlow 的每個子圖。
在推斷過程中,TensorFlow 先將執(zhí)行所有支持區(qū)域的圖,之后調用 TensorRT 去執(zhí)行那些經(jīng)過 TensorRT 優(yōu)化過的節(jié)點。舉個例子,如果你的圖包含 A,B,C 三段,其中 B 段被 TensorRT 優(yōu)化過,B 將被一個節(jié)點代替。那么在推理過程中,TensorFlow 將先執(zhí)行 A,之后調用 TensorRT 執(zhí)行 B,最后 TensorFlow 執(zhí)行 C。
這個用于優(yōu)化 TensorRT 的新加入的 TensorFlow API,以凍結的 TensorFlow 圖為輸入,針對該子圖進行優(yōu)化,最后將優(yōu)化過的推理子圖發(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 這個參數(shù)定義了 TensorFlow 允許使用的 GPU 顯存的比例,剩余的顯存將分配給 TensorRT。這個參數(shù)應該在 TensorFlow-TensorRT 進程第一次啟動的時候設定好。比如,per_process_gpu_fraction=0.67,那么 67% 的顯存會被分配給 TensorFlow,其余的 33% 會被分配給 TensorRT 引擎。
Create_inference_graph 函數(shù)將凍結住的 TensorFlow 圖作為輸入,返回一個經(jīng)過 TensorRT 節(jié)點優(yōu)化過的圖。我們看看這個函數(shù)的參數(shù):
Input_graph_def:凍結住的 TensorFlow 圖
Outputs:輸出節(jié)點名字的字符串列表,比如:[“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」
舉個例子,如果 GPU 有 12GB 顯存,想要給 TensorRT 引擎分配 4GB 顯存,那么應該設置 per_process_gpu_memory_fraction 為(12-4)/12=0.67,max_workspace_size_bytes=4,000,000,000.
我們來試著將這個新的 API 應用在 ResNet-50 上,看看經(jīng)過優(yōu)化后的模型在 TensorBoard 中看起來是什么樣的。左側的圖像是沒有經(jīng)過 TensorRT 優(yōu)化的 ResNet-50,右側是經(jīng)過優(yōu)化的。在這個設定下,大部分圖被 TensorRT 優(yōu)化,并用一個單一節(jié)點代替了(圖中高亮部分)。
經(jīng)過優(yōu)化的INT8推理性能
TensorRT 兼容單精度(FP32)和半精度(FP16)訓練的模型(也可以將它們量化為 INT8),同時能盡可能減少由精度降低而導致的準確率降低。INT8 模型能夠更快的計算,同時對帶寬的需求也會降低,但是因為可用的動態(tài)范圍降低了,這也對神經(jīng)網(wǎng)絡的權重和激活表示提出了很大的挑戰(zhàn)。
為了解決這個問題,TensorRT 使用了一個校正過程,以盡可能減小將 FP32 網(wǎng)絡近似成 8-bit 整型表示時的信息損失。在使用 TensorRT 優(yōu)化 TensorFlow 圖之后,可以使用下面的命令將圖傳遞給 TensorRT 進行校準,如下:
trt_graph=trt.calib_graph_to_infer_graph(calibGraph)
除此之外的網(wǎng)絡推理流程都沒有變化。這一步的輸出為一個可以被 TensorFlow 執(zhí)行的凍結圖。
在NVIDIA Volta GPU上自動使用Tensor核心
在 NVIDIA Volta GPU 的 Tensor 核心上通過 TensorRT 進行半精度 TensorFlow 模型推理,能夠提供相較于單精度模型八倍的吞吐量。相較于更高精度的 FP32 或者 FP64,半精度數(shù)據(jù)(FP16)能夠減少神經(jīng)網(wǎng)絡的顯存使用量,這使得開發(fā)者能夠訓練和部署更大規(guī)模的神經(jīng)網(wǎng)絡,同時 FP16 相比 FP32 和 FP64 有更少的傳輸時間。
如果每個 Tensor 核心執(zhí)行的是 D=A*B+C,其中 A 和 B 為半精度 4*4 矩陣,D 和 C 是單精度或者半精度 4*4 矩陣,那么 V100 上此時 Tensor 核心的峰值性能是雙精度(FP64)性能的 10 倍,是單精度(FP32)性能的 4 倍。
Google 目前已經(jīng)發(fā)布了 TensorFlow 1.7,同時也將跟 NVIDIA 更緊密地合作。希望這個新的解決方案額能夠在提供最強性能的同時,保持 TensorFlow 的易用性和靈活性。隨著 TensorRT 支持越來越多的網(wǎng)絡結構,大家只要更新就可以享受到這些好處,而無須改寫代碼。
使用標準pip install即可更新到 TensorFlow 1.7:
pip install tensorflow-gpu r1.7
詳細的安裝說明可在下面找到:
https://github.com/tensorflow/tensorflow/tree/r1.7/tensorflow/contrib/tensorrt
via: Google Develop Blog
雷鋒網(wǎng) AI 研習社編譯整理。
雷峰網(wǎng)版權文章,未經(jīng)授權禁止轉載。詳情見轉載須知。