模型部署 - TensorRT 简介及人脸推理简单示例
TensorRT是Nvidia开发的一个高性能深度学习推理引擎。 它主要用于优化和部署已经训练好的深度学习模型， 提高模型在NVIDIA GPU硬件上的推理性能。
TensorRT的主要功能和特点包括：
1. 模型优化
TensorRT提供了一系列优化技术， 包括层融合， 张量格式化， kernel自动调优等， 可以大幅度提高模型在NVIDIA GPU上的推理性能。
2. 硬件加速
TensorRT深度集成了NVIDIA GPU硬件的底层特性， 可以充分发挥GPU的并行计算能力， 提供高性能的推理加速。
3. 跨平台部署
TensorRT支持在Windows， Linux和Jeson等多个平台部署优化后的模型， 方便用户在不同环境中使用。
4. 丰富的API
TensorRT提供了C++和Python两种语言的API， 开发者可以灵活的将TensorRT集成到自己的应用程序中。
5. 支持多种模型格式
TensorRT可以导出ONNX， Caffe， Tensorflow 等主流深度学习框架导出的模型， 并对其进行优化。
6. 多硬件支持
处理NVIDIA GPU， TensorRT也支持在NVIDIA Jesson嵌入式平台上部署优化的模型。
中的来说， TensorRT是NVIDIA专门为深度学习推理场景打造的一款高性能引擎， 广泛应用于需要部署深度学习模型的场景， 如自动驾驶， 智能视觉， 自然语言处理等。 掌握TensorRT的使用对于深度学习工程师来说是非常重要的技能。
以下是一个TensorRT结合C++进行人脸识别的示例代码：
这个示例将涵盖以下步骤:
- 使用 TensorFlow 或 PyTorch 训练好人脸检测和人脸特征提取的深度学习模型。
- 将模型转换为 ONNX 格式。
- 使用 TensorRT 优化 ONNX 模型,生成 TensorRT 引擎。
- 在 C++ 中使用 TensorRT API 加载优化后的模型,进行人脸检测和特征提取。
- 将检测到的人脸信息与特征向量输出。
#include <NvInfer.h> #include <opencv2/opencv.hpp> using namespace nvinfer1; // 初始化 TensorRT 推理引擎 std::unique_ptr<ICudaEngine> initEngine(const std::string& onnx_path) { // 创建IBuilder 对象， 用于构建TensorRT引擎。 IBuilder* builder = createInferBuilder(gLogger); // 创建网络 INetworkDefinition* network = builder->createNetwork(); // 创建ONNX解析器 OnnxParser* parser = createOnnxParser(network); // 从指定的ONNX文件中解析拓扑网络 结果存储在network parser->parseFromFile(onnx_path.c_str(), static_cast<int>(ILogger::Severity::kWARNING)); // 设置批处理大小为1 builder->setMaxBatchSize(1); // 设置最大工作空间大小为1MB builder->setMaxWorkspaceSize(1 << 20); // 构建使用Builder对象，构建CUDA引擎 std::unique_ptr<ICudaEngine> engine(builder->buildCudaEngine(network)); // 清理网络定义和构建器，和解析器 network->destroy(); builder->destroy(); parser->destroy(); return engine; } // 使用 TensorRT 进行人脸检测和特征提取 std::vector<cv::Rect> detectFaces(cv::Mat& image, std::unique_ptr<ICudaEngine>& engine) { std::vector<cv::Rect> faces; // 将图像输入到 TensorRT 引擎进行推理 // 获取人脸检测结果 // 提取每个人脸的特征向量 std::vector<cv::Rect> faces; // 1. 将输入图像转换为 TensorRT 引擎所需的格式 cv::Mat resizedImage; cv::resize(image, resizedImage, cv::Size(300, 300)); float* inputData = new float[300 * 300 * 3]; for (int y = 0; y < 300; y++) { for (int x = 0; x < 300; x++) { cv::Vec3b pixel = resizedImage.at<cv::Vec3b>(y, x); inputData[y * 300 * 3 + x * 3 + 0] = (pixel[2] - 127.5) / 127.5; inputData[y * 300 * 3 + x * 3 + 1] = (pixel[1] - 127.5) / 127.5; inputData[y * 300 * 3 + x * 3 + 2] = (pixel[0] - 127.5) / 127.5; } } // 2. 创建 TensorRT 引擎执行上下文 IExecutionContext* context = engine->createExecutionContext(); // 3. 设置输入数据 context->setBindingDimensions(0, Dims4(1, 3, 300, 300)); void* bindings[] = { inputData }; // 4. 执行 TensorRT 推理 context->executeV2(bindings); // 5. 获取输出数据 ,输出张量的大小是 1 x 100 x 7 // 1 表示批处理大小为 1 // 100表示最多检测 100 个物体 // 7表示每个检测框包含 7 个值(x(左上, 归一化), y(左上，归一化), width(归一化), height(归一化), confidence(整体置信度), class_id（类别ID）, class_score(类别置信度)) float* detectionOutput = new float[1 * 100 * 7]; // 从 TensorRT 引擎的输出张量中将数据复制到我们之前分配的 detectionOutput 数组中 // context->getOutputHandle(0) 获取了第一个(也是唯一一个)输出张量的句柄。由于我们只有一个输出张量,索引为 0。 // copyToHost() 方法将输出张量的数据复制到 CPU 内存中的 detectionOutput 数组 context->getOutputHandle(0)->copyToHost(detectionOutput); // 6. 解析人脸检测结果 遍历这个数组,解析每个检测框的信息(位置、大小、置信度、类别等) for (int i = 0; i < 100; i++) { float score = detectionOutput[i * 7 + 4]; if (score > 0.5) { int x = static_cast<int>(detectionOutput[i * 7 + 0] * image.cols); int y = static_cast<int>(detectionOutput[i * 7 + 1] * image.rows); int width = static_cast<int>(detectionOutput[i * 7 + 2] * image.cols - x); int height = static_cast<int>(detectionOutput[i * 7 + 3] * image.rows - y); faces.push_back(cv::Rect(x, y, width, height)); } } // 7. 释放资源 delete[] inputData; delete[] detectionOutput; context->destroy(); return faces; } int main() { std::unique_ptr<ICudaEngine> engine = initEngine("path/to/your/model.onnx"); cv::Mat image = cv::imread("path/to/your/image.jpg"); std::vector<cv::Rect> faces = detectFaces(image, engine); // 将检测到的人脸信息和特征向量输出 return 0; }
