本文目的主要在于如何使用TensorRT 5.x的python api來進行神經網絡的推理。因為目前TensorRT只支持 ONNX ， Caffe 和 Uff (Universal Framework Format)這三種格式。這里以tensorflow的pb模型為例（可以無縫轉換為uff）進行說明。
0. TensoRT介紹
TensorRT 是英偉達(NVIDIA)開發的一個可以在NVIDIA旗下的GPU上進行高性能推理的C++庫。它的設計目標是與現有的深度學習框架無縫貼合：比如Mxnet, PyTorch, Tensorflow 以及Caffe等。TensorRT只關注 推理階段 (inference stage)的優化。
TensorRT的作用是將一個訓練好的神經網絡(需要部署的)經過TensorRT優化，然后放到TensorRT的Runtime引擎去執行。示例圖如下：
下圖可以看出，TensorRT支持很多的框架和對其下的GPU平臺都做了特定的支持。
總的來說，使用TensorRT進行模型推理加速的
工作流
(Workflow)可以分為以下幾個步驟：
-
① 訓練神經網絡，得到模型。(以Tensorflow為例，我們得到
***.pb格式的模型文件) - ② 將模型用TensorRT提供的工具進行 parsing (解析)。
- ③ 將parsing后的結構通過TensorRT內部的 優化選項 (optimization options)對計算圖結構進行優化。（包括不限于： 1. 算子融合 Layer Funsion : 通過將Conv + BN等層的融合降低數據的吞吐量。 2. 精度校準 Precision Calibration : 當用戶為了節省計算資源使用INT8進行推理的時候，需要作精度校準，這個操作TensorRT提供了官方的支持。 3. kernel auto-tuning : 根據計算邏輯，自動選擇TensorRT實現的更高效的矩陣乘法，卷積運算等邏輯。）
- ④ 通過上述步驟，得到了一個優化后的推理引擎。我們就可以拿這個引擎進行推理了~
1. TensorRT 的關鍵接口
本部分將介紹TensorRT中的關鍵接口，這些接口在我們的例子中都有著確實的應用，所以有必要了解設計他們的目的和他們的功能。
Network Definition
(高階用法)
網絡定義接口提供方法來指定網絡的定義。我們可以指定輸入輸出的Tensor Name以及增加layer。我們可以通過自定義來擴展TensorRT不支持的層和功能。關于網絡定義的API，請查看Network Definition API.
Builder
Builder接口讓我們可以從一個網絡定義中創建一個優化后的引擎。在這個步驟我們可以選擇最大batch，workspace size，精度級別等參數，如果你想了解更多，請查看Builder API.
Engine
引擎接口就允許應用來執行推理了。它支持對
引擎輸入和輸出的綁定
進行同步和異步執行、分析、枚舉和查詢。
It supports synchronous and asynchronous execution, profiling, and
enumeration and querying of the bindings for the engine inputs and
outputs.
此外，引擎接口還允許單個引擎擁有多個 執行上下文 (execution contexts). 這可以讓一個引擎同時對多組數據進行執行推理操作。如想了解更多，請查Execution API.
此外，TensorRT還提供了解析訓練模型并創建TensorRT內部支持的模型定義的parser：
- Caffe Parser
This parser can be used to parse a Caffe network created in BVLC Caffe
or NVCaffe 0.16. It also provides the ability to register a plugin
factory for custom layers. For more details on the C++ Caffe Parser,
see NvCaffeParser or the Python Caffe Parser.
- UFF Parser
This parser can be used to parse a network in UFF format. It also
provides the ability to register a plugin factory and pass field
attributes for custom layers. For more details on the C++ UFF Parser,
see NvUffParser or the Python UFF Parser.
- ONNX Parser
This parser can be used to parse an ONNX model. For more details on
the C++ ONNX Parser, see NvONNXParser or the Python ONNX Parser.
Restriction: Since the ONNX format is quickly developing, you may
encounter a version mismatch between the model version and the parser
version. The ONNX Parser shipped with TensorRT 5.1.x supports ONNX IR
(Intermediate Representation) version 0.0.3, opset version 9.
2. TensorRT python api安裝
環境準備:
- Ubuntu 18.04 (需要注意的是, TensorRT的python接口不支持windows，在windows上使用只能用C++接口)
- cuda 9.0
- cudnn 7.5.0
- python 3.6
本著用新不用舊的原則，這里還是用TensorRT 5.x版本，下載地址
[4]
：
下好之后，我們看到有下面這個文件躺在你設定的下載路徑中，將其解壓即可：
正式開始:
-
a. 添加環境變量.
我將TensorRT-5.1.2.2解壓后放在下面這個路徑中，如圖：
添加環境變量：
vim ~/.bashrc
...
export LD_LIBRARY_PATH=/home/samuel/gaodaiheng/TensorRT/TensorRT-5.1.2.2/lib:$LD_LIBRARY_PATH
...
source ~/.bashrc
可以看到，tensorrt的lib已經鏈接上了.
-
b. 安裝wheel包.
注意：只安裝跟你環境一致的wheel包，因為我的python版本是python3.6, 所以這里安裝的是cp36.whl）
TensorRT-5.1.2.2/python
TensorRT-5.1.2.2/uff
TensorRT-5.1.2.2/graphsurgeon
安裝好之后，檢查一下：
emmm, 成功。
- c. 安裝pycuda.
直接使用
pip3 install pycuda
會存在問題，其原因是：“
因為我安裝的cuda不是安裝的debian版本，所以它默認找不到
”
所以需要使用tar包編譯的方式安裝。
安裝過程：https://wiki.tiker.net/PyCuda/Installation/Linux/Ubuntu
tar包下載地址： https://pypi.org/project/pycuda/#files
tar包下載完解壓后的地址：
然后我們進行編譯(可能根據情況需要改你的
--python-exe
后面的內容)：
.
/
configure
.
py
--
python
-
exe
=
/
usr
/
bin
/
python3
--
cuda
-
root
=
/
usr
/
local
/
cuda
make
-
j
4
sudo python setup
.
py install
sudo pip install
.
notice : 如第一步遇
setuptools找不到的時候，用python3 configure.py --python-exe=/usr/bin/python3 --cuda-root=/usr/local/cuda
驗證是否安裝成功:
3. TensorRT python 使用
這里以Tensorflow的pb模型出發(我假設你已經有了一個pb格式的模型)，下面的內容將一步一步的闡明如何通過這個pb模型來轉換為TensorRT的引擎并執行推理。
3.1 pb轉uff*
首先，我們需要把pb格式的模型轉換為TensorRT接收的3種模型結構之一: caffe, onnx, uff的uff.
# coding: UTF-8
"""
@func: 將pb模型轉換為uff模型.
"""
import
tensorflow
as
tf
import
uff
# 需要改成你自己的output_names
output_names
=
[
'output'
]
frozen_graph_filename
=
'xxx.pb'
# 將frozen graph轉換為uff格式
uff_model
=
uff
.
from_tensorflow_frozen_model
(
frozen_graph_filename
,
output_names
)
轉化成功，在graph_opt.pb的同級目錄下，有了一個
xxx.uff
模型文件.
3.2 開始推理
■ 初始化
.
.
.
engine
=
build_engine
(
)
# 這個步驟比較重要, 下面會詳細說明.
# engine.get_binding_shape(0) 對應的是輸入的shape; 1對應的是輸出的shape.
print
(
engine
.
get_binding_shape
(
0
)
)
print
(
engine
.
get_binding_shape
(
1
)
)
# 1. 為輸入輸出指定host和device的空間. h_input的h表示host, d_input的d表示device:
h_input
=
cuda
.
pagelocked_empty
(
trt
.
volume
(
engine
.
get_binding_shape
(
0
)
)
,
dtype
=
trt
.
nptype
(
trt
.
float32
)
)
h_output
=
cuda
.
pagelocked_empty
(
trt
.
volume
(
engine
.
get_binding_shape
(
1
)
)
,
dtype
=
trt
.
nptype
(
trt
.
float32
)
)
# 2. 分配顯存空間.
d_input
=
cuda
.
mem_alloc
(
h_input
.
nbytes
)
d_output
=
cuda
.
mem_alloc
(
h_output
.
nbytes
)
# 3. 創建cuda stream, 其可以host和device之間動態的轉換數據(顯存到內存，內存到顯存), 并執行推理.
stream
=
cuda
.
Stream
(
)
context
=
engine
.
create_execution_context
(
)
.
.
.
■ 構造引擎
.
.
.
# TRT_LOGGER是trt.Builder中必須要傳的參數！
TRT_LOGGER
=
trt
.
Logger
(
trt
.
Logger
.
WARNING
)
.
.
.
def
build_engine
(
)
:
uff_model
=
uff
.
from_tensorflow_frozen_model
(
ModelData
.
FROZEN_GRAPH
,
ModelData
.
OUTPUT_NAME
)
# For more information on TRT basics, refer to the introductory samples.
with
trt
.
Builder
(
TRT_LOGGER
)
as
builder
,
builder
.
create_network
(
)
as
network
,
trt
.
UffParser
(
)
as
parser
:
# 設置batch size和工作區大小.
builder
.
max_batch_size
=
1
# builder.max_workspace_size = 1 ** 30
# Parse the Uff Network
parser
.
register_input
(
ModelData
.
INPUT_NAME
,
ModelData
.
INPUT_SHAPE
,
trt
.
UffInputOrder
.
NHWC
)
parser
.
register_output
(
'output'
)
# parse_buffer是parse一個已經在內存中的buffer.
# https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/parsers/Uff/pyUff.html
parser
.
parse_buffer
(
uff_model
,
network
)
# Build and return an engine.
return
builder
.
build_cuda_engine
(
network
)
.
.
.
好，現在引擎和初始化的流程已經完成。我們可以進行模型的推理了，完整代碼如下:
# coding: UTF-8
"""
@author: samuel ko
"""
import
os
import
cv2
import
time
import
tensorflow
as
tf
import
tensorrt
as
trt
import
pycuda
.
driver
as
cuda
import
pycuda
.
autoinit
import
uff
import
numpy
as
np
TRT_LOGGER
=
trt
.
Logger
(
trt
.
Logger
.
WARNING
)
class
ModelData
(
object
)
:
FROZEN_GRAPH
=
"xxx.pb"
# NHWC
INPUT_NAME
=
"image"
# ref: https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/infer/FoundationalTypes/Dims.html#dimsnchw
INPUT_SHAPE
=
[
432
,
848
,
3
]
OUTPUT_NAME
=
[
'output'
]
# 構建引擎.
def
build_engine
(
)
:
uff_model
=
uff
.
from_tensorflow_frozen_model
(
ModelData
.
FROZEN_GRAPH
,
ModelData
.
OUTPUT_NAME
)
# For more information on TRT basics, refer to the introductory samples.
with
trt
.
Builder
(
TRT_LOGGER
)
as
builder
,
builder
.
create_network
(
)
as
network
,
trt
.
UffParser
(
)
as
parser
:
# 設置batch size和工作區大小.
builder
.
max_batch_size
=
1
# builder.max_workspace_size = 1 ** 30
# Parse the Uff Network
parser
.
register_input
(
ModelData
.
INPUT_NAME
,
ModelData
.
INPUT_SHAPE
,
trt
.
UffInputOrder
.
NHWC
)
parser
.
register_output
(
'Openpose/concat_stage7'
)
# parse_buffer是parse一個已經在內存中的buffer.
# https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/parsers/Uff/pyUff.html
parser
.
parse_buffer
(
uff_model
,
network
)
# Build and return an engine.
return
builder
.
build_cuda_engine
(
network
)
# 加載數據并將其喂入提供的pagelocked_buffer中.
def
load_normalized_data
(
data_path
,
pagelocked_buffer
,
target_size
=
(
848
,
432
)
)
:
upsample_size
=
[
int
(
target_size
[
1
]
/
8
*
4.0
)
,
int
(
target_size
[
0
]
/
8
*
4.0
)
]
img
=
cv2
.
imread
(
data_path
)
img
=
cv2
.
resize
(
img
,
target_size
,
interpolation
=
cv2
.
INTER_CUBIC
)
# 此時img.shape為H * W * C: 432, 848, 3
print
(
"圖片shape"
,
img
.
shape
)
# Flatten the image into a 1D array, normalize, and copy to pagelocked memory.
np
.
copyto
(
pagelocked_buffer
,
img
.
ravel
(
)
)
return
upsample_size
# 初始化(創建引擎，為輸入輸出開辟&分配顯存/內存.)
def
init
(
)
:
engine
=
build_engine
(
)
print
(
engine
.
get_binding_shape
(
0
)
)
print
(
engine
.
get_binding_shape
(
1
)
)
# 1. Allocate some host and device buffers for inputs and outputs:
h_input
=
cuda
.
pagelocked_empty
(
trt
.
volume
(
engine
.
get_binding_shape
(
0
)
)
,
dtype
=
trt
.
nptype
(
trt
.
float32
)
)
h_output
=
cuda
.
pagelocked_empty
(
trt
.
volume
(
engine
.
get_binding_shape
(
1
)
)
,
dtype
=
trt
.
nptype
(
trt
.
float32
)
)
# Allocate device memory for inputs and outputs.
d_input
=
cuda
.
mem_alloc
(
h_input
.
nbytes
)
d_output
=
cuda
.
mem_alloc
(
h_output
.
nbytes
)
# Create a stream in which to copy inputs/outputs and run inference.
stream
=
cuda
.
Stream
(
)
context
=
engine
.
create_execution_context
(
)
return
context
,
h_input
,
h_output
,
stream
,
d_input
,
d_output
# 推理
def
inference
(
data_path
)
:
global
context
,
h_input
,
h_output
,
stream
,
d_input
,
d_output
upsample
=
load_normalized_data
(
data_path
,
pagelocked_buffer
=
h_input
)
t1
=
time
.
time
(
)
cuda
.
memcpy_htod_async
(
d_input
,
h_input
,
stream
)
# Run inference.
context
.
execute_async
(
bindings
=
[
int
(
d_input
)
,
int
(
d_output
)
]
,
stream_handle
=
stream
.
handle
)
# Transfer predictions back from the GPU.
cuda
.
memcpy_dtoh_async
(
h_output
,
d_output
,
stream
)
# Synchronize the stream
stream
.
synchronize
(
)
# Return the host output.
print
(
"推理時間"
,
time
.
time
(
)
-
t1
)
return
h_output
,
upsample
if
__name__
==
'__main__'
:
context
,
h_input
,
h_output
,
stream
,
d_input
,
d_output
=
init
(
)
img_path
=
"xxx.jpg"
image
=
cv2
.
imread
(
img_path
)
for
_
in
range
(
10
)
:
output
,
upsample
=
inference
(
data_path
=
img_path
)
print
(
type
(
output
)
)
print
(
output
.
shape
)
···
4. 遇到的問題
-
① 如何使用
NHWC格式的模型進行推理。
之前誤以為TensorRT只支持推理NCHWLayout的模型，這一步坑了好久，因為tensorRT不支持transpose算子，所以我將pb模型轉成pbtxt，把data format都改成NCHW，還把concat的axis都從3改成1. 結果還是不對。
其實就很簡單: 1. 在register_input的時候用
trt.UffInputOrder.NHWC
parser.register_input("image", [432, 848, 3], trt.UffInputOrder.NHWC)
-
② TensorRT和Tensorflow Session無法并存。
解決方案是在系統中把tf的session和Trt的引擎分別放在不同的GPU上。
5. 總結
到這里，python使用TensorRT 5.x來推理模型的一整套流程就已經完畢了。如果有更高階的需要，比如自定義層讓TensorRT支持等，就需要大家深挖文檔，進行魔改。遇到的最大的坑就是誤以為TensorRT無法使用
NHWC
的布局，差點讓我以
NCHW
的layour重新訓練模型，所幸找到了解決辦法。
參考資料
[1] 3. Working With TensorRT Using The Python API 中文翻譯
[2] 3. Working With TensorRT Using The Python API 官方
[3] TensorRT 5.1.5.0 Python文檔
[4] TensorRT 5.1.2.2 下載地址
更多文章、技術交流、商務合作、聯系博主
微信掃碼或搜索：z360901061
微信掃一掃加我為好友
QQ號聯系： 360901061
您的支持是博主寫作最大的動力，如果您喜歡我的文章，感覺我的文章對您有幫助，請用微信掃描下面二維碼支持博主2元、5元、10元、20元等您想捐的金額吧，狠狠點擊下面給點支持吧，站長非常感激您！手機微信長按不能支付解決辦法：請將微信支付二維碼保存到相冊，切換到微信，然后點擊微信右上角掃一掃功能，選擇支付二維碼完成支付。
【本文對您有幫助就好】元
