TensorRT学习笔记（一）——使用Python API调用TensorRT

前言

作为在英伟达自家GPU上的推理库，TensoRT仍然是使用英伟达显卡部署方案的最佳选择。TensorRT虽然支持Pytho和C++调用，但是在TensorRT8之前，python api只能在linux上使用，直到TensorRT8才支持python api在window下使用。
具体安装在官方文档都有说明，附上官方文档：TensorRT官方文档

参考文章：

1. TensorRT 7.2.2 下载、安装、使用（ubuntu 18.04, cuda 10.2, cudnn 8.0.5
2. ImportError: libnvinfer.so.7: cannot open shared object file
3. TensorRT使用教程（Python版）_ZhouJianGuos的博客-程序员宝宝_tensorrt使用教程
4. 【踩坑实录】TensorRT 傻瓜式部署流程

安装

首先需要进入tensorrt 官网： https://developer.nvidia.com/tensorrt
点击“立即下载”，需要英伟达账号，可以选择任意平台适配你的cuda版本的TensorRT.
注意：安装tensorrt需要先安装好配套的cuda和cudnn。请根据自己的硬件配置自行安装cuda和cudnn

1. 下载并解压tar文件

tar -zxvf TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz

2. 进入到解压路径下的python文件夹，使用pip安装

cd TensorRT-8.4.1.5/python
# 选择自己的python版本进行安装
pip3 install tensorrt-8.4.1.5-cp36-none-linux_x86_64.whl

3. 安装uff以支持tensorflow

cd TensorRT-8.4.1.5/uff
pip install uff-0.6.9-py2.py3-none-any.whl

4. 安装graphsurgeon以支持自定义结构

cd TensorRT-8.4.1.5/graphsurgeon
pip install graphsurgeon-0.4.6-py2.py3-none-any.whl

5. onnx_graphsurgeon以支持onnx

cd TensorRT-8.4.1.5/onnx_grahsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

6. 修改环境变量
   把刚才解压的TensorRT文件夹路径写进 ~/.bashrc中

#使用以下命令打开.bashrc文件
sudo gedit ~/.bashrc

#解压得到TensorRT-7.2.1.6的文件夹，将里边的lib绝对路径添加到环境变量中，用户名和TensorRT的版本号更换为自己的
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/youpath/TensorRT-8.4.1.5/lib

#保存文件并退出，执行以下命令，使路径生效
source ~/.bashrc

如果这个时候发现在终端调用tensorrt没有问题，但在pycharm里运行时，提示找不到libnvinfer.so.7（或者缺少其他so包，这些包能在lib下找到）。虽然在bashrc中设置了环境变量，但是也没有对pycharm起到作用，出现如下报错：

Traceback (most recent call last):
File "/media/lindsay/591d8930-456f-45c2-9059-4e020abd1169/tensorrtx/yolov5/streams_trt.py", line 15, in <module>
  import tensorrt as trt
File "/home/lindsay/anaconda3/envs/yolov5/lib/python3.8/site-packages/tensorrt/__init__.py", line 66, in <module>
  from .tensorrt import *
ImportError: libnvinfer.so.7: cannot open shared object file: No such file or directory

可以参考这篇文章：ImportError: libnvinfer.so.7: cannot open shared object file 在pycharm进行设置。

Python API 调用

1. 导出ONNX文件（Pytorch为例）

 # 构造模型实例
 model = HRNet()
 # 反序列化权重参数
 model.load_state_dict(torch.load(self.weight_path),strict=False)
 model.eval()
 # 定义输入名称，list结构，可能有多个输入
 input_names = ['input']
 # 定义输出名称，list结构，可能有多个输出
 output_names = ['output']
 # 构造输入用以验证onnx模型的正确性
 input = torch.rand(1, 3, self.resolution[0], self.resolution[1])
 # 导出
 torch.onnx.export(model, input, output_path,
                          export_params=True,
                          opset_version=11,
                          do_constant_folding=True,
                          input_names=input_names,
                          output_names=output_names)

参数介绍：

• model为pytorch模型实例
• input为测试输入数据（形状必须要和模型输入一致，但是数值可以是随机的）
• output_path为导出路径，xxx.onnx
• export_params为是否导出参数权重，必然是True
• opset_version=11 发行版本，11就可以了
• do_constant_folding是否对常量进行折叠，True就可以了
• input_names是模型输入的名称，list类型，因为模型可能有多个输入
• output_names同上，只不过这是针对输出的

2. trtexec.exe实现预推理

打开tensorrt安装包的子目录bin，里面有一个trtexec文件，trtexec 是 TensorRT sample 里的一个例子，时把 TensorRT 许多方法包装成了一个可执行文件。它可以把模型优化成 TensorRT Engine ，并且填入随机数跑 inference 进行速度测试。
输入以下指令：

trtexec --onnx=xxx.onnx --saveEngine=xxx.engine

这个命令把 onnx 模型优化成 Engine ，然后多次 inference 后统计并报时。
报时会报很多个， Enqueue Time 是 GPU 任务排队的时间，H2D Latency 是把网络输入数据从主存 host 拷贝进显存 device 的时间，D2H Latency 是把显存上的网络输出拷贝回主存的时间，只有 GPU Compute Time 是真正的网络 inference 时间。

一般来说模型的第一维大小是任意的（batch size 维度），而 TensorRT 不能把任意 batch size 都加速。可以指定一个输入范围，并且重点优化其中一个 batch size。例如网络输入格式是 [-1, 3, 244, 244]：

./trtexec --onnx=model.onnx --minShapes=input:1x3x244x244 --optShapes=input:16x3x244x244 --maxShapes=input:32x3x244x244 --shapes=input:5x3x244x244

还可以降低精度优化速度。一般来说大家写的模型内都是 float32 的运算，TensorRT 会默认开启 TF32 数据格式，它是截短版本的 FP32，只有 19 bit，保持了 fp16 的精度和 fp32 的指数范围。

另外，TensorRT 可以额外指定精度，把模型内的计算转换成 float16 或者 int8 的类型，可以只开一个也可以两个都开，trtexec 会倾向于速度最快的方式（有些网络模块不支持 int8）

./trtexec --onnx=model.onnx --saveEngine=xxx.trt --int8 --fp16

一般来说，只开 fp16 可以把速度提一倍并且几乎不损失精度；但是开 --int8 会大大损失精度，速度会比 fp16 快，但不一定能快一倍。

int8 优化涉及模型量化，需要校准（calibrate）提升精度。TensorRT 有两种量化方法：训练后量化和训练中量化。二者的校准方法不同，精度也不同，后者更高一些。具体参考 Developer Guide :: NVIDIA Deep Learning TensorRT Documentation。

trtexec 采用的是训练后量化，写起来更方便一些。不过看看源码就能发现，因为 trtexec 只为了测试速度，所以校准就象征性做了一下，真想自己部署 int8 模型还得自己写校准。

常用参数解释

• – onnx onnx路径
• – saveEngine执行计划（推理引擎）序列化地址
• – fp16开启 float16精度的推理（推荐此模式，一方面能够加速，另一方面精度下降比较小）
• – int8 开启 int8精度的推理（不太推荐，虽然更快，但是精度下降太厉害了）

使用 TensorRT Python API

简单版本

这个版本借助了torch2trt这个库，已经把推理代码的细节封装好了。这个仓库可以实现TensorRT和Pytorch的链接，具体信息请看仓库的更多介绍：https://github.com/NVIDIA-AI-IOT/torch2trt

import tensorrt as trt
from torch2trt import TRTModule

# 加载日志记录器
logger = trt.Logger(trt.Logger.INFO)
# 加载engine
with open("xxx.trt", "rb") as f, trt.Runtime(logger) as runtime:
  engine=runtime.deserialize_cuda_engine(f.read())

net = TRTModule(engine, input_names=["input"], output_names=['output'])

# 推理
image = cv2.imread(image_path)
image = cv2.resize(image, (224,224))
image = image.transpose(2,0,1)
img_input = image.astype(np.float32)
img_input = torch.from_numpy(img_input)
img_input = img_input.unsqueeze(0)
img_input = img_input.to(device)
# 运行模型
result_trt = trt_model(img_input)

整个流程都跟pytorch流程差不多，输入进去给网络的也是pytorch的张量，得到的结果也是pytorch的张量。

原始版本

#导入必用依赖
import tensorrt as trt
import pycuda.autoinit  #负责数据初始化，内存管理，销毁等
import pycuda.driver as cuda  #GPU CPU之间的数据传输

#创建logger：日志记录器
logger = trt.Logger(trt.Logger.WARNING)
#创建runtime并反序列化生成engine
with open(“sample.engine”, “rb”) as f, trt.Runtime(logger) as runtime:
    engine = runtime.deserialize_cuda_engine(f.read())

# 分配空间并绑定输入输出
host_inputs = []
cuda_inputs = []
host_outputs = []
cuda_outputs = []
bindings = []
for binding in engine:
    size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
    dtype = trt.nptype(engine.get_binding_dtype(binding))
    # 分配主机和设备buffers
    host_mem = cuda.pagelocked_empty(size, dtype)    
    cuda_mem = cuda.mem_alloc(host_mem.nbytes)      
    # 将设备buffer绑定到设备.
    bindings.append(int(cuda_mem))
    # 绑定到输入输出
    if engine.binding_is_input(binding):
         host_inputs.append(host_mem)           
         cuda_inputs.append(cuda_mem)    
    else:
         host_outputs.append(host_mem)
         cuda_outputs.append(cuda_mem)

#创建cuda流
stream = cuda.Stream()

# 拷贝输入图像到主机buffer
np.copyto(host_inputs[0], input_image.ravel())
# 将输入数据转到GPU.
cuda.memcpy_htod_async(cuda_inputs[0], host_inputs[0], stream)
# 推理.
context.execute_async(bindings=bindings, stream_handle=stream.handle)
# 将推理结果传到CPU.
cuda.memcpy_dtoh_async(host_outputs[0], cuda_outputs[0], stream)
# 同步 stream
stream.synchronize()
# 拿到推理结果 batch_size = 1
output = host_outputs[0]