1. 导出onnx
经过前面训练,得到了best.pt模型,现在想要使用tensorrt进行推理,需要先导出为onnx格式,再转化为engine格式。
yolo export model=best.pt format=onnx opset=12 simplify=True
2.解决错误
在导出过程中,可能会出现错误,
- cmake版本问题
安装好后,默认cmake版本为3.16,需要对cmake进行升级
sudo pip install --upgrade cmake==3.24.0 -i https://pypi.douban.com/simple
hash -r
使用cmake --version 查看最新版本号
- 没有onnxruntime、onnxsim问题
使用以下命令安装
pip3 install onnxsim onnxruntime -i https://pypi.douban.com/simple
3.使用gpu进行预测
yolo task=pose mode=predict model=/home/nvidia/best.onnx source=0 show=True save=True device=0
如果出现错误
onnxruntime_inference_collection.py:69: UserWarning: Specified provider 'CUDAExecutionProvider' is not in available provider names.Available providers: 'AzureExecutionProvider, CPUExecutionProvider'
warnings.warn(
说明没有安装onnxruntime-gpu,无法使用gpu进行预测。
解决办法:
在jetson orin nano中,先卸载onnruntime,
pip3 uninstall onnxruntime
然后在https://elinux.org/Jetson_Zoo,根据自己的jetpack版本、cuda版本、python版本,选择对应的onnxruntime-gpu,
例如,我的时jetpack5.1.1,cuda是11.4,python是3.8,所以选择onnxruntime1.12.1,
在官网https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html查看cuda对应的onnruntime-gpu版本。
下载完成后,使用命令安装
pip3 install onnxruntime_gpu-1.12.1-cp38-cp38-linux_aarch64.whl
4、导出为engine
首先设置tensorrt路径,打开~/.bashrc,添加
export PATH=/usr/src/tensorrt/bin:$PATH
到文件末尾,执行
source ~/.bashrc
更新后,可以执行
trtexec --onnx=/home/nvidia/best.onnx --saveEngine=/home/nvidia/best.engine
等待5-10分钟后,转换完毕