当前位置：首页 > news >正文

使用TensorRT对YOLOv8模型进行加速推理

news 来源：原创 2024/9/19 11:04:00

这里使用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 torchdef 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 argsdef train(task, yaml, epochs):if task == "detect":model = YOLO("yolov8n.pt") # load a pretrained modelelif task == "segment":model = YOLO("yolov8n-seg.pt") # load a pretrained modelelse:print(colorama.Fore.RED + "Error: unsupported task:", task)raiseresults = model.train(data=yaml, epochs=epochs, imgsz=640) # train the modelmetrics = 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 supportmodel.export(format="onnx", opset=12, simplify=True, dynamic=False, imgsz=640)model.export(format="torchscript") # libtorchmodel.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 int8if __name__ == "__main__":# python test_yolov8_train.py --yaml datasets/melon_new_detect/melon_new_detect.yaml --epochs 1000 --task detectcolorama.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 <iostream>
#include <filesystem>
#include <vector>
#include <fstream>
#include <sstream>
#include <random>
#include <map>
#include <memory>
#include <chrono>
#include <string>
#include <algorithm>#include <opencv2/opencv.hpp>
#include "yolo.hpp"namespace {constexpr float confidence_threshold{ 0.45f }; // confidence threshold
constexpr float nms_threshold{ 0.50f }; // nms threshold
constexpr char* engine_file{ "../../../data/best.transd.fp32.engine" };
constexpr char* images_dir{ "../../../data/images/predict" };
constexpr char* result_dir{ "../../../data/result" };
constexpr char* classes_file{ "../../../data/images/labels.txt" };std::vector<std::string> parse_classes_file(const char* name)
{std::vector<std::string> classes;std::ifstream file(name);if (!file.is_open()) {std::cerr << "Error: fail to open classes file: " << name << std::endl;return classes;}std::string line;while (std::getline(file, line)) {auto pos = line.find_first_of(" ");classes.emplace_back(line.substr(0, pos));}file.close();return classes;
}auto get_dir_images(const char* name)
{std::map<std::string, std::string> images; // image name, image path + image namefor (auto const& dir_entry : std::filesystem::directory_iterator(name)) {if (dir_entry.is_regular_file())images[dir_entry.path().filename().string()] = dir_entry.path().string();}return images;
}auto get_random_color(int labels_number)
{std::random_device rd;std::mt19937 gen(rd());std::uniform_int_distribution<int> dis(100, 255);std::vector<cv::Scalar> colors;for (auto i = 0; i < labels_number; ++i) {colors.emplace_back(cv::Scalar(dis(gen), dis(gen), dis(gen)));}return colors;
}} // namespaceint main()
{namespace fs = std::filesystem;if (!fs::exists(result_dir)) {fs::create_directories(result_dir);}auto classes = parse_classes_file(classes_file);if (classes.size() == 0) {std::cerr << "Error: fail to parse classes file: " << classes_file << std::endl;return -1;}std::cout << "classes: ";for (const auto& val : classes) {std::cout << val << " ";}std::cout << std::endl;auto colors = get_random_color(classes.size());auto model = yolo::load(engine_file, yolo::Type::V8, confidence_threshold, nms_threshold);for (auto i = 0; i < 10; ++i) {std::cout << "i: " << i << std::endl;for (const auto& [key, val] : get_dir_images(images_dir)) {cv::Mat frame = cv::imread(val, cv::IMREAD_COLOR);if (frame.empty()) {std::cerr << "Warning: unable to load image: " << val << std::endl;continue;}auto tstart = std::chrono::high_resolution_clock::now();auto objs = model->forward(yolo::Image(frame.data, frame.cols, frame.rows));auto tend = std::chrono::high_resolution_clock::now();std::cout << "elapsed millisenconds: " << std::chrono::duration_cast<std::chrono::milliseconds>(tend - tstart).count() << " ms" << std::endl;for (const auto& obj : objs) {cv::rectangle(frame, cv::Point(obj.left, obj.top), cv::Point(obj.right, obj.bottom), colors[obj.class_label], 2);std::string class_string = classes[obj.class_label] + ' ' + std::to_string(obj.confidence).substr(0, 4);cv::Size text_size = cv::getTextSize(class_string, cv::FONT_HERSHEY_DUPLEX, 1, 2, 0);cv::Rect text_box(obj.left, obj.top - 40, text_size.width + 10, text_size.height + 20);cv::rectangle(frame, text_box, colors[obj.class_label], cv::FILLED);cv::putText(frame, class_string, cv::Point(obj.left + 5, obj.top - 10), cv::FONT_HERSHEY_DUPLEX, 1, cv::Scalar(0, 0, 0), 2, 0);}std::string path(result_dir);path += "/" + key;cv::imwrite(path, frame);}}std::cout << "test finish" << std::endl;return 0;
}

执行结果如下图所示：