ubuntu18下TensorRT加速YOLOV5推理(超详细)

安装TensorRT

TensorRT是Nvidia官方给的C++推理加速工具,如同OpenVINO之于Intel。支持诸多的AI框架,如Tensorflow,Pytorch,Caffe,MXNet等。

可以通过以下链接进行TensorRT的下载安装,根据自己的cuda版本安装相对应的版本。

下载连接:https://developer.nvidia.com/tensorrt

下载这个需要注册nvidia账号,按要求注册就好了。
下载完成后解压:

tar -xvzf TensorRT-6.0.1.5.Ubuntu-16.04.x86_64-gnu.cuda-10.1.cudnn7.6.tar.gz

进行export:

export TRT_RELEASE=`pwd`/TensorRT-6.0.1.5
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$TRT_RELEASE/lib

安装Python模块:

cd  TensorRT-6.0.1.5/python/
pip install tensorrt-6.0.1.5-cp37-none-linux_x86_64.whl

安装完毕!

部署YOLOV5-tensorrt

Yolov5使用TensorRT的过程,实际就是在其API下复现yolov5的过程。已经有大佬实现了并且在github上开源了项目,大家可以直接下载使用了。

项目地址:https://github.com/wang-xinyu/tensorrtx

(1)下载yolov5和 yolov5 tensorrt:

git clone https://github.com/wang-xinyu/tensorrtx.git
git clone https://github.com/ultralytics/yolov5.git

(2)将tensorrt/yolov5/gen_wts.py拷贝至yolov5下:

sudo cp -r tenorrtx/yolov5/gen_wts.py {路径}yolov5

(3)生成yolov5s.pt对应的yolov5x.wts:

cd yolov5
python gen_wts.py -w yolov5ws.pt -o yolov5s.wts

(4)修改tensorrtx中yolov5下的CMakeLists.txt:

include_directories("/home/tensorrt_tar/TensorRT-6.0.1.5/include")    #具体路径根据实际安装路径来写
link_directories(/home/tensorrt_tar/TensorRT-6.0.1.5/lib)

5)以原始yolov5s模型为例运行:
C++接口:

cd tensorrtx/yolov5/
mkdir build && cd build
sudo cp -r yolov5/yolov5s.wts build
cmake ..
make

//由wts生成engine文件
./yolov5 -s yolov5s.wts yolov5s.engine s
//利用engine文件进行推理,此处输入图片所在文件夹,即可得到预测输出
./yolov5 -d yolov5s.engine ../samples

Python接口:

python yolov5_trt.py

说明:
得到engine文件后即可进行推理,但默认输入是图片文件夹,而一般来说实时性项目多用摄像头做输入源,故在yolov5.cpp中魔改一下再重新编译就行。此外IRuntime, ICudaEngine, IExecutionContext等核心类源码中每次有图片输入都会创建一次,严重拖慢速度,改成仅初始化阶段创建,后续再有图片输入时保持可以提高推理速度。

  • 2
    点赞
  • 24
    收藏
    觉得还不错? 一键收藏
  • 1
    评论
要在Ubuntu中使用TensorRT加速Yolov5需要安装Yolov5TensorRT。 1. 安装Yolov5 可以按照官方文档的步骤进行安装:https://github.com/ultralytics/yolov5 2. 安装TensorRT 可以按照NVIDIA的官方文档进行安装:https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html 3. 导出ONNX模型 将训练好的Yolov5模型导出为ONNX格式,可以使用官方提供的导出脚本进行导出: ``` python export.py --weights path/to/weights.pt --img 640 --batch 1 --name yolov5s --simplify --dynamic ``` 4. 转换为TensorRT模型 使用TensorRT提供的工具`trtexec`将ONNX模型转换为TensorRT模型: ``` trtexec --onnx=path/to/yolov5s.onnx --saveEngine=path/to/yolov5s.engine --explicitBatch --fp16 ``` 这里使用了`--explicitBatch`选项来指定显式批次大小,并使用了`--fp16`选项来启用FP16精度加速。 5. 加载TensorRT模型并推理 在Python代码中加载TensorRT模型,并进行推理: ```python import tensorrt as trt import pycuda.driver as cuda import pycuda.autoinit import numpy as np import cv2 # 加载TensorRT引擎 with open('path/to/yolov5s.engine', 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime: engine = runtime.deserialize_cuda_engine(f.read()) # 创建执行上下文 context = engine.create_execution_context() # 分配输入和输出内存 input_shape = engine.get_binding_shape(0) output_shape = engine.get_binding_shape(1) input_host_mem = cuda.pagelocked_empty(trt.volume(input_shape), dtype=np.float32) output_host_mem = cuda.pagelocked_empty(trt.volume(output_shape), dtype=np.float32) input_device_mem = cuda.mem_alloc(input_host_mem.nbytes) output_device_mem = cuda.mem_alloc(output_host_mem.nbytes) # 加载图像并预处理 img = cv2.imread('path/to/image.jpg') img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = cv2.resize(img, (input_shape[3], input_shape[2])) img = img.astype(np.float32) / 255.0 img = np.transpose(img, [2, 0, 1]) # HWC -> CHW img = np.expand_dims(img, axis=0) # add batch dimension # 将输入数据复制到设备内存 cuda.memcpy_htod(input_device_mem, input_host_mem) # 执行推理 context.execute_v2(bindings=[int(input_device_mem), int(output_device_mem)]) # 将输出数据复制到主机内存 cuda.memcpy_dtoh(output_host_mem, output_device_mem) # 解析输出结果 output = output_host_mem.reshape(output_shape) ``` 这里需要注意的是,输入和输出内存的分配需要使用`cuda.pagelocked_empty`函数分配固定大小的内存,并且需要使用`cuda.mem_alloc`函数在设备上分配对应大小的内存,这样能够提高数据传输效率。 参考链接: - https://developer.nvidia.com/blog/speeding-up-deep-learning-inference-using-tensorrt-integration-with-onnx-runtime/ - https://docs.nvidia.com/deeplearning/tensorrt/sample-support-guide/index.html#python_sample_code
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值