文章目录
- 1、环境配置
- 2、推理部分
- 2.1、检测
- 2.2、分类
- 2.3、分割
- 2.4、INT8 量化
1、环境配置
链接: TensorRT cuda环境安装
2、推理部分
下载yolov5对应版本的包
https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5
2.1、检测
1、源码模型下载
git clone -b v7.0 https://github.com/ultralytics/yolov5.git
git clone -b yolov5-v7.0 https://github.com/wang-xinyu/tensorrtx.git
cd yolov5/
wget https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt2、pt模型转换
# 将tensorrx文件夹下yolov5下的gen_wts.py拷贝到yolov5目录下
cp [PATH-TO-TENSORRTX]/yolov5/gen_wts.py .
python gen_wts.py -w yolov5s.pt -o yolov5s.wts3、构建tensorrtx/yolov5并运行
cd tensorrx目录下
mkdir build
cd build
cp [PATH-TO-ultralytics-yolov5]/yolov5s.wts .
cmake ..
make# 转换成engin
./yolov5_det -s yolov5s.wts yolov5s.engine s# 执行推理
./yolov5_det -d yolov5s.engine ../images# 自定义 model with depth_multiple=0.17, width_multiple=0.25 in yolov5.yaml
./yolov5_det -s yolov5_custom.wts yolov5.engine c 0.17 0.25
./yolov5_det -d yolov5.engine ../images4、在python中运行
# 安装python-tensorrt, pycuda等。
# 确保yolov5s.engin和libmyplugins已经构建好
python yolov5_det_trt.py# 另一个版本的python脚本,它使用CUDA python而不是pycuda
python yolov5_det_cuda_python.py
不拷贝到yolov5目录下直接在tensorrx目录下执行就会报错
推理结果
2.2、分类
# Download ImageNet labels
wget https://github.com/joannzhang00/ImageNet-dataset-classes-labels/blob/main/imagenet_classes.txt# Build and serialize TensorRT engine
./yolov5_cls -s yolov5s-cls.wts yolov5s-cls.engine s# Run inference
./yolov5_cls -d yolov5s-cls.engine ../images
2.3、分割
# Build and serialize TensorRT engine
./yolov5_seg -s yolov5s-seg.wts yolov5s-seg.engine s# Download the labels file
wget -O coco.txt https://raw.githubusercontent.com/amikelive/coco-labels/master/coco-labels-2014_2017.txt# Run inference with labels file
./yolov5_seg -d yolov5s-seg.engine ../images coco.txt
2.4、INT8 量化
1、准备校准图像,您可以从训练集中随机选择 1000 张图像。coco_calib对于coco,你也可以从百度盘下载我的校准图像pwd:链接: [a9wh](https://pan.baidu.com/share/init?surl=GOm_-JobpyLMAqZWCDUhKg)2、将其解压到 yolov5/build 中3、在 src/config.h 中设置宏USE_INT8并 make4、序列化模型并进行测试
说明:推理流程来源官方文档:https://github.com/wang-xinyu/tensorrtx/tree/master/yolov5