TensorRT 是Nvidia 提出的深度学习推理平台,能够在GPU 上实现低延迟、高吞吐量的部属。基于TensorRT 的推论运行速度会比仅使用CPU 快40倍,提供精度INT8 和FP16 优化,支援TensorFlow、Caffe、Mxnet、Pytorch 等深度学习框架,其中Mxnet、Pytorch 需先转换为ONNX 格式。
trt-info.png
TensorRT 函式库以C++撰写,提供C++ 和Python API。
- TensorRT 的 API 具有 C++ 和 Python 的语言绑定,具有几乎相同的功能。Python API 促进了与 Python 数据处理工具包和库(如 NumPy 和 SciPy)的互操作性,C++ API 可以更高效。
1.TensorRT 的优化方法
1)层间融合或张量融合
TensorRT通过对层间的横向或纵向合并(合并后的结构称为CBR,意指 convolution, bias, and ReLU layers are fused to form a single layer),使得层的数量大大减少。合并之后的计算图的层次更少了,占用的CUDA核心数也少了,因此整个模型结构会更小,更快,更高效。
- 横向合并可以把卷积、偏置和激活层合并成一个CBR结构,只占用一个CUDA核心。
- 纵向合并可以把结构相同,但是权值不同的层合并成一个更宽的层,也只占用一个CUDA核心。
2)数据精度校准
大部分深度学习框架在训练神经网络时网络中的张量都是32位浮点数的精度(Full 32-bit precision,FP32),一旦网络训练完成,在部署推理的过程中由于不需要反向传播,完全可以适当降低数据精度,比如降为FP16或INT8的精度。更低的数据精度将会使得内存占用和延迟更低,模型体积更小。
3)Kernel Auto-Tuning
网络模型在推理计算时,是调用 GPU 的 CUDA 进行计算的,TensorRT 可以针对不同的算法、不同的模型结构、不同的 GPU 平台等,进行 CUDA 调整,以保证当前模型在特定平台上以最优的性能计算。
- 假设在 3090 和 T4 上要分别部署,需要分别在这两个平台上进行 TensorRT 的转换,然后在对应的平台上使用,而不能在相同同的平台上转换,在不同的平台上使用。
4)Dynamic Tensor Memory
在每个 tensor 使用期间,TensorRT 会为其指定显存,避免显存重复申请,减少内存占用和提高重复使用效率。
2.相关软件
-
NVIDIA Triton™推理服务器是一个更高级别的库,可提供跨 CPU 和 GPU 的优化推理。它提供了启动和管理多个模型的功能,以及用于服务推理的 REST 和 gRPC 端点。
-
NVIDIA DALI ®为预处理图像、音频和视频数据提供高性能原语。TensorRT 推理可以作为自定义算子集成到 DALI 管道中。可以在此处找到作为 DALI 的一部分集成的 TensorRT 推理的工作示例。
-
TensorFlow-TensorRT (TF-TRT)是将 TensorRT 直接集成到 TensorFlow 中。它选择 TensorFlow 图的子图由 TensorRT 加速,同时让图的其余部分由 TensorFlow 本地执行。结果仍然是您可以照常执行的 TensorFlow 图。有关 TF-TRT 示例,请参阅TensorFlow 中的 TensorRT 示例。
-
PyTorch 量化工具包提供了以降低精度训练模型的工具,然后可以将其导出以在 TensorRT 中进行优化 。
-
PyTorch Automatic SParsity (ASP)工具提供了用于训练具有结构化稀疏性的模型的工具,然后可以将其导出并允许 TensorRT 在 NVIDIA Ampere GPU 上利用更快的稀疏策略。
-
TensorRT 与 NVIDIA 的分析工具、NVIDIA Nsight™ Systems和NVIDIA® Deep Learning Profiler (DLProf)集成。
二、安装TensorRT及其依赖
1.安装TensorRT
1)查看tensorrt能力支持
根据tensorrt compatibility确定对CUDA、cuDNN、Pytorch和ONNX版本的支持。例如:
[图片上传失败...(image-7469eb-1650335485672)]
2)执行安装
TensorRT 有四种安装方式: 使用Debian, RPM, Tar, Zip 档案,其中Zip 档案只支持Windows。
通过TensorRT官网https://developer.nvidia.com/nvidia-tensorrt-download下载所需的TensorRT版本,以Ubuntu 18.04.5 LTS,CUDA 10.2, cuDNN 7.6.5, python 3.7.9 环境为例:
直接安装到系统
# 要记得换成刚下载的deb 安装包
sudo dpkg -i nv-tensorrt-repo-ubuntu1804-cuda10.2-trt7.1.3.4-ga-20200617_1–1_amd64.deb
sudo apt-key add /var/nv-tensorrt-repo-cuda10.2-trt7.1.3.4-ga-20200617/7fa2af80.pub
sudo apt-get update
sudo apt-get install tensorrt
# 如果所需依赖项因为版本匹配问题导致无法安装,需要指定版本,如:
apt-get install libnvinfer-dev=7.1.3-1+cuda10.2 -y
apt-get install libnvinfer-plugin-dev=7.1.3-1+cuda10.2
apt-get install libnvparsers-dev=7.1.3-1+cuda10.2 -y
apt-get install libnvonnxparsers-dev=7.1.3-1+cuda10.2 -y
apt-get install libnvinfer-samples=7.1.3-1+cuda10.2 -y
# If using Python 2.7:
sudo apt-get install python-libnvinfer-dev
# If using Python 3.x:
sudo apt-get install python3-libnvinfer-dev
# 验证安装
$ dpkg -l | grep TensorRT
通过tar包安装tensorrt
# 下载TensorRT-7.1.3.4.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz后,进行解压
tar -xzvf TensorRT-7.1.3.4.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn8.0.tar.gz
cd TensorRT-7.1.3.4/python
# 安装跟python 版本一样的whl
pip install tensorrt-7.1.3.4-cp37-none-linux_x86_64.whl
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/share/TensorRT-7.1.3.4/lib
# 用于编译C++ tensorrt代码
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/share/TensorRT-7.1.3.4/targets/x86_64-linux-gnu/lib/
export C_INCLUDE_PATH=/home/share/TensorRT-7.1.3.4/include
export CPLUS_INCLUDE_PATH=/home/share/TensorRT-7.1.3.4/include
2.安装pytorch(及cuda)
基于Anaconda,根据所需的cuda版本安装pytorch。
# CUDA 10.2
conda install pytorch torchvision torchaudio cudatoolkit=10.2 -c pytorch
# CUDA 11.3
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
# CPU
conda install pytorch torchvision torchaudio cpuonly -c pytorch
# 查看pytorch版本
$ python
import torch
torch.__version__
'1.8.1+cu102'
torch.version.cuda
10.2
torch.backends.cudnn.version() # cudnn版本
7605
# 查看cuda版本
$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2019 NVIDIA Corporation
Built on Wed_Oct_23_19:24:38_PDT_2019
Cuda compilation tools, release 10.2, V10.2.89
# 检验驱动是否安装成功
$ nvidia-smi -L
GPU 0: GeForce RTX 2080 Ti (UUID: GPU-09b6f910-63d7-e56d-0885-0862fea0bc0c)
GPU 1: GeForce RTX 2080 Ti (UUID: GPU-38d16434-c477-cea8-150a-b88cce8aaacb)
GPU 2: GeForce RTX 2080 Ti (UUID: GPU-94eb0194-e443-7214-b6a5-d430989dd703)
3.安装cudnn
点击cuDNN Archive下载链接,按照要求进行下载。
<img src="https://img-blog.csdnimg.cn/img_convert/7ccaf4652a53f1de1a4c66e571a8492a.png" style="zoom: 67%;" />
# 安装cudnn, 根据实际版本修改
tar -xzvf cudnn-11.2-linux-x64-v8.1.1.33.tgz
sudo cp cuda/include/* /usr/local/cuda/include/
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64/
sudo chmod a+r /usr/local/cuda/include/cudnn.h
sudo chmod a+r /usr/local/cuda/lib64/libcudnn*
# 查看cudnn版本
cat /usr/local/cuda/include/cudnn.h | grep CUDNN_MAJOR -A 2
# 或
$ python
import torch
torch.backends.cudnn.version()
三、TensorRT 样例
1.验证demo
https://blog.csdn.net/irving512/article/details/107165757
进入TensorRT的解压目录
# 数据准备
$ cd TensorRT-7.1.3.4/data/mnist
$ python download_pgms.py
# 下载了一些 *.pgm
$ ls *.pgm
0.pgm 1.pgm 2.pgm 3.pgm 4.pgm 5.pgm 6.pgm 7.pgm 8.pgm 9.pgm
# 代码编译与运行
$ cd TensorRT-7.1.3.4/samples/sampleMNIST
$ make
# 可执行文件生成在 ./bin/ 目录下
$ cd TensorRT-7.1.3.4/
$ ./bin/sample_mnist
[04/12/2022-11:48:05] [I] Building and running a GPU inference engine for MNIST
[04/12/2022-11:48:11] [I] [TRT] Detected 1 inputs and 1 output network tensors.
[04/12/2022-11:48:11] [I] Input:
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@#- -#@@@@@@@@
@@@@@@@@@@@@@@# @@@@@@@@
@@@@@@@@@@@@@#. #@@@@@@@
@@@@@@@@@@@@#. :* +@@@@@@
@@@@@@@@@@@- *: -@@@@@@
@@@@@@@@@@# :+ .%* -@@@@@@
@@@@@@@@@# :@*+@@@ #@@@@@
@@@@@@@@%- .*@@@@@@ -@@@@@
@@@@@@@@: #@%@@@@@@ :@@@@@
@@@@@@@# #@@@@@@@@@ :@@@@@
@@@@@@@: :@@@@@@@@@@ :@@@@@
@@@@@@* +@@@@@@@@@@ =@@@@@
@@@@@@* %@@@@@@@@@= :@@@@@@
@@@@@@* .@@@@@@@@@= .#@@@@@@
@@@@@@* =@@@@@@@#- -@@@@@@@@
@@@@@@* .@@@@@@+ -@@@@@@@@@
@@@@@@* =#%*:. .-#@@@@@@@@@
@@@@@@* .. :=@@@@@@@@@@@
@@@@@@%: =@@@@@@@@@@@@@
@@@@@@@%= =%@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
源码的基本流程
- 第一步:将训练好的神经网络模型转换为TensorRT的形式,并用TensorRT Optimizer进行优化。
- 第二步:将在TensorRT Engine中运行优化好的TensorRT网络结构。
<img src="https://img-blog.csdnimg.cn/img_convert/38ec5fb7f2376f6a2e9fa4b954142cf3.png" style="zoom:67%;" />
2.python demo
$ cd TensorRT-7.1.3.4/samples/python/yolov3_onnx
# 安装依赖
$ pip install -r requirements.txt
# 下载yolov3 权重并转为onnx
$ python yolov3_to_onnx.py
# build TensorRT engine from the generated ONNX file and run inference on a sample image
$ python onnx_to_tensorrt.py
Loading ONNX file from path yolov3.onnx...
Beginning ONNX file parsing
Completed parsing of ONNX file
Building an engine from file yolov3.onnx; this may take a while...
Completed creating Engine
Running inference on image dog.jpg...
[[135.14841345 219.59879372 184.30209058 324.0265199 ]
[ 98.30804939 135.72612864 499.71263113 299.2558099 ]
[478.00607009 81.25702312 210.57786012 86.91503109]] [0.99854713 0.99880403 0.93829261] [16 1 7]
$ python onnx_to_tensorrt.py
Reading engine from file yolov3.trt
Running inference on image dog.jpg...
[[135.14841345 219.59879372 184.30209058 324.0265199 ]
[ 98.30804939 135.72612864 499.71263113 299.2558099 ]
[478.00607009 81.25702312 210.57786012 86.91503109]] [0.99854713 0.99880403 0.93829261] [16 1 7]
Saved image with bounding boxes of detected objects to dog_bboxes.png.
四、pytorch转TensorRT部署
1.ONNX简介
ONNX 是 Open Neural Network Exchange 的简称,也叫开放神经网络交换,是一个用于表示深度学习模型的标准,可使模型在不同框架直接转换。
在深度学习模型落地的过程中,会面临将模型部署到边端设备的问题,模型训练使用不同的框架,则推理的时候也需要使用相同的框架,但不同类型的平台,调优和实现起来非常困难,因为每个平台都有不同的功能和特性。使用ONNX可以通过将不同框架训练的模型转换成通用的 ONNX 模型,再进而转换成各个平台支持的格式,就可以实现简化部署。
ONNX 目前支持的框架有:Caffe2、PyTorch、TensorFlow、MXNet、TensorRT、CNTK 等。
[图片上传失败...(image-7442c4-1650335485672)]
常用部署方案
-
cpu: pytorch/tensorflow/Caffe->onnx->onnxruntime
-
gpu: pytorch/tensorflow/Caffe->onnx->onnx2trt->tensorRT
-
arm: pytorch/tensorflow/Caffe->onnx->ncnn/mace/mnn等
2.PyTorch to ONNX
只要模型中所有OP均被ONNX支持,即可利用Pytorch中的ONN库进行转换。
需要提供的有:加载好的Pytorch模型、一个输入样例。
- 其中模型需要按照自己的方式导入并加载模型;
- 输入样例的格式为BCHW,B为batch_size,CHW为通道、高、宽,CHW的值需要与你自己的模型相匹配,否则后面转换成功后输出结果也不对。
#--*-- coding:utf-8 --*--
import onnx
import torch
import torchvision
import netron
net = torchvision.models.resnet18(pretrained=True).cuda() # pretrained=True:使用模型自带预训练模型
# net.eval()
export_onnx_file = "./resnet18.onnx"
x=torch.onnx.export(net, # 待转换的网络模型和参数
torch.randn(1, 3, 224, 224, device='cuda'), # 虚拟的输入,用于确定输入尺寸和推理计算图每个节点的尺寸
export_onnx_file, # 输出文件的名称
verbose=False, # 是否以字符串的形式显示计算图
input_names=["input"]+ ["params_%d"%i for i in range(120)], # 输入节点的名称,这里也可以给一个list,list中名称分别对应每一层可学习的参数,便于后续查询
output_names=["output"], # 输出节点的名称
opset_version=10, # onnx 支持采用的operator set, 应该和pytorch版本相关,目前我这里最高支持10
do_constant_folding=True, # 是否压缩常量
dynamic_axes={"input":{0: "batch_size", 2: "h"}, "output":{0: "batch_size"},} #设置动态维度,此处指明input节点的第0维度可变,命名为batch_size
)
# import onnx # 注意这里导入onnx时必须在torch导入之前,否则会出现segmentation fault
net = onnx.load("./resnet18.onnx") # 加载onnx 计算图
onnx.checker.check_model(net) # 检查文件模型是否正确
onnx.helper.printable_graph(net.graph) # 输出onnx的计算图
# 使用onnxruntime进行cpu推理onnx
import onnxruntime
import numpy as np
# netron.start("./resnet18.onnx")
session = onnxruntime.InferenceSession("./resnet18.onnx") # 创建一个运行session,类似于tensorflow
out_r = session.run(None, {"input": np.random.rand(16, 3, 256, 224).astype('float32')}) # 模型运行,注意这里的输入必须是numpy类型
print(len(out_r)) # 输出是list类型
print(out_r[0].shape)
- 如果出现“RuntimeError: ONNX export failed: Couldn't export Python operator XXXX”错误提示,说明模型中有ONNX不支持的OP,可以尝试升级Pytorch版本,或者编写自定义op。
- 注意:dynamic_axes可能会导致转tenorrt时失败,谨慎使用。
3.onnxruntime验证onnx模型【可选】
onnxruntime基于cpu推理,以试验转换的onnx模型是否正常。
import cv2
import numpy as np
import onnxruntime
def image_process(image_path):
mean = np.array([[[0.485, 0.456, 0.406]]]) # 训练的时候用来mean和std
std = np.array([[[0.229, 0.224, 0.225]]])
img = cv2.imread(image_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (96, 96)) # (96, 96, 3)
image = img.astype(np.float32)/255.0
image = (image - mean)/ std
image = image.transpose((2, 0, 1)) # (3, 96, 96)
image = image[np.newaxis,:,:,:] # (1, 3, 96, 96)
image = np.array(image, dtype=np.float32)
return image
def onnx_runtime():
imgdata = image_process('test.jpg')
sess = onnxruntime.InferenceSession('test.onnx')
input_name = sess.get_inputs()[0].name
output_name = sess.get_outputs()[0].name
pred_onnx = sess.run([output_name], {input_name: imgdata}) # 定义一路输入,一路输出
# pred_onnx = session.run(None, input_feed={input_name: imgdata}) # 定义一路输入,可多路输出
print("outputs:")
print(np.array(pred_onnx))
onnx_runtime()
4.ONNX to TensorRT
https://github.com/qq995431104/Pytorch2TensorRT
根据生成的ONNX文件,将其转换为TensorRT engine,过程如下:
-
先以trt的Logger为参数,使用TensorRT创建一个builder,然后用builder创建一个network。
with trt.Builder(G_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, \ trt.OnnxParser(network, G_LOGGER) as parser:
-
指定builder的参数设置,如:max_batch_size、max_workspace_size;
- 如需转为特定格式,如fp16或int8,需指定相应参数:fp16_mode或int8_mode设为True;
-
利用对应的Parser(OnnxParser、CaffPaser、UffParser)加载ONNX文件,解析得到网络架构,并填充计算图。
- 由于onnx文件包含模型网络信息,因此并不需要重构网络。
- 当然也可以使用tensorrt的API手动构建网络(c程序wts转engine需要这种方式)。
with open(args.onnx_file_path, 'rb') as model: print('Beginning ONNX file parsing') parser.parse(model.read())
-
builder以network为参数,创建engine。
engine = builder.build_cuda_engine(network)
-
将engine序列化为字符串,写入到engine文件。
with open(args.engine_file_path, "wb") as f: f.write(engine.serialize())
-
可以读取engine文件并反序列化为engine。【可选】
G_LOGGER = trt.Logger(trt.Logger.WARNING) # 反序列化引擎 with open(filepath, "rb") as f, trt.Runtime(G_LOGGER) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) return engine
1)通过python代码转换
import tensorrt as trt
def ONNX2TRT(args, calib=None):
''' convert onnx to tensorrt engine, use mode of ['fp32', 'fp16', 'int8']
:return: trt engine
'''
assert args.mode.lower() in ['fp32', 'fp16', 'int8'], "mode should be in ['fp32', 'fp16', 'int8']"
G_LOGGER = trt.Logger(trt.Logger.WARNING)
# TRT7中的onnx解析器的network,需要指定EXPLICIT_BATCH
EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
with trt.Builder(G_LOGGER) as builder, builder.create_network(EXPLICIT_BATCH) as network, \
trt.OnnxParser(network, G_LOGGER) as parser:
builder.max_batch_size = args.batch_size
builder.max_workspace_size = 1 << 30
if args.mode.lower() == 'int8':
assert (builder.platform_has_fast_int8 == True), "not support int8"
builder.int8_mode = True
builder.int8_calibrator = calib
elif args.mode.lower() == 'fp16':
assert (builder.platform_has_fast_fp16 == True), "not support fp16"
builder.fp16_mode = True
print('Loading ONNX file from path {}...'.format(args.onnx_file_path))
with open(args.onnx_file_path, 'rb') as model:
print('Beginning ONNX file parsing')
if not parser.parse(model.read()):
for e in range(parser.num_errors):
print(parser.get_error(e))
raise TypeError("Parser parse failed.")
print('Completed parsing of ONNX file')
print('Building an engine from file {}; this may take a while...'.format(args.onnx_file_path))
engine = builder.build_cuda_engine(network)
print("Created engine success! ")
# 保存计划文件
print('Saving TRT engine file to path {}...'.format(args.engine_file_path))
with open(args.engine_file_path, "wb") as f:
f.write(engine.serialize())
print('Engine file has already saved to {}!'.format(args.engine_file_path))
return engine
- 对于Int8格式,需要:
- 准备一个校准集,用于在转换过程中寻找使得转换后的激活值分布与原来的FP32类型的激活值分布差异最小的阈值;
- 并写一个校准器类,该类需继承trt.IInt8EntropyCalibrator2父类,并重写get_batch_size, get_batch, read_calibration_cache, write_calibration_cache这几个方法。具体做法参考脚本
myCalibrator.py
. - 使用时,需额外指定cache_file,该参数是校准集cache文件的路径,会在校准过程中生成,方便下一次校准时快速提取。
2)通过trtexec工具转换
cd TensorRT-7.1.3.4/bin
#生成静态batchsize的engine
./trtexec --onnx=<onnx_file> \ #指定onnx模型文件
----maxBatch=1 # max batch size,默认为1
--explicitBatch \ #在构建引擎时使用显式批大小(默认=隐式)显示批处理
--saveEngine=<tensorRT_engine_file> \ #输出engine
--workspace=<size_in_megabytes> \ #设置工作空间大小单位是MB(默认为16MB)
--fp16 #除了fp32之外,还启用fp16精度(默认=禁用)
#生成动态batchsize的engine
./trtexec --onnx=<onnx_file> \ #指定onnx模型文件
--minShapes=input:<shape_of_min_batch> \ #最小的batchsize x 通道数 x 输入尺寸x x 输入尺寸y
--optShapes=input:<shape_of_opt_batch> \ #最佳输入维度,跟maxShapes一样就好
--maxShapes=input:<shape_of_max_batch> \ #最大输入维度
--workspace=<size_in_megabytes> \ #设置工作空间大小单位是MB(默认为16MB)
--saveEngine=<engine_file> \ #输出engine
--fp16
# 示例:
./trtexec --onnx=yolov4_-1_3_416_416_dynamic.onnx \
--minShapes=input:1x3x416x416 \
--optShapes=input:8x3x416x416 \
--maxShapes=input:8x3x416x416 \
--workspace=4096 \
--saveEngine=yolov4_-1_3_416_416_dynamic_b8_fp16.engine \
--fp16
./trtexec --onnx=onnx/mobilenet_v2.onnx --saveEngine=engine/mobilenet_v2_b1_fp16.engine --workspace=4096 --fp16
----------------------------------------------------------------
Input filename: onnx/mobilenet_v2.onnx
ONNX IR version: 0.0.6
Opset version: 9
Producer name: pytorch
Producer version: 1.8
Domain:
Model version: 0
Doc string:
----------------------------------------------------------------
[04/16/2022-08:40:21] [I] Host Latency
[04/16/2022-08:40:21] [I] min: 0.490479 ms (end to end 0.496094 ms)
[04/16/2022-08:40:21] [I] max: 4.66577 ms (end to end 4.79541 ms)
[04/16/2022-08:40:21] [I] mean: 0.50234 ms (end to end 0.748152 ms)
[04/16/2022-08:40:21] [I] median: 0.497772 ms (end to end 0.793213 ms)
[04/16/2022-08:40:21] [I] percentile: 0.538818 ms at 99% (end to end 0.828125 ms at 99%)
[04/16/2022-08:40:21] [I] throughput: 2295.34 qps
[04/16/2022-08:40:21] [I] walltime: 3.00129 s
[04/16/2022-08:40:21] [I] Enqueue Time
[04/16/2022-08:40:21] [I] min: 0.217041 ms
[04/16/2022-08:40:21] [I] max: 4.73901 ms
[04/16/2022-08:40:21] [I] median: 0.243286 ms
[04/16/2022-08:40:21] [I] GPU Compute
[04/16/2022-08:40:21] [I] min: 0.414551 ms
[04/16/2022-08:40:21] [I] max: 4.58044 ms
[04/16/2022-08:40:21] [I] mean: 0.425963 ms
[04/16/2022-08:40:21] [I] median: 0.422241 ms
[04/16/2022-08:40:21] [I] percentile: 0.454712 ms at 99%
[04/16/2022-08:40:21] [I] total compute time: 2.93446 s
- 注意:保存为.trt和.engine文件没有区别
5.TensorRT Python推理
推理过程完全独立于原先模型所依赖的框架,基本过程如下:
-
按照原模型的输入输出格式,准备数据,如:输入的shape、均值、方差,输出的shape等;
-
根据得到的引擎文件,利用TensorRT Runtime反序列化为引擎engine;
-
创建上下文环境;
context = engine.create_execution_context()
-
使用Pycuda的mem_alloc对输入输出分配cuda内存;
d_input = cuda.mem_alloc(1 * input.size * input.dtype.itemsize) d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize) bindings = [int(d_input), int(d_output)]
-
创建Stream,即pycuda操作缓冲区;
stream = cuda.Stream()
-
使用memcpy_htod_async将IO数据放入device(一般为GPU);
cuda.memcpy_htod_async(d_input, input, stream)
-
使用context.execute_async执行推理(异步);
context.execute_async(batch_size, bindings, stream.handle, None)
-
使用memcpy_dtoh_async取出结果:从cuda从缓冲区取出结果并复制到cpu;
cuda.memcpy_dtoh_async(output, d_output, stream)
-
对输出结果进行后处理,因模型而异。
1)直接编写推理程序
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from torchvision import transforms
import numpy as np
from PIL import Image
import time
import argparse
def loadEngine2TensorRT(filepath):
G_LOGGER = trt.Logger(trt.Logger.WARNING)
# 反序列化引擎
with open(filepath, "rb") as f, trt.Runtime(G_LOGGER) as runtime:
engine = runtime.deserialize_cuda_engine(f.read())
return engine
def do_inference(engine, batch_size, input, output_shape):
# 创建上下文
context = engine.create_execution_context()
output = np.empty(output_shape, dtype=np.float32)
# 分配内存
d_input = cuda.mem_alloc(1 * input.size * input.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)
bindings = [int(d_input), int(d_output)]
# pycuda操作缓冲区
stream = cuda.Stream()
# 将输入数据放入device
cuda.memcpy_htod_async(d_input, input, stream)
start = time.time()
# 执行模型
context.execute_async(batch_size, bindings, stream.handle, None)
# 将预测结果从从缓冲区取出
cuda.memcpy_dtoh_async(output, d_output, stream)
end = time.time()
# 线程同步
stream.synchronize()
print("\nTensorRT {} test:".format(engine_path.split('/')[-1].split('.')[0]))
print("output:", output)
print("time cost:", end - start)
def get_shape(engine):
for binding in engine:
if engine.binding_is_input(binding):
input_shape = engine.get_binding_shape(binding)
else:
output_shape = engine.get_binding_shape(binding)
return input_shape, output_shape
if __name__ == "__main__":
parser = argparse.ArgumentParser(description = "TensorRT do inference")
parser.add_argument("--batch_size", type=int, default=1, help='batch_size')
parser.add_argument("--img_path", type=str, default='test_image/1.jpg', help='cache_file')
parser.add_argument("--engine_file_path", type=str, default='my_files/test.engine', help='engine_file_path')
args = parser.parse_args()
engine_path = args.engine_file_path
engine = loadEngine2TensorRT(engine_path)
img = Image.open(args.img_path)
input_shape, output_shape = get_shape(engine)
transform = transforms.Compose([
transforms.Resize([input_shape[1], input_shape[2]]), # [h,w]
transforms.ToTensor()
])
img = transform(img).unsqueeze(0)
img = img.numpy()
do_inference(engine, args.batch_size, img, output_shape)
2)使用官方库torch2trt进行推理
https://github.com/NVIDIA-AI-IOT/torch2trt
def torch2trt(module,
inputs,
input_names=None,
output_names=None,
log_level=trt.Logger.ERROR,
max_batch_size=1,
fp16_mode=False,
max_workspace_size=1<<25,
strict_type_constraints=False,
keep_network=True,
int8_mode=False,
int8_calib_dataset=None,
int8_calib_algorithm=DEFAULT_CALIBRATION_ALGORITHM,
int8_calib_batch_size=1,
use_onnx=False,
**kwargs):
import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet
# create some regular pytorch model...
model = alexnet(pretrained=True).eval().cuda()
# create example data
x = torch.ones((1, 3, 224, 224)).cuda()
# convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])
优点
- 使用简单,对于Mobilenet、Unet、Resnet等可以直接使用
缺点
- 不支持多输出
- 自定义的tensor或者list都没有_trt这个属性
6.TensorRT C++ 推理
把模型部署到内存有限的嵌入式板的过程:
- 电脑上安装的有anaconda, pytorch等,但是在电脑上转的不能直接在板子上用。
- 板子的内存有限,不能安装anaconda, pytorch这些,但是需要部署模型上去。这时就可以现在电脑上把pth转成wts,再把wts传到板子上,在板子上转成tensorrt。
1)pth转wts
import torch
import torch.nn as nn
from torchvision import models
import struct
from torchsummary import summary
def get_model():
net = getattr(models.quantization, 'mobilenet_v2')(pretrained=False, num_classes=2, quantize=False)
net.load_state_dict(torch.load('weights/xxx.pth'))
net = net.eval().cuda()
return net
def pth_to_wts(model, wts_name):
f = open(wts_name, 'w')
f.write('{}\n'.format(len(model.state_dict().keys())))
for k, v in model.state_dict().items():
vr = v.reshape(-1).cpu().numpy()
f.write('{} {} '.format(k, len(vr)))
for vv in vr:
f.write(' ')
f.write(struct.pack('>f',float(vv)).hex())
f.write('\n')
if __name__ == '__main__':
model = get_model()
summary(model, (3, 256, 256)) # 可选,将模型视觉化,了解模型每一层输入输出
pth_to_wts(model, "wts/mobilenet_v2.wts")
2)wts转tensorrt
关键头文件
-
#include "cuda_runtime_api.h"
- 路径:/usr/local/cuda-10.2/targets/x86_64-linux/include/cuda_runtime_api.h
-
#include "NvInfer.h"
- 本地路径:/usr/include/x86_64-linux-gnu/NvInfer.h,在线API地址
- 提供IRuntime、IBuilder、IHostMemory、IExecutionContext、ICudaEngine和英伟达gpu 算子等调用接口。
常用的模型转换可参考Tensorrt C API。
wts转tensorrt的原理
- 从wts文件把weight给load出来,存到一个map里,key是网络每层的名称,value就是对应的权重
- 利用tensorrt的API把网络重建出来,同时导入key对应的value,也就是weightMap的形式
- 定义网络的输出,设置内存空间
- build engine输出一个engine文件
wts转tensorrt示例程序
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";
// Load weights from files shared with TensorRT samples.
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
// Open weights file
std::ifstream input(file);
assert(input.is_open() && "Unable to load weight file.");
// Read number of weight blobs
int32_t count;
input >> count;
assert(count > 0 && "Invalid weight map file.");
while (count--)
{
Weights wt{DataType::kFLOAT, nullptr, 0};
uint32_t size;
// Read name and type of blob
std::string name;
input >> name >> std::dec >> size;
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
for (uint32_t x = 0, y = size; x < y; ++x)
{
input >> std::hex >> val[x];
}
wt.values = val;
wt.count = size;
weightMap[name] = wt;
}
return weightMap;
}
// Creat the engine using only the API and not any parser. 重建网络
ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt)
{
INetworkDefinition* network = builder->createNetworkV2(0U);
// Create input tensor of shape { 3, INPUT_H, INPUT_W } with name INPUT_BLOB_NAME
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W}); // 定义输入层
assert(data);
std::map<std::string, Weights> weightMap = loadWeights("../mobilenet.wts");
Weights emptywts{DataType::kFLOAT, nullptr, 0};
// 输入传入各卷积层
auto ew1 = convBnRelu(network, weightMap, *data, 32, 3, 2, 1, "features.0.");
ILayer* ir1 = invertedRes(network, weightMap, *ew1->getOutput(0), "features.1.", 32, 16, 1, 1);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.2.", 16, 24, 2, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.3.", 24, 24, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.4.", 24, 32, 2, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.5.", 32, 32, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.6.", 32, 32, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.7.", 32, 64, 2, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.8.", 64, 64, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.9.", 64, 64, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.10.", 64, 64, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.11.", 64, 96, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.12.", 96, 96, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.13.", 96, 96, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.14.", 96, 160, 2, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.15.", 160, 160, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.16.", 160, 160, 1, 6);
ir1 = invertedRes(network, weightMap, *ir1->getOutput(0), "features.17.", 160, 320, 1, 6);
IElementWiseLayer* ew2 = convBnRelu(network, weightMap, *ir1->getOutput(0), 1280, 1, 1, 1, "features.18.");
// 池化
IPoolingLayer* pool1 = network->addPoolingNd(*ew2->getOutput(0), PoolingType::kAVERAGE, DimsHW{7, 7});
assert(pool1);
// 全连接
IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool1->getOutput(0), 1000, weightMap["classifier.1.weight"], weightMap["classifier.1.bias"]);
assert(fc1);
fc1->getOutput(0)->setName(OUTPUT_BLOB_NAME); // 定义输出
std::cout << "set name out" << std::endl;
network->markOutput(*fc1->getOutput(0)); // 指定网络输出内存空间
// Build engine
builder->setMaxBatchSize(maxBatchSize); // 设置一个engine内存空间
config->setMaxWorkspaceSize(1 << 20);
ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
std::cout << "build out" << std::endl;
// Don't need the network any more
network->destroy();
// Release host memory
for (auto& mem : weightMap)
{
free((void*) (mem.second.values));
}
return engine;
}
void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream)
{
// Create builder
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
// Create model to populate the network, then set the outputs and create an engine
ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT);
assert(engine != nullptr);
// Serialize the engine
(*modelStream) = engine->serialize();
// Close everything down
engine->destroy();
builder->destroy();
config->destroy();
}
int main(int argc, char** argv)
{
if (std::string(argv[1]) == "-s") {
IHostMemory* modelStream{nullptr};
APIToModel(1, &modelStream); // 加载wts文件,利用tensorrt的API把网络重建,并转为engine
assert(modelStream != nullptr);
std::ofstream p("mobilenet.engine", std::ios::binary);
if (!p)
{
std::cerr << "could not open plan output file" << std::endl;
return -1;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size()); // 写入engine数据到文件
modelStream->destroy();
return 1;
}
}
3)engine推理
void doInference(IExecutionContext& context, float* input, float* output, int batchSize)
{
const ICudaEngine& engine = context.getEngine();
// Pointers to input and output device buffers to pass to engine.
// Engine requires exactly IEngine::getNbBindings() number of buffers.
assert(engine.getNbBindings() == 2);
void* buffers[2];
// In order to bind the buffers, we need to know the names of the input and output tensors.
// Note that indices are guaranteed to be less than IEngine::getNbBindings()
const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME);
// Create GPU buffers on device
CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float)));
CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float)));
// Create stream
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host
CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream));
context.enqueue(batchSize, buffers, stream, nullptr); // 上下文将数据入列,执行推理
CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); // 将预测结果从从缓冲区取出
cudaStreamSynchronize(stream); // 数据同步
// Release stream and buffers
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
}
int main(int argc, char** argv)
{
char *trtModelStream{nullptr};
if (std::string(argv[1]) == "-d") {
std::ifstream file("mobilenet.engine", std::ios::binary);
if (file.good()) {
file.seekg(0, file.end);
size = file.tellg();
file.seekg(0, file.beg);
trtModelStream = new char[size];
assert(trtModelStream);
file.read(trtModelStream, size); // 读取engine文件到trtModelStream
file.close();
}
}
// trtModelStream转ICudaEngine,并创建上下文
IRuntime* runtime = createInferRuntime(gLogger);
assert(runtime != nullptr);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr);
assert(engine != nullptr);
IExecutionContext* context = engine->createExecutionContext();
assert(context != nullptr);
delete[] trtModelStream;
// Subtract mean from image
static float data[3 * INPUT_H * INPUT_W];
for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++)
data[i] = 1.0;
// Run inference
static float prob[OUTPUT_SIZE];
doInference(*context, data, prob, 1); // 构建图像进行推理
}
4)tensorrtx
tensorrtx提供多种模型的C版本的的wts转engine和gpu推理程序, github地址:https://github.com/wang-xinyu/tensorrtx。
tensorrtx工程目录
$ git clone https://github.com/wang-xinyu/tensorrtx.git
$ tree -L 1 tensorrtx
├── alexnet
├── arcface
├── crnn
├── dbnet
├── Dockerfile
├── googlenet
├── hrnet
├── inceptionv3
├── lenet
├── LICENSE
├── mnasnet
├── mobilenetv2
├── mobilenetv3
├── psenet
├── README.md
├── resnet
├── retinaface
├── retinafaceAntiCov
├── senet
├── shufflenetv2
├── squeezenet
├── tutorials
├── ufld
├── vgg
├── yolov3
├── yolov3-spp
├── yolov3-tiny
├── yolov4
└── yolov5