pytorch yolov7 模型转tensorRT

无聊时看看书

已于 2023-08-21 16:25:59 修改

阅读量414

点赞数

CC 4.0 BY-SA版权

文章标签： linux ubuntu YOLO pytorch

于 2023-08-21 16:25:21 首次发布

本文链接：https://blog.youkuaiyun.com/qq_33554090/article/details/130868640

TensorRT安装

下载

下载地址 https://developer.nvidia.com/nvidia-tensorrt-download

选择需要的版本，点击进入：

勾选同意。。。选择具体的版本

选择下载8.4.2.4 版本的tar.gz安装包：

查看下载的安装包：

安装

解压安装包：

tar -xzvf TensorRT-8.4.2.4.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz

设置环境变量：vi ~/.bashrc,写入下面设置

export LD_LIBRARY_PATH=/path/TensorRT-8.4.2.4/lib:$LD_LIBRARY_PATH
export LIBRARY_PATH=/path/TensorRT-8.4.2.4/lib::$LIBRARY_PATH

source ~/.bashrc

验证TensorRT是否安装正确，cd进入samples文件夹：

cd samples

make -j8

如果环境什么都没出错的话，这里编译是会很顺利的。然后cd进入bin文件夹，

cd ..
cd bin
./sample_onnx_mnist

安装python包（可选）

cd /path/TensorRT-8.4.2.4/python

根据自己需要安装对应python的安装包，例如：

pip install tensorrt-8.4.2.4-cp37-none-linux_x86_64.whl

还需要安装pycuda

pip install pycuda==2019.1.1

检查python中tensorrt是否安装成功

>>> import tensorrt as tr
>>> tr.__version__
'8.4.2.4'

模型转换

pt -> onnx

安装依赖包：

pip install onnx==1.14.0

pip install nvidia-pyindex

pip install onnx_graphsurgeon

pip install coremltools

转换

git clone https://github.com/WongKinYiu/yolov7.git

cd yolov7

# 640
python export.py --weights yolov7-tiny.pt  --dynamic  --grid
python export.py --weights yolov7.pt  --dynamic  --grid
python export.py --weights yolov7x.pt  --dynamic  --grid
# 1280
python export.py --weights yolov7-w6.pt  --dynamic  --grid --img-size 1280

onnx - > trt

git clone https://github.com/FeiYull/TensorRT-Alpha.git

cd TensorRT-Alpha/data/yolov7

# 640
/path/TensorRT-8.4.2.4/bin/trtexec    --onnx=yolov7-tiny.onnx  --saveEngine=yolov7-tiny.trt  --buildOnly --minShapes=images:1x3x640x640 --optShapes=images:4x3x640x640 --maxShapes=images:8x3x640x640
/path/TensorRT-8.4.2.4/bin/trtexec    --onnx=yolov7.onnx   	--saveEngine=yolov7.trt       --buildOnly --minShapes=images:1x3x640x640 --optShapes=images:4x3x640x640 --maxShapes=images:8x3x640x640
/path/TensorRT-8.4.2.4/bin/trtexec    --onnx=yolov7x.onnx   	--saveEngine=yolov7x.trt      --buildOnly --minShapes=images:1x3x640x640 --optShapes=images:4x3x640x640 --maxShapes=images:8x3x640x640
# 1280
/path/TensorRT-8.4.2.4/bin/trtexec    --onnx=yolov7-w6.onnx    --saveEngine=yolov7-w6.trt    --buildOnly --minShapes=images:1x3x1280x1280 --optShapes=images:4x3x1280x1280 --maxShapes=images:8x3x1280x1280

# note:if report an error（Error Code 1: Cuda Runtime (an illegal memory access was encountered "bool context = m_context->executeV2((void**)bindings)" returns false） 
when running the model(yolov7-w6), just lower the batch_size.

注意：若是报错so包找不到：

cd /path/TensorRT-8.4.2.4/bin/

ldd trtexec

查看so是否没有找不到，若so包正常，说明export环境变量没生效，请 source ~/.bashrc，或者手动执行

export LD_LIBRARY_PATH=/path/TensorRT-8.4.2.4/lib:$LD_LIBRARY_PATH
export LIBRARY_PATH=/path/TensorRT-8.4.2.4/lib::$LIBRARY_PATH

转换成功