本文目的主要在于如何使用TensorRT 5.x的python api來進(jìn)行神經(jīng)網(wǎng)絡(luò)的推理。因為目前TensorRT只支持 ONNX ， Caffe 和 Uff (Universal Framework Format)這三種格式。這里以tensorflow的pb模型為例（可以無縫轉(zhuǎn)換為uff）進(jìn)行說明。
0. TensoRT介紹
TensorRT 是英偉達(dá)(NVIDIA)開發(fā)的一個可以在NVIDIA旗下的GPU上進(jìn)行高性能推理的C++庫。它的設(shè)計目標(biāo)是與現(xiàn)有的深度學(xué)習(xí)框架無縫貼合：比如Mxnet, PyTorch, Tensorflow 以及Caffe等。TensorRT只關(guān)注 推理階段 (inference stage)的優(yōu)化。
TensorRT的作用是將一個訓(xùn)練好的神經(jīng)網(wǎng)絡(luò)(需要部署的)經(jīng)過TensorRT優(yōu)化，然后放到TensorRT的Runtime引擎去執(zhí)行。示例圖如下：
下圖可以看出，TensorRT支持很多的框架和對其下的GPU平臺都做了特定的支持。
總的來說，使用TensorRT進(jìn)行模型推理加速的
工作流
(Workflow)可以分為以下幾個步驟：
-
① 訓(xùn)練神經(jīng)網(wǎng)絡(luò)，得到模型。(以Tensorflow為例，我們得到
***.pb格式的模型文件) - ② 將模型用TensorRT提供的工具進(jìn)行 parsing (解析)。
- ③ 將parsing后的結(jié)構(gòu)通過TensorRT內(nèi)部的 優(yōu)化選項 (optimization options)對計算圖結(jié)構(gòu)進(jìn)行優(yōu)化。（包括不限于： 1. 算子融合 Layer Funsion : 通過將Conv + BN等層的融合降低數(shù)據(jù)的吞吐量。 2. 精度校準(zhǔn) Precision Calibration : 當(dāng)用戶為了節(jié)省計算資源使用INT8進(jìn)行推理的時候，需要作精度校準(zhǔn)，這個操作TensorRT提供了官方的支持。 3. kernel auto-tuning : 根據(jù)計算邏輯，自動選擇TensorRT實(shí)現(xiàn)的更高效的矩陣乘法，卷積運(yùn)算等邏輯。）
- ④ 通過上述步驟，得到了一個優(yōu)化后的推理引擎。我們就可以拿這個引擎進(jìn)行推理了~
1. TensorRT 的關(guān)鍵接口
本部分將介紹TensorRT中的關(guān)鍵接口，這些接口在我們的例子中都有著確實(shí)的應(yīng)用，所以有必要了解設(shè)計他們的目的和他們的功能。
Network Definition
(高階用法)
網(wǎng)絡(luò)定義接口提供方法來指定網(wǎng)絡(luò)的定義。我們可以指定輸入輸出的Tensor Name以及增加layer。我們可以通過自定義來擴(kuò)展TensorRT不支持的層和功能。關(guān)于網(wǎng)絡(luò)定義的API，請查看Network Definition API.
Builder
Builder接口讓我們可以從一個網(wǎng)絡(luò)定義中創(chuàng)建一個優(yōu)化后的引擎。在這個步驟我們可以選擇最大batch，workspace size，精度級別等參數(shù)，如果你想了解更多，請查看Builder API.
Engine
引擎接口就允許應(yīng)用來執(zhí)行推理了。它支持對
引擎輸入和輸出的綁定
進(jìn)行同步和異步執(zhí)行、分析、枚舉和查詢。
It supports synchronous and asynchronous execution, profiling, and
enumeration and querying of the bindings for the engine inputs and
outputs.
此外，引擎接口還允許單個引擎擁有多個 執(zhí)行上下文 (execution contexts). 這可以讓一個引擎同時對多組數(shù)據(jù)進(jìn)行執(zhí)行推理操作。如想了解更多，請查Execution API.
此外，TensorRT還提供了解析訓(xùn)練模型并創(chuàng)建TensorRT內(nèi)部支持的模型定義的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安裝
環(huán)境準(zhǔn)備:
- Ubuntu 18.04 (需要注意的是, TensorRT的python接口不支持windows，在windows上使用只能用C++接口)
- cuda 9.0
- cudnn 7.5.0
- python 3.6
本著用新不用舊的原則，這里還是用TensorRT 5.x版本，下載地址
[4]
：
下好之后，我們看到有下面這個文件躺在你設(shè)定的下載路徑中，將其解壓即可：
正式開始:
-
a. 添加環(huán)境變量.
我將TensorRT-5.1.2.2解壓后放在下面這個路徑中，如圖：
添加環(huán)境變量：
vim ~/.bashrc
...
export LD_LIBRARY_PATH=/home/samuel/gaodaiheng/TensorRT/TensorRT-5.1.2.2/lib:$LD_LIBRARY_PATH
...
source ~/.bashrc
可以看到，tensorrt的lib已經(jīng)鏈接上了.
-
b. 安裝wheel包.
注意：只安裝跟你環(huán)境一致的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版本，所以它默認(rèn)找不到
”
所以需要使用tar包編譯的方式安裝。
安裝過程：https://wiki.tiker.net/PyCuda/Installation/Linux/Ubuntu
tar包下載地址： https://pypi.org/project/pycuda/#files
tar包下載完解壓后的地址：
然后我們進(jìn)行編譯(可能根據(jù)情況需要改你的
--python-exe
后面的內(nèi)容)：
.
/
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模型出發(fā)(我假設(shè)你已經(jīng)有了一個pb格式的模型)，下面的內(nèi)容將一步一步的闡明如何通過這個pb模型來轉(zhuǎn)換為TensorRT的引擎并執(zhí)行推理。
3.1 pb轉(zhuǎn)uff*
首先，我們需要把pb格式的模型轉(zhuǎn)換為TensorRT接收的3種模型結(jié)構(gòu)之一: caffe, onnx, uff的uff.
# coding: UTF-8
"""
@func: 將pb模型轉(zhuǎn)換為uff模型.
"""
import
tensorflow
as
tf
import
uff
# 需要改成你自己的output_names
output_names
=
[
'output'
]
frozen_graph_filename
=
'xxx.pb'
# 將frozen graph轉(zhuǎn)換為uff格式
uff_model
=
uff
.
from_tensorflow_frozen_model
(
frozen_graph_filename
,
output_names
)
轉(zhuǎn)化成功，在graph_opt.pb的同級目錄下，有了一個
xxx.uff
模型文件.
3.2 開始推理
■ 初始化
.
.
.
engine
=
build_engine
(
)
# 這個步驟比較重要, 下面會詳細(xì)說明.
# engine.get_binding_shape(0) 對應(yīng)的是輸入的shape; 1對應(yīng)的是輸出的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. 創(chuàng)建cuda stream, 其可以host和device之間動態(tài)的轉(zhuǎn)換數(shù)據(jù)(顯存到內(nèi)存，內(nèi)存到顯存), 并執(zhí)行推理.
stream
=
cuda
.
Stream
(
)
context
=
engine
.
create_execution_context
(
)
.
.
.
■ 構(gòu)造引擎
.
.
.
# TRT_LOGGER是trt.Builder中必須要傳的參數(shù)！
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
:
# 設(shè)置batch size和工作區(qū)大小.
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一個已經(jīng)在內(nèi)存中的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
)
.
.
.
好，現(xiàn)在引擎和初始化的流程已經(jīng)完成。我們可以進(jìn)行模型的推理了，完整代碼如下:
# 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'
]
# 構(gòu)建引擎.
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
:
# 設(shè)置batch size和工作區(qū)大小.
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一個已經(jīng)在內(nèi)存中的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
)
# 加載數(shù)據(jù)并將其喂入提供的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
# 初始化(創(chuàng)建引擎，為輸入輸出開辟&分配顯存/內(nèi)存.)
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格式的模型進(jìn)行推理。
之前誤以為TensorRT只支持推理NCHWLayout的模型，這一步坑了好久，因為tensorRT不支持transpose算子，所以我將pb模型轉(zhuǎn)成pbtxt，把data format都改成NCHW，還把concat的axis都從3改成1. 結(jié)果還是不對。
其實(shí)就很簡單: 1. 在register_input的時候用
trt.UffInputOrder.NHWC
parser.register_input("image", [432, 848, 3], trt.UffInputOrder.NHWC)
-
② TensorRT和Tensorflow Session無法并存。
解決方案是在系統(tǒng)中把tf的session和Trt的引擎分別放在不同的GPU上。
5. 總結(jié)
到這里，python使用TensorRT 5.x來推理模型的一整套流程就已經(jīng)完畢了。如果有更高階的需要，比如自定義層讓TensorRT支持等，就需要大家深挖文檔，進(jìn)行魔改。遇到的最大的坑就是誤以為TensorRT無法使用
NHWC
的布局，差點(diǎn)讓我以
NCHW
的layour重新訓(xùn)練模型，所幸找到了解決辦法。
參考資料
[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 下載地址
更多文章、技術(shù)交流、商務(wù)合作、聯(lián)系博主
微信掃碼或搜索：z360901061
微信掃一掃加我為好友
QQ號聯(lián)系： 360901061
您的支持是博主寫作最大的動力，如果您喜歡我的文章，感覺我的文章對您有幫助，請用微信掃描下面二維碼支持博主2元、5元、10元、20元等您想捐的金額吧，狠狠點(diǎn)擊下面給點(diǎn)支持吧，站長非常感激您！手機(jī)微信長按不能支付解決辦法：請將微信支付二維碼保存到相冊，切換到微信，然后點(diǎn)擊微信右上角掃一掃功能，選擇支付二維碼完成支付。
【本文對您有幫助就好】元
