TensorRT~部署

最新推荐文章于 2024-08-25 11:42:55 发布

感谢型歌姬亚瑟

最新推荐文章于 2024-08-25 11:42:55 发布

阅读量61

点赞数

文章标签： java 开发语言

把整体进行框架性的封装, 最后实现下面的使用方法

// 创建一个worker的实例, 在创建的时候就完成初始化
auto worker   = thread::create_worker(onnxPath, level, params);

// 根据worker中的task类型进行推理
worker->inference("data/cat.png");

#封装logger

logger的作用是用于记录和打印日志的工具, 跟之前一样是继承了nvinfer1::ILogger的, 这个里面是封装了6个不同的级别的严重程度，在写的时候可以定义说，这里的DEBUG可以换成很多不同的，如LOGF、LOGE、LOGW、LOG、LOGV和LOGD

根据日志级别返回对应的严重性级别,这里的Severity是来源于nvinfer，这个类还提供了一些方法，如get_severity和get_level，用于在Level和Severity之间进行转换

转换的原因是因为你的理解是正确的。Severity是TensorRT库中定义的，它是nvinfer1::ILogger接口的一部分。当TensorRT运行时遇到问题或需要提供信息时，它会调用ILogger::log方法，并传入一个Severity参数来表示消息的重要性。

另一方面，Level是这个项目中定义的，它是Logger类的一部分。这个项目中的其他部分使Level来控制日志的打印级别。

Logger::Logger(Level level) {
    m_level = level;
    m_severity = get_severity(level);
}

Logger::Severity Logger::get_severity(Level level) {
    switch (level) {
        case Level::FATAL: return Severity::kINTERNAL_ERROR;
        case Level::ERROR: return Severity::kERROR;
        case Level::WARN:  return Severity::kWARNING;
        case Level::INFO:  return Severity::kINFO;
        case Level::VERB:  return Severity::kVERBOSE;
        default:           return Severity::kVERBOSE;
    }
}

Level Logger::get_level(Severity severity) {
    string str;
    switch (severity) {
        case Severity::kINTERNAL_ERROR: return Level::FATAL;
        case Severity::kERROR:          return Level::ERROR;
        case Severity::kWARNING:        return Level::WARN;
        case Severity::kINFO:           return Level::INFO;
        case Severity::kVERBOSE:        return Level::VERB;
    }
}

这些都是用来确认东西的, 确认完东西之后就需要把这些东西打印出来, 下面就是输出的内容了

void Logger::log (Severity severity, const char* msg) noexcept{
    /* 
        有的时候TensorRT给出的log会比较多并且比较细，所以我们选择将TensorRT的打印log的级别稍微约束一下
        - TensorRT的log级别如果是FATAL, ERROR, WARNING, 按照正常方式打印
        - TensorRT的log级别如果是INFO或者是VERBOSE的时候，只有当logger的level在大于VERBOSE的时候再打出
    */
    if (severity <= get_severity(Level::WARN)
        || m_level >= Level::DEBUG)
        __log_info(get_level(severity), "%s", msg);
}

void Logger::__log_info(Level level, const char* format, ...) {
    char msg[1000];
    va_list args;
    va_start(args, format);
    int n = 0;
    
    switch (level) {    
        case Level::DEBUG: n += snprintf(msg + n, sizeof(msg) - n, DGREEN "[debug]" CLEAR); break;
        case Level::VERB:  n += snprintf(msg + n, sizeof(msg) - n, PURPLE "[verb]" CLEAR); break;
        case Level::INFO:  n += snprintf(msg + n, sizeof(msg) - n, YELLOW "[info]" CLEAR); break;
        case Level::WARN:  n += snprintf(msg + n, sizeof(msg) - n, BLUE "[warn]" CLEAR); break;
        case Level::ERROR: n += snprintf(msg + n, sizeof(msg) - n, RED "[error]" CLEAR); break;
        default:           n += snprintf(msg + n, sizeof(msg) - n, RED "[fatal]" CLEAR); break;
    }

    n += vsnprintf(msg + n, sizeof(msg) - n, format, args);

    va_end(args);

    if (level <= m_level) 
        fprintf(stdout, "%s\n", msg);

    if (level <= Level::ERROR) {
        fflush(stdout);
        exit(0);
    }
}

#理解本框架三个最重要的模块: worker, classifier, model

trt_worker.cpp

这个文件实现了Worker类，这个类可以看作是一个工作线程，它负责加载ONNX模型，创建分类器，并执行推理。在Worker类的构造函数中，根据传入的任务类型参数，会创建相应的分类器实例。在inference方法中，会调用分类器的load_image和inference方法进行推理。

trt_model.cpp

这个文件实现了Model类，这是一个抽象基类，定义了所有模型类的公共接口和一些共享的实现。这个类的主要方法包括load_image（加载图像）、init_model（初始化模型）、build_engine（构建TensorRT引擎）、load_engine（加载已经构建的TensorRT引擎）和inference（执行推理）。这个类的实现主要是关于如何使用TensorRT库来加载和执行模型的。

trt_classifier.cpp

这个文件实现了Classifier类，这是一个具体的模型类，继承自Model类，用于分类任务。这个类重写了Model类的setup、preprocess_cpu、preprocess_gpu、postprocess_cpu和postprocess_gpu等方法，以实现分类任务的特定需求。从这三个文件的角色和功能来看，它们之间的关系可以简单地描述为：Worker类使用Model类（或其子类，如Classifier类）的实例来执行具体的任务。Model类定义了所有模型类的公共接口和一些共享的实现，而Classifier类则是Model类的一个具体实现，用于分类任务。

#Worker类的分析

从程序的执行流程来看，main.cpp中首先创建了Worker类的实例，然后通过Worker类的方法来使用Model类，所以我们先看怎么去使用worker这个工作线程，学习怎么去使用这个Model, 然后再去深入分析Model这个类

在Worker类中，有以下几个主要的成员：

m_logger
一个Logger类的智能指针，用于处理TensorRT的日志信息。
m_params
一个Params类的智能指针，用于存储模型的参数,里面包含了设备信息，类别信息，预处理的方式: 双线性插值/最近邻插值... 模型输入，任务类型, 精度，WorkSpace。
m_classifier
一个Classifier类的智能指针，用于执行分类任务，里面有做分类任务的参数，以及各种前后处理函数，包括cpu版本跟GPU版本。
m_scores
一个浮点数向量，用于保存分类任务的结果。Worker类的实现在 trt_worker.cpp 中。在Worker类的构造函数中，会根据传入的任务类型参数，创建相应的分类器实例。在inference方法中，会调用分类器的load_image和inference方法进行推理。

#ifndef __WORKER_HPP__
#define __WORKER_HPP__

#include <memory>
#include <vector>
#include "trt_model.hpp"
#include "trt_logger.hpp"
#include "trt_classifier.hpp"

namespace thread{

class Worker {
public:
    Worker(std::string onnxPath, logger::Level level, model::Params params);
    void inference(std::string imagePath);

public:
    std::shared_ptr<logger::Logger>          m_logger;
    std::shared_ptr<model::Params>           m_params;

    std::shared_ptr<model::classifier::Classifier>  m_classifier; // 因为今后考虑扩充为multi-task，所以各个task都是worker的成员变量
    std::vector<float>                       m_scores;     // 因为今后考虑会将各个multi-task间进行互动，所以worker需要保存各个task的结果
};

std::shared_ptr<Worker> create_worker(
    std::string onnxPath, logger::Level level, model::Params params);

}; //namespace thread

#endif //__WORKER_HPP__

worker实例化的时候读取onnx文件, 然后log的等级，还有传入的是param的参数，这些参数，然后剩下的是inference推理和create_worker创造一个工作的线程

inference整体上就是调用model命名空间下classifier命名空间下的Classfier(这里是多个命名控件嵌套), 这个类别里面有load_image和inference方法来完成读取图片和推理

然后创造一个工作线程create_worker()函数, 这个方法会使用一个智能指针指向一个实例, 然后这个实例就可以执行inference了具体的使用方法在main.cpp里面可以看到，简单的两行就可以实现图片的推理

// 创建一个worker的实例, 在创建的时候就完成初始化
auto worker   = thread::create_worker(onnxPath, level, params);

// 根据worker中的task类型进行推理
worker->inference("data/cat.png");

#include "trt_worker.hpp"
#include "trt_classifier.hpp"
#include "trt_logger.hpp"
#include "memory"

using namespace std;

namespace thread{

Worker::Worker(string onnxPath, logger::Level level, model::Params params) {
    m_logger = logger::create_logger(level);

    // 这里根据task_type选择创建的trt_model的子类，今后会针对detection, segmentation扩充
    if (params.task == model::task_type::CLASSIFICATION) 
        m_classifier = model::classifier::make_classifier(onnxPath, level, params);

}

void Worker::inference(string imagePath) {
    if (m_classifier != nullptr) {
        m_classifier->load_image(imagePath);
        m_classifier->inference();
    }
}

shared_ptr<Worker> create_worker(
    std::string onnxPath, logger::Level level, model::Params params) 
{
    // 使用智能指针来创建一个实例
    return make_shared<Worker>(onnxPath, level, params);
}

}; // namespace thread

#Model模块

首先在头文件里面这里定义了多个枚举 task_type, device, precision 来表示任务类型（如分类、检测）、计算设备（CPU/GPU）、精度（FP32/FP16）。

enum task_type {
    CLASSIFICATION,
    DETECTION,
    SEGMENTATION,
};

enum device {
    CPU,
    GPU
};

enum precision {
    FP32,
    FP16,
};

这里其实希望每一个model都希望有自己的image_info, 因为可能不同的model会有不同的优化策略, 前后处理的策略, 可能说都是分类但是会做不同的定制化处理

结构体定义:

image_info: 存储图像的高度、宽度和通道数。

Params: 存储模型构建时的参数，如设备类型、类别数、处理策略等。

// 对Params设定一些默认值, 这些Params是在build model中所需要的
struct Params {
    device            dev      = GPU;
    int               num_cls  = 1000;
    process::tactics  tac      = process::tactics::GPU_BILINEAR;
    image_info        img      = {224, 224, 3};
    task_type         task     = CLASSIFICATION;
    int               ws_size  = WORKSPACESIZE;
    precision         prec     = FP32;
};

这边加了一个函数，用来管理trt的指针，全部的trt api创建出来的指针都通过这个释放

/* 构建一个针对trt的shared pointer. 所有的trt指针的释放都是通过ptr->destroy完成*/
template<typename T>
void destroy_trt_ptr(T* ptr){
    if (ptr) {
        std::string type_name = typeid(T).name();
        LOGD("Destroy %s", type_name.c_str());
        ptr->destroy(); 
    };
}

当然后面会看到类似下面的输出, 这里就能够看出来我们把这里的runtime, engine, context的指针释放掉了, 不过要把输出改成debug模式才可以

[debug]Destroy N8nvinfer117IExecutionContextE
[debug]Destroy N8nvinfer111ICudaEngineE
[debug]Destroy N8nvinfer18IRuntimeE

Model类分析:

下面是一些通用的方法, 这些方法因为分类检测模型其实是差不多的，所以我们并不是很需要重写这些东西

load_image(std::string image_path): 加载用于推理的图像，保存其路径到 m_imagePath。

init_model(): 初始化模型，包括构建推理引擎、分配内存、创建执行上下文(context)和设置bindings（数据输入和输出的GPU地址）。

inference(): 执行模型推理，包括数据预处理、将数据传递到GPU（enqueue）进行推理计算、和后处理。

build_engine(): 构建TensorRT推理引擎，这是一个关键步骤，包括设置网络结构、优化配置和序列化引擎。load_engine(): 加载已存在的TensorRT引擎，如果引擎文件存在，则直接加载，否则会调用 build_engine 来构建。save_plan(nvinfer1::IHostMemory& plan): 保存序列化后的TensorRT引擎到磁盘，以便将来可以重新加载和使用。

print_network(nvinfer1::INetworkDefinition &network, bool optimized): 打印TensorRT网络的详细信息，可以选择在优化前后进行打印。

enqueue_bindings(): 执行推理计算，将预处理后的数据传递给GPU，并从GPU获取推理后的结果。

public:
    Model(std::string onnx_path, logger::Level level, Params params); 
    virtual ~Model() {};
    void load_image(std::string image_path);
    void init_model(); //初始化模型，包括build推理引擎, 分配内存，创建context, 设置bindings
    void inference();  //推理部分，preprocess-enqueue-postprocess

public:
    bool build_engine();
    bool load_engine();
    void save_plan(nvinfer1::IHostMemory& plan);
    void print_network(nvinfer1::INetworkDefinition &network, bool optimized);

    // 这里的dnn推理部分，只要设定好了m_bindings的话，不同的task的infer_dnn的实现都是一样的
    bool enqueue_bindings();

下面是一些纯虚方法, 这些方法因为不同类型的任务他是不一样的，所以需要在后面继承的时候自己写一个特定的方法, 例如说检测是640x640, 分类224x224

// 以下都是子类自己实现的内容, 通过定义一系列虚函数来实现
// setup负责分配host/device的memory, bindings, 以及创建推理所需要的上下文。
// 由于不同task的input/output的tensor不一样，所以这里的setup需要在子类实现
virtual void setup(void const* data, std::size_t size) = 0;

// 不同的task的前处理/后处理是不一样的，所以具体的实现放在子类
virtual bool preprocess_cpu()      = 0;
virtual bool preprocess_gpu()      = 0;
virtual bool postprocess_cpu()     = 0;
virtual bool postprocess_gpu()     = 0;

下面是一些nvinfer的成员变量

nvinfer1::Dims m_inputDims;
nvinfer1::Dims m_outputDims;
cudaStream_t   m_stream;

std::shared_ptr<logger::Logger>               m_logger;
std::shared_ptr<timer::Timer>                 m_timer;
std::shared_ptr<nvinfer1::IRuntime>           m_runtime;
std::shared_ptr<nvinfer1::ICudaEngine>        m_engine;
std::shared_ptr<nvinfer1::IExecutionContext>  m_context;
std::shared_ptr<nvinfer1::INetworkDefinition> m_network;

Model类的具体实现

构造函数初始化

Model::Model(string onnx_path, logger::Level level, Params params) {
    m_onnxPath      = onnx_path;
    m_enginePath    = getEnginePath(onnx_path);
    m_workspaceSize = WORKSPACESIZE;
    m_logger        = make_shared<logger::Logger>(level);
    m_timer         = make_shared<timer::Timer>();
    m_params        = new Params(params);
}

加载模型图片的信息输出到控制台

void Model::load_image(string image_path) {
    if (!fileExists(image_path)){
        LOGE("%s not found", image_path.c_str());
    } else {
        m_imagePath = image_path;
        LOG("Model:      %s", getFileName(m_onnxPath).c_str());
        LOG("Image:      %s", getFileName(m_imagePath).c_str());
    }
}

输出如下

还是跟之前一样的结构:

createInferBuilder->builder
builder->network
builder->config
network, config->parser
config->setMaxWorkspaceSize
config->setProfilingVerbosity
builder->platformHasFastFp16
config->setFlag
parser->parseFromFile
builder->buildEngineWithConfig
builder->buildSerializedNetwork
createInferRuntime->runtime
save_plan
setup
print_network

重点是这里的API调用都使用了之前的自动释放的函数

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;
}

下面是加载engine, 如果没有在文件路径找到这个engine, 就停止，然后自己的反序列化一个engine, 因为我们的目的是为了后面直接inference一个engine, 所以这里的话需要生成一个context, 创建一个cuda stream, 在下面没有看到是因为我们会把这个操作放在子类中的setup去实现，之前解释过，setup是必须实现的因为是纯虚函数

bool Model::load_engine() {
    // 同样的，我们也希望在load一个engine的时候就把一系列初始化全部做完，其中包括
    //  1. deserialize一个engine
    //  2. 创建一个context
    //  3. 创建推理所用的stream
    //  4. 创建推理所需要的device空间
    // 这样，我们就可以在load结束以后，就可以直接推理了。这样的写法会比较干净
    
    if (!fileExists(m_enginePath)) {
        LOGE("engine does not exits! Program terminated");
        return false;
    }

    vector<unsigned char> modelData;
    modelData     = loadFile(m_enginePath);
    
    // 根据runtime初始化engine, context, 以及memory
    setup(modelData.data(), modelData.size());

    return true;
}

这里是把engine保存到磁盘上的一个操作

void Model::save_plan(IHostMemory& plan) {
    auto f = fopen(m_enginePath.c_str(), "wb");
    fwrite(plan.data(), 1, plan.size(), f);
    fclose(f);
}

下面这里是一个推理，推理之前需要提前做一个预处理，这个预处理可以是CPU的也可以是GPU的，多种选择

void Model::inference() {
    if (m_params->dev == CPU) {
        preprocess_cpu();
    } else {
        preprocess_gpu();
    }

    enqueue_bindings();

    if (m_params->dev == CPU) {
        postprocess_cpu();
    } else {
        postprocess_gpu();
    }
}

bool Model::enqueue_bindings() {
    m_timer->start_gpu();
    if (!m_context->enqueueV2((void**)m_bindings, m_stream, nullptr)){
        LOG("Error happens during DNN inference part, program terminated");
        return false;
    }
    m_timer->stop_gpu();
    m_timer->duration_gpu("trt-inference(GPU)");
    return true;
}

下面就是打印信息，这里就是在构建引擎的时候是否选择打印优化过的层的信息

#Classfier模块

Classifier类的分析

这里就是Classifier继承了Model类，但是要留意的是这里的Classifer仍然是处于model的namespace命名空间，这里的构造函数就是我们在main.cpp会写到的东西, 使用就是使用这些

Classifier(std::string onnx_path, logger::Level level, Params params) : 
        Model(onnx_path, level, params) {};

这里是要继承的纯虚函数的内容override重写一次, 因为不同的模型写不同的setup, 前后处理是不一样的

virtual void setup(void const* data, std::size_t size) override;
virtual bool preprocess_cpu() override;
virtual bool preprocess_gpu() override;
virtual bool postprocess_cpu() override;
virtual bool postprocess_gpu() override;

Classifier的实现

这里不同模型之间可能性不一样的地方应该是input, output的dimension, 虽然现在检测器的输出也是一个了，但是我们提高性能的时候就直接sigmoid后面直接取了，就会有多个输出头，详情可以参考这个repo: NVIDIA-AI-IOT/yolov5_gpu_optimization

void Classifier::setup(void const* data, size_t size) {
    m_runtime     = shared_ptr<IRuntime>(createInferRuntime(*m_logger), destroy_trt_ptr<IRuntime>);
    m_engine      = shared_ptr<ICudaEngine>(m_runtime->deserializeCudaEngine(data, size), destroy_trt_ptr<ICudaEngine>);
    m_context     = shared_ptr<IExecutionContext>(m_engine->createExecutionContext(), destroy_trt_ptr<IExecutionContext>);
    m_inputDims   = m_context->getBindingDimensions(0);
    m_outputDims  = m_context->getBindingDimensions(1);
    // 考虑到大多数classification model都是1 input, 1 output, 这边这么写。如果像BEVFusion这种有多输出的需要修改

    CUDA_CHECK(cudaStreamCreate(&m_stream));
    
    m_inputSize     = m_params->img.h * m_params->img.w * m_params->img.c * sizeof(float);
    m_outputSize    = m_params->num_cls * sizeof(float);
    m_imgArea       = m_params->img.h * m_params->img.w;

    // 这里对host和device上的memory一起分配空间
    CUDA_CHECK(cudaMallocHost(&m_inputMemory[0], m_inputSize));
    CUDA_CHECK(cudaMallocHost(&m_outputMemory[0], m_outputSize));
    CUDA_CHECK(cudaMalloc(&m_inputMemory[1], m_inputSize));
    CUDA_CHECK(cudaMalloc(&m_outputMemory[1], m_outputSize));

    // //创建m_bindings，之后再寻址就直接从这里找
    m_bindings[0] = m_inputMemory[1];
    m_bindings[1] = m_outputMemory[1];
}

然后是前处理后处理的选择，这里可以选择GPU, 也可以选择CPU的这里的内容太多了，如果在写框架的时候不确定是否有对应的，那这里就直接return cpu版本吧，或者可以直接写不分设备的前后处理接口然后再重写，这里严谨只是为了之后添加模型一定有这些步骤

bool Classifier::preprocess_cpu() 
bool Classifier::preprocess_gpu() 
bool Classifier::postprocess_cpu()
bool Classifier::postprocess_gpu()

#总结

在本次文章总基于之前的的代码进行了许多改动，以提高代码的可复用性、可读性、安全性、可扩展性和可调试性。

代码可复用性：设计了一个推理框架，可以支持多种任务，如分类、检测、分割、姿态估计等。所有这些任务都有一个共同的流程：前处理 -> DNN推理 -> 后处理。这个框架设计了一个基类来实现这个基本操作，然后不同的任务可以继承这个基类，完成每个任务需要单独处理的内容。这是通过C++工厂模式的设计思路实现的。

可读性：为了提高代码的可读性，设计了一个接口worker进行推理。在主函数中，只需要创建一个worker，然后让worker读取图片，然后让worker做推理。worker内部可以根据主函数传入的参数，启动多种不同的任务，如分类推理、检测、分割等。

安全性：在设计框架的时候，需要做很多初始化，释放内存，以及处理错误调用。为了避免忘记释放内存、忘记处理某种错误调用、没有分配内存却释放了内存等问题，可以使用unique pointer或者shared pointer这种智能指针帮助我们管理内存的释放。同时，使用RAII设计机制，将资源的申请封装在一个对象的生命周期内，方便管理。这里通过下面这种方法调用destory来实现

auto engine        = shared_ptr<ICudaEngine>(builder->buildEngineWithConfig(*network, *config), destroy_trt_ptr<ICudaEngine>);

可扩展性：一个好的框架需要有很强的扩展性。这就意味着设计需要尽量模块化。当有新的任务出现的时候，可以做到最小限度的代码更改。在这里我们后面的检测器增加了setup, 前后处理就可以添加进来了。

可调试性：在设计框架的时候，希望能够为了让开发效率提高，在设计中比较关键的几个点设置debug信息，方便我们查看我们的模型是否设计的有问题。可以实现一个logger来方便我们管理这些。logger可以通过传入的不同参数显示不同级别的日志信息。这里我们通过转换Level, Servity来实现跟TensorRT上下文通信解决这个问题。

原创作者: whaosoft 转载于: https://blog.51cto.com/whaosoft/11642410