TensorRT进阶 | 从零搭建一个TensorRT分类框架

点击下方卡片,关注“自动驾驶之心”公众号

ADAS巨卷干货,即可获取

>>点击进入→自动驾驶之心【模型部署】技术交流群

论文作者 | EasonBob

编辑 | 自动驾驶之心

写在前面

本文是我在学习韩博《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继承过来的, 因为里面有一个必须实现的虚函数

本文内容均出自《CUDA与TensorRT部署实战课程》

fa57b6801f4bd769a162f27b4228687b.png

然后在下面自己把这些封装一下

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模型构建和序列化引擎的过程,主要步骤包括:

  1. 判断引擎文件是否已经存在,如果存在直接加载,否则开始构建。

  2. 创建TensorRT的构建器(IBuilder)、网络定义(INetworkDefinition)、构建配置(IBuilderConfig)。

  3. 使用ONNX解析器(IParser)解析ONNX模型,生成TensorRT网络定义。

  4. 设置最大工作空间大小,输出构建日志详情。

  5. 根据精度模式设置构建配置的优化方式。

  6. 使用构建器和配置构建引擎(ICudaEngine)。

  7. 序列化引擎为plan文件,并保存到磁盘。

  8. 反序列化plan文件获取ICudaEngine。

  9. 记录输入和输出的维度信息。

  10. 打印网络结构的优化前后信息。

  11. 返回构建成功的标志。

所以主要步骤是:

  1. 使用ONNX解析器解析模型得到TensorRT网络定义

  2. 使用TensorRT构建器和配置构建优化后的引擎

  3. 序列化并保存引擎计划到磁盘

  4. 反序列化加载引擎

通过这些步骤完成了使用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. 加载引擎

// 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()));
  1. 创建一个流来做后面的异步处理

cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
  1. 从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);
  1. 给预处理分配足够的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));
  1. 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
  1. 做图像的预处理工作 归一化 + BGR2RGB + hwc2chw

先看一下hwc跟chw的区别

  1. HWC格式:

  • HEIGHT: 图像的高度,有多少行

  • WIDTH: 图像的宽度,每行有多少列

  • CHANNEL: RGB三个通道,每个像素点有RGB三个值

可以想象为一叠书,每本书是一张图像,书页的高和宽就是H和W,一页书由墨水的RGB三原色混合而成,是三个channel。

  1. 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上面把数据做了

  1. 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;
}
  1. 从后处理拿到结果并且输出到控制台

想要后处理,还得把数据从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实现图像分类的推理框架。总体来说,该框架存在以下问题:

  1. 设计模式缺失,导致代码复用性、扩展性、可读性较差。建议使用面向对象设计模式,提高模块化和封装程度。

  2. 封装不够,外部函数需要处理内部逻辑,降低了抽象程度。建议降低接口间的耦合,隐藏内部实现细节。

  3. 内存复用方面有优化空间,存在重复分配和释放开销。建议预分配资源,重用上下文对象,减少内存操作。

  4. 存在多方面功能不完善之处,如INT8量化、TensorRT插件、并行推理等。后续需要持续扩展,使框架功能更加完备。

  5. 当前仅支持单张图像推理,无法扩展到批处理。建议抽象批处理逻辑,实现动态大小的批处理。

  6. 仅使用GPU推理,可进一步利用CPU并行来优化。

综上,该框架尚需在模块化、性能、功能等多方面进一步优化。但作为初步实践和总结,已经反映出框架设计和实现的多方面考量,对后续改进具有很好的启发作用。

引用

[1] https://wrzpl.xet.tech/s/T30p4

① 全网独家视频课程

BEV感知、毫米波雷达视觉融合多传感器标定多传感器融合多模态3D目标检测点云3D目标检测目标跟踪Occupancy、cuda与TensorRT模型部署协同感知语义分割、自动驾驶仿真、传感器部署、决策规划、轨迹预测等多个方向学习视频(扫码即可学习

07a39e8cefe6cff523a3c932120c029e.png 视频官网:www.zdjszx.com

② 国内首个自动驾驶学习社区

近2000人的交流社区,涉及30+自动驾驶技术栈学习路线,想要了解更多自动驾驶感知(2D检测、分割、2D/3D车道线、BEV感知、3D目标检测、Occupancy、多传感器融合、多传感器标定、目标跟踪、光流估计)、自动驾驶定位建图(SLAM、高精地图、局部在线地图)、自动驾驶规划控制/轨迹预测等领域技术方案、AI模型部署落地实战、行业动态、岗位发布,欢迎扫描下方二维码,加入自动驾驶之心知识星球,这是一个真正有干货的地方,与领域大佬交流入门、学习、工作、跳槽上的各类难题,日常分享论文+代码+视频,期待交流!

f64b5a9123e4960b5bf089489606e763.png

③【自动驾驶之心】技术交流群

自动驾驶之心是首个自动驾驶开发者社区,聚焦目标检测、语义分割、全景分割、实例分割、关键点检测、车道线、目标跟踪、3D目标检测、BEV感知、多模态感知、Occupancy、多传感器融合、transformer、大模型、点云处理、端到端自动驾驶、SLAM、光流估计、深度估计、轨迹预测、高精地图、NeRF、规划控制、模型部署落地、自动驾驶仿真测试、产品经理、硬件配置、AI求职交流等方向。扫码添加汽车人助理微信邀请入群,备注:学校/公司+方向+昵称(快速入群方式)

10eaa74bd584b3ea41c3fa86558778e0.jpeg

④【自动驾驶之心】平台矩阵,欢迎联系我们!

79fbe72a5ab3aa33e5bbfa74bf71a56e.jpeg

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值