/PPYOLOE_Tensorrt

A simple implementation of Tensorrt PPYOLOE

Primary LanguagePython

B站视频教程

https://www.bilibili.com/video/BV12a411H73Q/

Introduction

PPYOLOE是百度提出的单阶段Anchor-free模型,超越了多种YOLO模型,本仓库实现PPYOLOE端到端的Tensorrt部署

Environment

  • Tensorrt 8.4.1.5
  • Cuda 10.2 Cudnn 8.4.1
  • onnx 1.8.0
  • onnx-simplifier 0.3.10
  • onnx-graphsurgeon 0.3.19
  • Torch 1.10.0

Benchmark

Model Size mAPval 0.5:0.95 GTX1650 FP16(ms) ONNX (EfficientNMS)
PPYOLOE-S 640 43.0 10.4 model
PPYOLOE-M 640 49.0 18.4 model
PPYOLOE-L 640 51.4 28.3 model
PPYOLOE-X 640 52.3 48.9 model

说明:此处FP16预测时间包含preprocess+inference+nms,测速方法为warmup10次,预测100次取平均值,并未使用trtexec测速,与官方测速不同;mAPval为原始Paddle模型精度,转换后精度未测试。

Quick Start

下载PaddleDetection仓库,安装相关依赖,并导出训练模型。因为初始NMS中部分算子不支持Tensorrt,需要转换,所以导出时设置exclude_nms=True去除NMS

git clone https://github.com/PaddlePaddle/PaddleDetection.git
cd PaddleDetection
python3 tools/export_model.py \
  -c configs/ppyoloe/ppyoloe_crn_s_300e_coco.yml \
  -o weights=https://paddledet.bj.bcebos.com/models/ppyoloe_crn_s_300e_coco.pdparams \
  trt=True \
  exclude_nms=True \
  TestReader.inputs_def.image_shape=[3,640,640] \
  --output_dir inference_model

安装paddle2onnx,并通过下方命令导出onnx模型

paddle2onnx --model_dir inference_model/ppyoloe_crn_s_300e_coco \
            --model_filename model.pdmodel \
            --params_filename model.pdiparams \
            --save_file examples/ppyoloe_crn_s_300e_coco.onnx \
            --input_shape_dict "{'image':[1, 3, 640, 640], 'scale_factor': [1, 2]}"

使用本仓库中EfficientNMS_TRT.py,为onnx模型添加EfficientNMS算子,注意修改文件头部常量

INPUT_PATH = './onnx_model/ppyoloe_crn_s_300e_coco.onnx'
WEIGHTS_TYPE = "s"
SAVE_PATH = "./onnx_model/ppyoloes_nms.onnx"
CLASS_NUM = 80
SCORE_THRESHOLD = 0.5
IOU_THRESHOLD = 0.4
python EfficientNMS_TRT.py

使用Tensorrt官方工具trtexec转化添加完EfficientNMS的onnx模型。FP32预测删除--fp16参数即可

trtexec --onnx=./ppyoloes_nms.onnx --saveEngine=./ppyoloes_nms_fp16.engine --fp16 --workspace=30

等待生成序列化模型后,修改本仓库infer.py模型路径和图片路径

trt_engine = TRT_engine("./trt_model/ppyoloes_nms_fp16.engine")
img1 = cv2.imread("./pictures/zidane.jpg")
python infer.py

Reference

https://github.com/PaddlePaddle/PaddleDetection

https://github.com/ultralytics/yolov5

https://github.com/Linaom1214/tensorrt-python

https://github.com/triple-Mu