TensorRT进阶（一）：从零搭建一个TensorRT分类框架
作者： 人工智能技术与时代人物风云 来源： 人工智能技术与时代人物风云
点击下方卡片 ，关注“自动驾驶Daily ”公众号
ADAS巨卷干货，即可获取
**»****点击进入→自动驾驶Daily技术交流群**
写在前面
本文是我在学习韩博《CUDA与TensorRT部署实战课程》第六章的课程部分输出的个人学习笔记，欢迎大家一起讨论学习！这里写一个未封装的TensorRT推理分类图形，不封装的原因是为了学习的时候能够更好的了解到这里面的过程，看完后会有一个汇总，如果你正在学习相关内容，强烈建议收藏！
1. 先看简单封装一个Model做了什么
其实这里就是做了一个model Load的一个过程，然后build model再然后infer很多张图片, 这些看似封装了但其实封装的很差
#include <iostream>
#include <memory>
#include "model.hpp"
#include "utils.hpp"
using namespace std;
int main(int argc, char const *argv[])
{
Model model("models/onnx/resnet50.onnx", Model::precision::FP32);
if(!model.build()){
LOGE("fail in building model");
return 0;
}
if(!model.infer("data/fox.png")){
LOGE("fail in infering model");
return 0;
}
if(!model.infer("data/cat.png")){
LOGE("fail in infering model");
return 0;
}
if(!model.infer("data/eagle.png")){
LOGE("fail in infering model");
return 0;
}
if(!model.infer("data/gazelle.png")){
LOGE("fail in infering model");
return 0;
}
return 0;
}
2. 看一下Model这个类做了什么
2.1 工具
首先先写一个Logger, 这里是从nvinfer1::ILogger继承过来的, 因为里面有一个必须实现的虚函数
然后在下面自己把这些封装一下
virtual void log(Severity severity, AsciiChar const* msg) noexcept = 0;
class Logger : public nvinfer1::ILogger{
public:
virtual void log (Severity severity, const char* msg) noexcept override{
string str;
switch (severity){
case Severity::kINTERNAL_ERROR: str = RED "[fatal]" CLEAR;
case Severity::kERROR: str = RED "[error]" CLEAR;
case Severity::kWARNING: str = BLUE "[warn]" CLEAR;
case Severity::kINFO: str = YELLOW "[info]" CLEAR;
case Severity::kVERBOSE: str = PURPLE "[verb]" CLEAR;
}
// if (severity <= Severity::kINFO)
// cout << str << string(msg) << endl;
}
};
下面这段代码用于确保为TensorRT对象分配的内存在拥有它们的唯一指针超出范围时被正确地回收。
首先智能指针std::unique_ptr 和 std::shared_ptr 是 C++ 标准库中的智能指针，用于管理动态分配的资源（如堆上的对象）。它们提供了自动资源管理，以帮助避免内存泄漏和释放后的访问错误。
unique_ptr管理的资源只能被一个unique_Ptr拥有，当这个指针管理的作用域或被销毁时，它会自动释放所管理的资源。
这边可以理解为是基于智能指针的第二层的保险
struct InferDeleter
{
template <typename T>
void operator()(T* obj) const
{
delete obj;
}
};
template <typename T>
using make_unique = std::unique_ptr<T, InferDeleter>;
举个例子, make_uniquenvinfer1::IBuilder创建了一个智能指针builder，它管理一个nvinfer1::IBuilder对象，该对象是通过nvinfer1::createInferBuilder(logger)函数创建的。
auto builder = make_unique<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
2.2 构建模型, 从onnx到engine
使用TensorRT的API根据ONNX模型构建和序列化引擎的过程,主要步骤包括:
-
判断引擎文件是否已经存在,如果存在直接加载,否则开始构建。
-
创建TensorRT的构建器(IBuilder)、网络定义(INetworkDefinition)、构建配置(IBuilderConfig)。
-
使用ONNX解析器(IParser)解析ONNX模型,生成TensorRT网络定义。
-
设置最大工作空间大小,输出构建日志详情。
-
根据精度模式设置构建配置的优化方式。
-
使用构建器和配置构建引擎(ICudaEngine)。
-
序列化引擎为plan文件,并保存到磁盘。
-
反序列化plan文件获取ICudaEngine。
-
记录输入和输出的维度信息。
-
打印网络结构的优化前后信息。
-
返回构建成功的标志。
所以主要步骤是:
-
使用ONNX解析器解析模型得到TensorRT网络定义
-
使用TensorRT构建器和配置构建优化后的引擎
-
序列化并保存引擎计划到磁盘
-
反序列化加载引擎
通过这些步骤完成了使用TensorRT对ONNX模型进行解析、优化、构建和序列化的工作,得到了一个高性能的TensorRT引擎。
bool Model::build_engine() {
// 我们也希望在build一个engine的时候就把一系列初始化全部做完，其中包括
// 1. build一个engine
// 2. 创建一个context
// 3. 创建推理所用的stream
// 4. 创建推理所需要的device空间
// 这样，我们就可以在build结束以后，就可以直接推理了。这样的写法会比较干净
auto builder = shared_ptr<IBuilder>(createInferBuilder(*m_logger), destroy_trt_ptr<IBuilder>);
auto network = shared_ptr<INetworkDefinition>(builder->createNetworkV2(1), destroy_trt_ptr<INetworkDefinition>);
auto config = shared_ptr<IBuilderConfig>(builder->createBuilderConfig(), destroy_trt_ptr<IBuilderConfig>);
auto parser = shared_ptr<IParser>(createParser(*network, *m_logger), destroy_trt_ptr<IParser>);
config->setMaxWorkspaceSize(m_workspaceSize);
config->setProfilingVerbosity(ProfilingVerbosity::kLAYER_NAMES_ONLY); //这里也可以设置为kDETAIL;
if (!parser->parseFromFile(m_onnxPath.c_str(), 1)){
return false;
}
if (builder->platformHasFastFp16() && m_params->prec == model::FP16) {
config->setFlag(BuilderFlag::kFP16);
config->setFlag(BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
}
auto engine = shared_ptr<ICudaEngine>(builder->buildEngineWithConfig(*network, *config), destroy_trt_ptr<ICudaEngine>);
auto plan = builder->buildSerializedNetwork(*network, *config);
auto runtime = shared_ptr<IRuntime>(createInferRuntime(*m_logger), destroy_trt_ptr<IRuntime>);
// 保存序列化后的engine
save_plan(*plan);
// 根据runtime初始化engine, context, 以及memory
setup(plan->data(), plan->size());
// 把优化前和优化后的各个层的信息打印出来
LOGV("Before TensorRT optimization");
print_network(*network, false);
LOGV("After TensorRT optimization");
print_network(*network, true);
return true;
}
2.3 infer 模型推理
简单流程, 比较方便大家理解，如果推理一个TensorRT的文件
runtime->engine
engine->context
context->intput_dims, (224X224, 640x640)
context->output_dims
context->enquenev2
这里为了方便理解，我把这个地方分成了15个步骤
- 加载引擎
// 1. 加载引擎
vector<unsigned char> modelData = loadFile(mEnginePath);
auto runtime = make_unique<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));
auto engine = make_unique<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(modelData.data(), modelData.size()));
- 创建一个流来做后面的异步处理
cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
- 从engine拿到模型的输入和输出, 也要计算出来input, output的size, 为了后面的cudaMemcpy做的
int input_width = input_dims.d[3];
int input_height = input_dims.d[2];
int input_channel = input_dims.d[1];
int num_classes = output_dims.d[1];
int input_size = input_channel * input_width * input_height * sizeof(float);
int output_size = num_classes * sizeof(float);
- 给预处理分配足够的Device空间和Host空间, 这里主要是用了之前的CUDA_CHECK
/*Preprocess -- 分配host和device的内存空间*/
float* input_host = nullptr;
float* input_device = nullptr;
float* output_host = nullptr;
float* output_device = nullptr;
CUDA_CHECK(cudaMalloc(&input_device, input_size));
CUDA_CHECK(cudaMalloc(&output_device, output_size));
CUDA_CHECK(cudaMallocHost(&input_host, input_size));
CUDA_CHECK(cudaMallocHost(&output_host, output_size));
- preprocess打印信息
/*Preprocess -- 读取数据*/
cv::Mat input_image;
input_image = cv::imread(imagePath);
if (input_image.data == nullptr) {
LOGE("file not founded! Program terminated");
return false;
} else {
LOG("Model: %s", getFileName(mOnnxPath).c_str());
LOG("Precision: %s", getPrecision(mPrecision).c_str());
LOG("Image: %s", getFileName(imagePath).c_str());
}
会有下面的输出
[info]models/engine/resnet50_fp32.engine has been generated!
[info]Model: resnet50.onnx
[info]Precision: FP32
[info]Image: fox.png
- 做图像的预处理工作 归一化 + BGR2RGB + hwc2chw
先看一下hwc跟chw的区别
- HWC格式:
-
HEIGHT: 图像的高度,有多少行
-
WIDTH: 图像的宽度,每行有多少列
-
CHANNEL: RGB三个通道,每个像素点有RGB三个值
可以想象为一叠书,每本书是一张图像,书页的高和宽就是H和W,一页书由墨水的RGB三原色混合而成,是三个channel。
- CHW格式:
-
CHANNEL: 把图像看成三个通道的集合,R通道一堆,G通道一堆,B通道一堆
-
HEIGHT: 每个通道堆里有多少行,行数等于图像高度
-
WIDTH: 每行有多少列,列数等于图像宽度
可以想象为把书分开三摞,每摞书是一种颜色的页面,每本书的页数(高度)和每页的文字数量(宽度)不变。
输入时HWC格式比较直观,但CHW格式在计算机中的存储更连续,利于访问。这个转换 rearrange了元素的位置,但总数据量未变。
总结是h的合集，w的合集，根据不同通道
实现hwc2chw的地方是通过下面循环，
for (int i = 0; i < input_height; i++) {
for (int j = 0; j < input_width; j++) {
index = i * input_width * input_channel + j * input_channel;
BGR2RGB就是在写的时候把通道换一下就好了，归一化在这个里面做的
/*Preprocess -- host端进行normalization + BGR2RGB + hwc2cwh)*/
int index;
int offset_ch0 = input_width * input_height * 0;
int offset_ch1 = input_width * input_height * 1;
int offset_ch2 = input_width * input_height * 2;
for (int i = 0; i < input_height; i++) {
for (int j = 0; j < input_width; j++) {
index = i * input_width * input_channel + j * input_channel;
input_host[offset_ch2++] = (input_image.data[index + 0] / 255.0f - mean[0]) / std[0];
input_host[offset_ch1++] = (input_image.data[index + 1] / 255.0f - mean[1]) / std[1];
input_host[offset_ch0++] = (input_image.data[index + 2] / 255.0f - mean[2]) / std[2];
}
}
总结: 这里其实就是在input_host上面把数据做了
- TensorRT执行推理
把数据从host挪到Device上面去，然后执行enqueueV2, 这里就是TRT的黑盒子过程了，这里就会把build engine的时候所计算出来最好的方法拿出来用
CUDA_CHECK(cudaMemcpyAsync(input_device, input_host, input_size, cudaMemcpyKind::cudaMemcpyHostToDevice, stream));
/*Inference -- device端进行推理*/
float* bindings[] = {input_device, output_device};
if (!context->enqueueV2((void**)bindings, stream, nullptr)){
LOG("Error happens during DNN inference part, program terminated");
return false;
}
- 从后处理拿到结果并且输出到控制台
想要后处理，还得把数据从Device放到Host上面去, 这里不要忘记了还要同步一下
CUDA_CHECK(cudaMemcpyAsync(output_host, output_device, output_size, cudaMemcpyKind::cudaMemcpyDeviceToHost, stream));
CUDA_CHECK(cudaStreamSynchronize(stream));
然后取出来labels最大的那个类别
/*Postprocess -- 寻找label*/
ImageNetLabels labels;
int pos = max_element(output_host, output_host + num_classes) - output_host;
float confidence = output_host[pos] * 100;
然后就可以看到对应的结果了
3. 总结
上面完成的是一个基于TensorRT实现图像分类的推理框架。总体来说,该框架存在以下问题:
-
设计模式缺失,导致代码复用性、扩展性、可读性较差。建议使用面向对象设计模式,提高模块化和封装程度。
-
封装不够,外部函数需要处理内部逻辑,降低了抽象程度。建议降低接口间的耦合,隐藏内部实现细节。
-
内存复用方面有优化空间,存在重复分配和释放开销。建议预分配资源,重用上下文对象,减少内存操作。
-
存在多方面功能不完善之处,如INT8量化、TensorRT插件、并行推理等。后续需要持续扩展,使框架功能更加完备。
-
当前仅支持单张图像推理,无法扩展到批处理。建议抽象批处理逻辑,实现动态大小的批处理。
-
仅使用GPU推理,可进一步利用CPU并行来优化。
综上,该框架尚需在模块化、性能、功能等多方面进一步优化。但作为初步实践和总结,已经反映出框架设计和实现的多方面考量,对后续改进具有很好的启发作用。
引用
[1] https://wrzpl.xet.tech/s/T30p4
投稿作者为『**自动驾驶之心知识星球** 』特邀嘉宾，如果您希望分享到自动驾驶之心平台，欢迎联系我们！
① 全网独家视频课程
BEV感知、毫米波雷达视觉融合、多传感器标定、多传感器融合、多模态3D目标检测、点云3D目标检测、目标跟踪、Occupancy、cuda与TensorRT模型部署、协同感知、语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频（扫码学习）
视频官网：www.zdjszx.com
② 国内首个自动驾驶学习社区
近2000人的交流社区，涉及30+自动驾驶技术栈学习路线，想要了解更多自动驾驶感知（2D检测、分割、2D/3D车道线、BEV感知、3D目标检测、Occupancy、多传感器融合、多传感器标定、目标跟踪、光流估计）、自动驾驶定位建图（SLAM、高精地图、局部在线地图）、自动驾驶规划控制/轨迹预测等领域技术方案、AI模型部署落地实战、行业动态、岗位发布，欢迎扫描下方二维码，加入自动驾驶之心知识星球，这是一个真正有干货的地方，与领域大佬交流入门、学习、工作、跳槽上的各类难题，日常分享论文+代码+视频 ，期待交流！
③【自动驾驶之心】技术交流群
自动驾驶之心是首个自动驾驶开发者社区，聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D目标检测、BEV感知、多模态感知、Occupancy、多传感器融合、transformer、大模型、点云处理、端到端自动驾驶、SLAM、光流估计、深度估计、轨迹预测、高精地图、NeRF、规划控制、模型部署落地、自动驾驶仿真测试、产品经理、硬件配置、AI求职交流 等方向。扫码添加汽车人助理微信邀请入群，备注：学校/公司+方向+昵称（快速入群方式）
**④【自动驾驶之心】平台矩阵，**欢迎联系我们！
更多AI工具，参考Github-AiBard123，国内AiBard123
