TensorRT 4 之TF-TRT

其力斷金……

——子棐

前言

筆者在前面TensorRT系列中已經介紹了如何使用TensorRT加速Tensorflow模型，不知道各位小夥伴是否還有印象，如果有沒有的話可以查看歷史文章。在TensorRT3版本的TensorFlow的支持是通過UFF Converter和UFF Parser來實現的，也就是說通過將Frozen的TensorFlow模型轉化成為UFF模型，之後通過UFF Parser將模型導入TensorRT，之後進行優化生成Plan流圖，並最後部署該流圖。

那麼問題來了，是否有Tensorflow原生支持的TensorRT呢？現在很高興的和大家分享——TensorRT已經集成進TensorFlow 1.7版本中，可以通過以下命令安裝使用：

pip install tensorflow-gpu r1.7

在TensorFlow1.7正式發布前，可以通過以下鏈接安裝使用TF-TRT：

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

優化核心

TensorRT3已經支持TensorFlow的模型，而現在Tensorflow中原生支持的TensorRT4又是怎麼回事呢？

上圖分別比較了TensorFlow推理流程、TensorRT3的TensorFlow模型導出模式流程以及TensorRT4的TensorFlow原生支持模式流程。其中可以看到該兩種流程的區別在於，TRT3是通過支持導入TF模型來實現的，而TRT4是通過TF內部支持集成TRT4優化好的圖以實現直接使用TF完成推理流程，當然TRT4也是支持模型導出模式的。

新增的模式是什麼樣子呢？舉個例子：對於一個推理圖分別有A、B和C三部分，其中B部分可以使用TRT加速，因此在最新的版本推理流程為：

TF執行推理圖A部分；

TF調用TRT執行推理圖B部分（已優化）；

TF繼續執行推理圖C部分。

上圖是使用TensorBoard視覺化的Resnet50優化圖，其中左圖是TensorFlow推理圖，右圖是TensorRT的推理圖。可以看到TensorRT加速了幾乎整個圖，左圖大部分的內容在右圖中用單個節點「my_trt_op0」予以替換。

新基礎指南

TensorRT4是整合進TensorFlow中，因此可以直接使用TensorFlow進行推理部署，而且還能享受到TensorRT的優化，而這隻需要在代碼中加上以下兩部分代碼：

指定TensorFlow可以使用的GPU顯存比例，剩餘的部分將由TensorRT使用；

讓TensorRT分析TensorFlow推理圖，執行優化並替換掉可優化的計算節點。

對於第一點，使用GPUOptions函數中新的參數per_process_gpu_memory_fraction來指定TensorRT可以使用的顯存比列。這個參數需要在初始化整個流程時予以配置。例如，將0.67賦值per_process_gpu_memory_fraction，意味著67%的GPU顯存將分配給TensorFlow使用，而TensorRT使用餘下的33%。配置函數為：

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction = number_between_0_and_1)

對於第二點，TensorRT4使用的是新的函數create_inference_graph來優化Forzen的TensorFlow推理圖。調用方法如下：

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)

以上的輸入參數分別是：

frozen_graph_def: frozen的TensorFlow推理圖

put_node_name: 輸出的網路節點名稱列表

max_batch_size: 整數，最大的batchsize

max_workspace_size_bytes: 整數，TensorRT可使用的最大的GPU顯存空間

precision_mode: 字元串，可以使用 「FP32」, 「FP16」 or 「INT8」

需要注意的是，您需要同時配置per_process_gpu_memory_fraction和max_workspace_size_bytes參數來獲得最好的性能。例如，對於12G顯存的GPU，將4G顯存分配給TensorRT使用時，需要將per_process_gpu_memory_fraction 設置為0.56，將max_workspace_size_bytes 設置為4000000000。

用於對比，對於TensorFlow模型導出模式，其主要API為：

可以看到模型導出模式是在TensorRT Runtime環境下執行由TensorRT 模型Engine創建的context，而原生支持模式是直接使用TensorFlow進行推理部署，TensorFlow具有直接解析TensorRT優化推理圖的能力。

大殺器之Tensor Core

關於V100中新增的Tensor Core，筆者已經在前面多篇文章中予以解釋，這邊就不再贅述了，如果對Tensor Core還不了解的小夥伴，請先查看歷史文章。

當設置FP16計算以及擁有Tensor Core硬體支持時，TensorRT將自動使用Tensor Core進行推理運算。因此這邊只需要將create_inference_graph中的precision_mode設置為「FP16」即可。參考代碼如下所示：

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")

INT8來矣~

以上是TensorRT 4中TensorFlow的模型INT8推理流程，與FP16推理不同的是

需要先生成校準推理圖並進行校準；

創建TensorRT優化過的推理圖。

第一步，調用函數create_inference_graph function並將precision_mode參數設置為「INT8」以校準模型。這個函數的輸出為Frozen的TensorFlow推理圖並執行校準，函數調用範例：

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")

執行校準的時候需要使用校準數據，校準這部分的原理和以前的文章闡明的一致，需要注意的是在第一次使用校準的模型時需要先觀察統計下整體錯誤率。

對推理圖完成校準後，需要調用函數calib_graph_to_infer_graph來將TensorRT優化的節點應用至TensorFlow推理圖中，函數調用範例：

trt_graph=trt.calib_graph_to_infer_graph(calibGraph)

該函數可以使用INT8的子圖替換掉TensorFlow推理圖的子圖，該函數的輸出是Frozen的可用於正常執行的TensorFlow推理圖。

結語

行文至此，本篇主要介紹的是TensorRT4中關於TensorFlow原生支持的問題，通過這個新加的特性可以實現直接使用TensorFlow運行TensorRT優化過的推理圖。實驗也證明了，TensorRT的優化能極大的帶來性能的提升。在接下來的文章中，我會逐步的為大家介紹TensorRT4的新特性。

以上

子棐

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！