模型部署 - YoloV5模型与TensorRT结合使用的完整流程 - 从模型训练到模型部署 - 人脸检测任务
将YOLOv5与TensorRT结合使用可以显著提高推理速度， 特别是在NVIDIA GPU上。 下面是一个基本的步骤指南， 帮助你了解YOLOv5与TensorRT结合使用的步骤。
人脸检测是一种常见的计算机视觉任务， 下面将详细介绍从模型训练到模型部署的完整流程， 包括使用YOLOv5进行人脸检测的示例， 最后给出C++示例代码。
主要流程包括：
1. 数据准备
2. 模型训练
3. 模型导出为ONNX格式
4. 模型转换为TensorRT引擎
5. C++部署示例
1. 数据准备
准备一个包含人脸图像和相应标注的数据集，常用的数据集包括WIDER FACE等。标注格式可以是YOLO格式或者COCO格式。
2. 模型训练
使用YOLOv5进行模型训练。 假设已经安装了YOLOv5的依赖。
# 克隆 YOLOv5 仓库 git clone https://github.com/ultralytics/yolov5 cd yolov5 # 安装必要的依赖 pip install -r requirements.txt
创建一个自定义的配置文件 data/faces.yaml，内容如下：
train: /path/to/train/images val: /path/to/val/images nc: 1 names: ['face']
然后开始训练：
python train.py --img 640 --batch 16 --epochs 100 --data data/face.yaml --weights yolov5s.pt --cache
3. 模型导出为ONNX格式
训练完成后， 将模型导出为ONNX格式：
python export.py --weights runs/train/exp/weights/best.pt --img 640 --batch 1 --device 0 --include onnx
4. 模型转换为TrensorRT引擎
使用TensorRT工具将ONNX模型转换为TensorRT引擎：
/usr/src/tensorrt/bin/trtexec --onnx=best.onnx --saveEngine=best.trt --fp16
5. C++部署示例
以下是一个完整的C++示例代码， 展示如何加载TensorRT引擎进行推理。
必备条件：
确保你已经安装了以下工具和库：
- TensorRT
- CUDA
- OpenCV
C++ 示例代码
#include <iostream> #include <fstream> #include <vector> #include <NvInfer.h> #include <NvInferRuntime.h> #include <cuda_runtime_api.h> #include <opencv2/opencv.hpp> using namespace nvinfer1; class Logger: public ILogger { public: void log(Serverity severity, const char* msg) noexcept override { if(severity != Severity::kINFO) std::cout << msg << std::endl; } }; void * safeCudaMalloc(size_t memsize) { void deviceMem; cudaError_t stats = cudaMalloc(&deviceMem, memSize); if(status != cudaSuccess) { std::cerr << "cudaMalloc Failed!" << std::endl; return nullptr; } return deviceMem; } bool readEngineFile(const std::string& engineFile, IRuntime*& runtime, ICudaEngine*& engine) { std::ifstream file(engineFile, std::ios::binary); if (!file) { std::cerr << "Failed to open engine file: " << engineFile << std::endl; return false; } file.seekg(0, file.end); size_t size = file.tellg(); file.seekg(0, file.beg); std::vector<char> buffer(size); file.read(buffer.data(), size); file.close(); Logger logger; runtime = createInferRuntime(logger); if (!runtime) { std::cerr << "Failed to create TensorRT runtime" << std::endl; return false; } engine = runtime->deserializeCudaEngine(buffer.data(), size); if (!engine) { std::cerr << "Failed to deserialize CUDA engine" << std::endl; return false; } return true; } void doInference(IExecutionContext& context, float* input, float* output, int batchSize, int inputSize, int outputSize) { const ICudaEngine& engine = context.getEngine(); assert(engine.getNbBindings() == 2); // Allocate GPU buffers void* buffers[2]; const int inputIndex = engine.getBindingIndex("images"); const int outputIndex = engine.getBindingIndex("output"); buffers[inputIndex] = safeCudaMalloc(inputSize * batchSize * sizeof(float)); buffers[outputIndex] = safeCudaMalloc(outputSize * batchSize * sizeof(float)); // Create stream cudaStream_t stream; cudaStreamCreate(&stream); // Copy input data to GPU cudaMemcpyAsync(buffers[inputIndex], input, inputSize * batchSize * sizeof(float), cudaMemcpyHostToDevice, stream); // Execute inference context.enqueue(batchSize, buffers, stream, nullptr); // Copy output data to CPU cudaMemcpyAsync(output, buffers[outputIndex], outputSize * batchSize * sizeof(float), cudaMemcpyDeviceToHost, stream); cudaStreamSynchronize(stream); // Release stream and buffers cudaStreamDestroy(stream); cudaFree(buffers[inputIndex]); cudaFree(buffers[outputIndex]); } int main() { const std::string engineFile = "best.trt"; IRuntime* runtime = nullptr; ICudaEngine* engine = nullptr; if (!readEngineFile(engineFile, runtime, engine)) { return -1; } IExecutionContext* context = engine->createExecutionContext(); if (!context) { std::cerr << "Failed to create execution context" << std::endl; return -1; } // Load an image using OpenCV cv::Mat img = cv::imread("face.jpg"); cv::resize(img, img, cv::Size(640, 640)); img.convertTo(img, CV_32FC3, 1.0 / 255); // Prepare input data std::vector<float> input(3 * 640 * 640); std::memcpy(input.data(), img.data, input.size() * sizeof(float)); // Prepare output data const int outputSize = 25200 * 85; // Adjust based on your model's output std::vector<float> output(outputSize); // Run inference doInference(*context, input.data(), output.data(), 1, 3 * 640 * 640, outputSize); // Process output data (this is just an example, you'd need to parse and use the output appropriately) for (int i = 0; i < outputSize; i += 85) { float confidence = output[i + 4]; if (confidence > 0.5) { std::cout << "Detection: " << i / 85 << " Confidence: " << confidence << std::endl; } } // Clean up context->destroy(); engine->destroy(); runtime->destroy(); return 0; }
Logger 类：自定义的日志记录器，用于捕捉 TensorRT 的日志。
safeCudaMalloc：一个简单的 CUDA 内存分配函数。
readEngineFile：从文件中读取 TensorRT 引擎并反序列化。
doInference：执行推理函数，包括内存分配、数据传输和推理执行。
main 函数：加载 TensorRT 引擎，准备输入数据，执行推理，并处理输出结果。
编译和运行
确保你已经安装并配置好了 CUDA 和 TensorRT，然后使用以下命令编译和运行代码：
g++ -std=c++11 -o yolov5_trt yolov5_trt.cpp -lnvinfer -lnvinfer_plugin -lcudart -lopencv_core -lopencv_imgcodecs -lopencv_imgproc ./yolov5_trt
通过以上步骤，你可以完成从模型训练到部署的整个流程，并使用 C++ 代码进行人脸检测的推理。你可以根据需要进一步优化和扩展代码。
