这是看了韩博《CUDA与TensorRT部署实战课程》第六章的课程部分输出的学习笔记,在这写一个未封装的TensorRT推理分类图形
1. 简单封装一个Model做了什么
其实这里就是做了一个model Load的一个过程,然后build model再然后infer很多张图片, 这些看似封装了但其实封装的很差
2. 看一下Model这个类做了什么
2.1 工具
首先先写一个Logger, 这里是从nvinfer1::ILogger继承过来的, 因为里面有一个必须实现的虚函数
然后在下面自己把这些封装一下
下面这段代码用于确保为TensorRT对象分配的内存在拥有它们的唯一指针超出范围时被正确地回收。
首先智能指针std::unique_ptr 和 std::shared_ptr 是 C++ 标准库中的智能指针,用于管理动态分配的资源(如堆上的对象)。它们提供了自动资源管理,以帮助避免内存泄漏和释放后的访问错误。
unique_ptr管理的资源只能被一个unique_Ptr拥有,当这个指针管理的作用域或被销毁时,它会自动释放所管理的资源。
这边可以理解为是基于智能指针的第二层的保险
举个例子, make_uniquenvinfer1::IBuilder创建了一个智能指针builder,它管理一个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引擎。
2.3 infer 模型推理
简单流程, 比较方便大家理解,如果推理一个TensorRT的文件
这里为了方便理解,我把这个地方分成了15个步骤
- 加载引擎
2.创建一个流来做后面的异步处理
3.从engine拿到模型的输入和输出, 也要计算出来input, output的size, 为了后面的cudaMemcpy做的
4. 给预处理分配足够的Device空间和Host空间, 这里主要是用了之前的CUDA_CHECK
5. preprocess打印信息
会有下面的输出
6. 做图像的预处理工作 归一化 + BGR2RGB + hwc2chw
先看一下hwc跟chw的区别
- HWC格式:
- HEIGHT: 图像的高度,有多少行
- WIDTH: 图像的宽度,每行有多少列
- CHANNEL: RGB三个通道,每个像素点有RGB三个值
可以想象为一叠书,每本书是一张图像,书页的高和宽就是H和W,一页书由墨水的RGB三原色混合而成,是三个channel。
2. CHW格式:
- CHANNEL: 把图像看成三个通道的集合,R通道一堆,G通道一堆,B通道一堆
- HEIGHT: 每个通道堆里有多少行,行数等于图像高度
- WIDTH: 每行有多少列,列数等于图像宽度
可以想象为把书分开三摞,每摞书是一种颜色的页面,每本书的页数(高度)和每页的文字数量(宽度)不变。
输入时HWC格式比较直观,但CHW格式在计算机中的存储更连续,利于访问。这个转换 rearrange了元素的位置,但总数据量未变。
总结是h的合集,w的合集,根据不同通道
实现hwc2chw的地方是通过下面循环,
BGR2RGB就是在写的时候把通道换一下就好了,归一化在这个里面做的
总结: 这里其实就是在input_host上面把数据做了
6. TensorRT执行推理
把数据从host挪到Device上面去,然后执行enqueueV2, 这里就是TRT的黑盒子过程了,这里就会把build engine的时候所计算出来最好的方法拿出来用
7. 从后处理拿到结果并且输出到控制台
想要后处理,还得把数据从Device放到Host上面去, 这里不要忘记了还要同步一下
然后取出来labels最大的那个类别
然后就可以看到对应的结果了
3. 总结
上面完成的是一个基于TensorRT实现图像分类的推理框架。总体来说,该框架存在以下问题:
- 设计模式缺失,导致代码复用性、扩展性、可读性较差。建议使用面向对象设计模式,提高模块化和封装程度。
- 封装不够,外部函数需要处理内部逻辑,降低了抽象程度。建议降低接口间的耦合,隐藏内部实现细节。
- 内存复用方面有优化空间,存在重复分配和释放开销。建议预分配资源,重用上下文对象,减少内存操作。
- 存在多方面功能不完善之处,如INT8量化、TensorRT插件、并行推理等。后续需要持续扩展,使框架功能更加完备。
- 当前仅支持单张图像推理,无法扩展到批处理。建议抽象批处理逻辑,实现动态大小的批处理。
- 仅使用GPU推理,可进一步利用CPU并行来优化。
综上,该框架尚需在模块化、性能、功能等多方面进一步优化。但作为初步实践和总结,已经反映出框架设计和实现的多方面考量,对后续改进具有很好的启发作用。