pt模型转onnx模型,onnx模型转engine模型,pt模型转engine模型详细教程(TensorRT,jetpack)

背景

背景是需要在nvidia jetpack4.5.1的arm64设备上跑yolov8,用TensorRT加速,需要用*.engine格式的模型,但是手头上的是pt格式模型,众所周知小板子的内存都很小,连安装ultralytics依赖库的容量都没有,所以我想到在windows系统上先把模型转好,然后再直接放在板子上用。

下载ultralytics yolov8

https://github.com/ultralytics/ultralytics

找到exporter.py

我以为是直接在这个exporter.py上直接改参数来转模型,找了半天找不到,然后直接运行这个python文件也是失败报错,后来才看到这个exporter.py上面有教如何使用

  • 安装依赖库
bash 复制代码
pip install "ultralytics[export]

等待安装好,这需要很长时间

然后不是直接运行这个exporter.py文件,而是像提示里面一样,创造一个py文件,在文件里面复制这几行然后改参数就行了

python 复制代码
    from ultralytics import YOLO
    model = YOLO('yolov8n.pt')
    results = model.export(format='onnx')

结果我运行了会报错找不到gpu

bash 复制代码
assert self.im.device.type != 'cpu', "export running on CPU but must be on GPU, i.e. use 'device=0'"
AssertionError: export running on CPU but must be on GPU, i.e. use 'device=0

查了是需要指定gpu,然后我尝试在python里指定gpu,都还是报这个错,突然看到有提示需要device=0,我想这是加在哪里的呢,后来发现是在下面那个cli命令行加

尝试一步到位转成engine模型

bash 复制代码
yolo mode=export model=Gap.pt format=engine

报错assert self.im.device.type != 'cpu', "export running on CPU but must be on GPU, i.e. use 'device=0'"

bash 复制代码
assert self.im.device.type != 'cpu', "export running on CPU but must be on GPU, i.e. use 'device=0'"

解决办法就是在后面加一个device=0

bash 复制代码
yolo mode=export model=Gap.pt format=engine device=0  

把模型放在小板子上用结果报错ERROR: coreReadArchive.cpp (31) - Serialization Error in verifyHeader: 0 (Magic tag does not match)

查了下,我真是慌了,意思是编译engine时候的使用的tensorrt版本与使用trt推理时候的tensorrt版本不一致,需要一致才行

,排查下各自是什么版本

  • linux查tensorRT版本
bash 复制代码
ldd 可执行文件

可以看到生成可执行文件的tensorRT是7版本

这个命令行也可以查询到

bash 复制代码
dpkg -l | grep TensorRT

能看到是7.1.3版本的tensorRT

  • windows查tensorRT版本
python 复制代码
import tensorrt as trt
print(f'TensorRT version: {trt.__version__}')

可以看到我在编译engine时候的使用的tensorrt版本是8.4,版本不匹配,所以trt推理失败。

然后我寻思arm64环境是不方便更改的,那我就降低windows端的tensorRT版本,我就去nvidia官网下载(https://developer.nvidia.com/tensorrt/download)结果,官网都不提供7.1版本的tensorRT下载了,一下子又不知道该怎么解决了。

我突然想到既然直接使用engine不行,那我用生成出来的onnx模型再在arm64上转成engine行不行

aarch64上onnx转engine模型

找到你本机上的trt转模型可执行文件

bash 复制代码
sudo find / -name trtexec

复制到你的模型文件夹

bash 复制代码
./trtexec --onnx=/home/adlink/wjp/YoloV8-TensorRT-Jetson_Nano-main/models/yolov8n.onnx  --saveEngine=yolov8n.engine --fp16=fp32

结果还是报错volume mismatch,Input dimensions [1,33,8400] have volume 277200 and output dimensions [1,4,16,8400] have volume 537600.

仿佛又没有办法解决了,难道非要统一版本,但是这上面尝试过已知道是非常困难的,我尝试用其他办法,后来尝试更换参数,发现是opset版本太高的问题,好像默认是17,而aarch64上的版本很低,需要在yolo转pt为onnx的步骤加指定opset的参数

bash 复制代码
yolo export model=yolov8s.pt format=onnx opset=11 simplify=True

再把生成的onnx文件用trtexec转

bash 复制代码
./trtexec --onnx=/home/adlink/wjp/YoloV8-TensorRT-Jetson_Nano-main/models/yolov8n.onnx  --saveEngine=yolov8n.engine --fp16=fp32

就转成功了,跑模型也没问题!

特别说明下

bash 复制代码
[TRT] Some tactics do not have sufficient workspace memory to run. Increasing workspace size may increase performance, please check verbose output.

中间会在这个步骤停很长时间,这个不是报错,静静等待就好了,过一会就能成功转换onnx为engine模型了

相关推荐
BoBoZz197 小时前
Glyph2D 同一个图形根据点云的输入产生不同位置的输出
python·vtk·图形渲染·图形处理
一笑code8 小时前
pycharm vs vscode安装python的插件
vscode·python·pycharm
liwulin05068 小时前
【PYTHON-YOLOV8N】yoloface+pytorch+cnn进行面部表情识别
python·yolo·cnn
(●—●)橘子……8 小时前
记力扣1471.数组中的k个最强值 练习理解
数据结构·python·学习·算法·leetcode
_OP_CHEN8 小时前
用极狐 CodeRider-Kilo 开发俄罗斯方块:AI 辅助编程的沉浸式体验
人工智能·vscode·python·ai编程·ai编程插件·coderider-kilo
Wpa.wk8 小时前
自动化测试 - 文件上传 和 弹窗处理
开发语言·javascript·自动化测试·经验分享·爬虫·python·selenium
_OP_CHEN8 小时前
【Python基础】(二)从 0 到 1 入门 Python 语法基础:从表达式到运算符的全面指南
开发语言·python
我命由我123458 小时前
Python Flask 开发:在 Flask 中返回字符串时,浏览器将其作为 HTML 解析
服务器·开发语言·后端·python·flask·html·学习方法
拾忆,想起8 小时前
设计模式:软件开发的可复用武功秘籍
开发语言·python·算法·微服务·设计模式·性能优化·服务发现
沃斯堡&蓝鸟9 小时前
DAY28 元组和OS模块
python·元组与os模块