使用TensorRT对YOLOv8模型进行加速推理
这里使用GitHub上shouxieai的 infer框架 对YOLOv8模型进行加速推理，操作过程如下所示：
1.配置环境，依赖项，包括：
(1).CUDA: 11.8
(2).cuDNN: 8.7.0
(3).TensorRT: 8.5.3.1
(4).ONNX: 1.16.0
(5).OpenCV: 4.10.0
2.clone infer代码：https://github.com/shouxieai/infer
3.使用 https://blog.csdn.net/fengbingchun/article/details/140691177 中采用的数据集生成best.onnx，训练代码如下所示：
import argparse
import colorama
from ultralytics import YOLO
import torch
def parse_args():
parser = argparse.ArgumentParser(description="YOLOv8 train")
parser.add_argument("--yaml", required=True, type=str, help="yaml file")
parser.add_argument("--epochs", required=True, type=int, help="number of training")
parser.add_argument("--task", required=True, type=str, choices=["detect", "segment"], help="specify what kind of task")
args = parser.parse_args()
return args
def train(task, yaml, epochs):
if task == "detect":
model = YOLO("yolov8n.pt") # load a pretrained model
elif task == "segment":
model = YOLO("yolov8n-seg.pt") # load a pretrained model
else:
print(colorama.Fore.RED + "Error: unsupported task:", task)
raise
results = model.train(data=yaml, epochs=epochs, imgsz=640) # train the model
metrics = model.val() # It'll automatically evaluate the data you trained, no arguments needed, dataset and settings remembered
# model.export(format="onnx") #, dynamic=True) # export the model, cannot specify dynamic=True, opencv does not support
model.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=640)
model.export(format="torchscript") # libtorch
model.export(format="engine", imgsz=640, dynamic=False, verbose=False, batch=1, workspace=2) # tensorrt fp32
# model.export(format="engine", imgsz=640, dynamic=True, verbose=True, batch=4, workspace=2, half=True) # tensorrt fp16
# model.export(format="engine", imgsz=640, dynamic=True, verbose=True, batch=4, workspace=2, int8=True, data=yaml) # tensorrt int8
if __name__ == "__main__":
# python test_yolov8_train.py --yaml datasets/melon_new_detect/melon_new_detect.yaml --epochs 1000 --task detect
colorama.init()
args = parse_args()
if torch.cuda.is_available():
print("Runging on GPU")
else:
print("Runting on CPU")
train(args.task, args.yaml, args.epochs)
print(colorama.Fore.GREEN + "====== execution completed ======")4.将best.onnx文件通过infer中的v8trans.py转换为best.transd.onnx，执行如下命令：增加Transpose层，YOLOv5不需要
python v8trans.py best.onnx注：yolov8 onnx的输出为NHW，而inter框架的输出只支持NWH，因此需要在原始onnx的输出之前添加一个Transpose节点
5.从 https://docs.nvidia.com/deeplearning/cudnn/archives/cudnn-870/install-guide/index.html#install-zlib-windows 下载zlib123dllx64.zip，解压缩将其中的zlibwapi.dll拷贝到C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.8\bin目录下
6.通过TensorRT中的trtexec.exe将best.transd.onnx转换为best.transd.engine，分别执行如下命令：
trtexec.exe --onnx=best.transd.onnx --saveEngine=best.transd.fp32.engine
trtexec.exe --onnx=best.transd.onnx --fp16 --saveEngine=best.transd.fp16.engine
trtexec.exe --onnx=best.transd.onnx --int8 --saveEngine=best.transd.int8.engine注：产生的best.transd.fp32.engine和best.transd.fp16.engine大小类似，推理耗时和准确度也类似；best.transd.int8.engine大小约是best.transd.fp32.engine的四分之一，推理耗时也小，但是准确度非常低
7.测试代码TensorRT_infer.cpp如下：工程见：TensorRT_Infer
#include
#include
#include
#include
#include
#include
#include 执行结果如下图所示：
检测结果如下图所示：
trtexec.exe是一个快速使用TensorRT的工具，无需开发自己的应用程序。此工具有三个主要用途：
(1).根据随机或用户提供的输入数据对网络进行基准测试。
(2).从模型生成序列化引擎(engine)。
(3).从构建器生成序列化时序缓存(serialized timing cache)。
trtexec.exe常用flags说明:
1.构建阶段flags
(1).--onnx=
(2).--minShapes=
(3).–-memPoolSize=
(4).--saveEngine=
(5).--fp16, --bf16, --int8, --fp8, --noTF32, and --best：指定network-level精度。
(6).--stronglyTyped：创建strongly typed网络。
(7).--sparsity=[disable|enable|force]：指定是否使用支持结构化稀疏性(structured sparsity)的策略。
(8).--noCompilationCache：禁用构建中的编译缓存(默认是启用编译缓存)。
(9).--verbose：开启详细日志。
(10).--skipInference：构建并保存engine而不运行推理。
(11).--dumpLayerInfo, --exportLayerInfo=
(12).--precisionConstraints=spec：控制精度约束设置。指定的值可为：none、prefer、obey。
(13).--layerPrecisions=spec：控制每层精度约束。仅当precisionConstraints设置为obey或prefer时才有效。规范从左到右读取，后面的会覆盖前面的。"*"可用作layerName，以指定所有未指定层的默认精度。
如：--layerPrecisions=*:fp16,layer_1:fp32 将除layer_1之外的所有层的精度设置为FP16，而layer_1的精度将设置为FP32。
(14).--layerOutputTypes=spec：控制每层输出类型约束。仅当precisionConstraints设置为obey或prefer时才有效。规范从左到右读取，后面的会覆盖前面的。"*"可用作layerName，以指定所有未指定层的默认精度。
(15).--versionCompatible, --vc：为engine构建和推理启用版本兼容模式。
(16).--tempdir=
2.推理阶段flags
(1).--loadEngine=
(2).--shapes=
(3).--loadInputs=
(4).--noDataTransfers：关闭host to device和device to host的数据传输。
(5).--verbose：开启详细日志。
(6).--dumpProfile, --exportProfile=
(7).--dumpLayerInfo, --exportLayerInfo=
GitHub：https://github.com/fengbingchun/NN_Test