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
  1. 进入到解压路径下的python文件夹,使用pip安装
cd TensorRT-8.4.1.5/python
# 选择自己的python版本进行安装
pip3 install tensorrt-8.4.1.5-cp36-none-linux_x86_64.whl
  1. 安装uff以支持tensorflow
cd TensorRT-8.4.1.5/uff
pip install uff-0.6.9-py2.py3-none-any.whl
  1. 安装graphsurgeon以支持自定义结构
cd TensorRT-8.4.1.5/graphsurgeon
pip install graphsurgeon-0.4.6-py2.py3-none-any.whl
  1. onnx_graphsurgeon以支持onnx
cd TensorRT-8.4.1.5/onnx_grahsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl
  1. 修改环境变量
    把刚才解压的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]
  • 5
    点赞
  • 33
    收藏
    觉得还不错? 一键收藏
  • 4
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 4
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值