六. 部署分类器-deploy-classification-advanced

前言

自动驾驶之心推出的《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第六章—部署分类器，一起来学习优化上节课的分类器部署代码

课程大纲可以看下面的思维导图

在这里插入图片描述

0. 简述

本小节目标：优化分类器部署代码

学习推理框架的设计，优化 6.1 案例小节代码

下面我们开始本次课程的学习🤗

1. 案例运行

在正式开始课程之前，博主先带大家跑通 6.2-deploy-classification-advanced 这个小节的案例🤗

源代码获取地址：https://github.com/kalfazed/tensorrt_starter

首先大家需要把 tensorrt_starter 这个项目给 clone 下来，指令如下：

git clone https://github.com/kalfazed/tensorrt_starter.git

也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2024/7/14 日，若有改动请参考最新）

整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV，大家可以参考 Ubuntu20.04软件安装大全进行相应软件的安装，博主这里不再赘述

假设你的项目、环境准备完成，下面我们一起来运行下 6.2-deploy-classification-advanced 小节案例代码

开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter6-deploy-classification-and-inference-design/6.2-deploy-classification-advanced 小节中创建一个 models 文件夹，接着在 models 文件夹下创建一个 onnx 和 engine 文件夹，总共三个文件夹需要创建

创建完后 6.2 小节整个目录结构如下：

在这里插入图片描述

接着我们需要执行 python 脚本创建一个 ONNX 模型，先进入到 6.2 小节中：

cd tensorrt_starter/chapter6-deploy-classification-and-inference-design/6.2-deploy-classification-advanced

执行如下指令：

python src/python/export_pretrained.py -d ./models/onnx/

Note：大家需要准备一个虚拟环境，安装好 torch、onnx、onnxsim 等第三方库

输出如下：

在这里插入图片描述

生成好的 resnet50.onnx 模型文件保存在 models/onnx 文件夹下，大家可以查看

接着我们需要利用 ONNX 生成对应的 engine 完成推理，在此之前我们需要修改下整体的 Makefile.config，指定一些库的路径：

# tensorrt_starter/config/Makefile.config
# CUDA_VER                    :=  11
CUDA_VER                    :=  11.6# opencv和TensorRT的安装目录
OPENCV_INSTALL_DIR          :=  /usr/local/include/opencv4
# TENSORRT_INSTALL_DIR        :=  /mnt/packages/TensorRT-8.4.1.5
TENSORRT_INSTALL_DIR        :=  /home/jarvis/lean/TensorRT-8.6.1.6

Note：大家查看自己的 CUDA 是多少版本，修改为对应版本即可，另外 OpenCV 和 TensorRT 修改为你自己安装的路径即可

接着我们就可以来执行编译，指令如下：

make -j64

输出如下：

在这里插入图片描述

接着执行：

./bin/trt-infer

输出如下：

在这里插入图片描述

我们这里可以看到每张图片推理的结果以及置信度

Note：这里大家可以测试下其它的 ONNX 模型看下推理结果，这里博主准备了导出好的各个分类模型的 ONNX，大家可以点击 here 下载

如果大家能够看到上述输出结果，那就说明本小节案例已经跑通，下面我们就来看看具体的代码实现

2. 补充说明

在分析代码之前我们先来看下韩君老师在这小节中写的 README 文档

本小节的代码是基于 6.1-deploy-classification 案例的优化，这里做了很多改动，具体体现在以下几个方面：

1. 代码可复用性
2. 可读性
3. 安全性
4. 可扩展性
5. 可调试性

这里围绕着这些点简单展开一下：

1. 代码可复用性

我们设计一个推理框架时会希望它能够支持多 task，比如 classification、detection、segmentation、pose estimation 等等，这些 task 都有 前处理➡DNN 推理➡后处理 这么一套流程，只不过不同的 task 的前后处理会有所不同。同时我们在 build 一个推理引擎的时候都是通过 builder->network->config->parser->serialize->save file 这么一套流程，在 infer 的时候则是 load file->deserialize->engine->context->enqueue 这么一套流程，在初始化的时候都是要根据 trt engine 的信息 分配 host memory，分配 device memory

那么我们就会很自然的想到是否可以设计一个类来实现这一套最基本的操作，之后不同的 task 来继承这个类去完成每一个 task 各自需要单独处理的内容，这个我们可以通过 C++ 工厂模式 的设计思路去搭建模型，实现封装

2. 可读性

我们希望我们的代码有比较好的可读性，这意味着我们在设计框架的时候需要尽量通过接口来暴露或者隐蔽一些功能，比如说我们可以使用 worker 作为接口进行推理，在 main 中我们只需要做到 创建一个 worker➡worker 读取图片➡worker 做推理 就好了。同时 worker 也只暴露这些接口，在 worker 内部我们可以让 worker 根据 main 函数传入的参数启动多种不同的 task，比如 worker➡classifier 负责做分类推理，worker➡detector 负责做检测，worker➡segmentor 负责做分割，worker 中会根据不同的 task 做对应的实现

3. 安全性

我们在设计框架的时候需要做初始化、内存释放以及针对错误调用时的处理，当代码规模较小时安全性意识其实并不是很复杂，但是当代码的规模比较大，有很多个 task 需要我们实现时难免会遇到类似于忘记释放内存，忘记对某一种调用做 error handler，没有分配内存却释放了内存等等情况，为了避免这些情况的发生，我们可以使用 unique pointer 或者 shared pointer 这种智能指针帮助我们管理内存的使用。同时使用 RAII 设计机制 将资源的申请封装在一个对象的生命周期内方便管理，RAII 是 Resource Acquisition Is Initialization 的缩写，中文译为资源获取即初始化，比较常见的方法就是在一个类的构造函数的时候把一系列初始化完成

4. 可扩展性

一个比较好的框架需要有很强的扩展性，这就意味着我们的设计需要尽量模块化，当有新的 task 出现的时候我们可以做到最小限度的代码更改，比如说：

worker(...);                                         //根据模型的种类(分类、检测、分割)来初始化一个模型worker->classifier(...);                             //资源获取即初始化，在这里创建engine，并且建立推理上下文。如果已经有了engine的话就直接load这个engine，并且建立推理上下文
worker->classifier.load_image(...);                  //在这里我们读取图片，并分配pinned memory，分配device memory
score = worker->classifier.infer_classifier(...);    //在这里我们进行预处理，推理，后处理的部分worker->detector(...);                               //资源获取即初始化，在这里创建engine，并且建立推理上下文。如果已经有了engine的话就直接load这个engine，并且建立推理上下文
worker->detector.load_image(...);                    //在这里我们读取图片，并分配pinned memory，分配device memory
bboxes = worker->detector.infer_detector(...);       //在这里我们进行预处理，推理，后处理的部分worker->segmenter(...);                              //资源获取即初始化，在这里创建engine，并且建立推理上下文。如果已经有了engine的话就直接load这个engine，并且建立推理上下文
worker->segmenter.load_image(...);                   //在这里我们读取图片，并分配pinned memory，分配device memory
mask = worker->segmenter.infer_segmentor(...);       //在这里我们进行预处理，推理，后处理的部分worker->drawBBox(...);                               //worker负责将bbox的信息绘制在原图上
worker->drawMask(...);                               //worker负责将mask的信息融合在原图上
worker->drawScore(...);                              //worker负责将score的信息绘制在原图上

当然这个小节案例下的 worker 的内容是很空的，因为目前只是一个单独的 classification 的任务，整体上的结构比较 simple，但这么设计的目的是为了今后的扩展，比如说针对视频流的异步处理，多线程的处理，multi-stage model 的处理等等，在今后的应用中很可能会出现下面这种情况：

// 1st stage detection
worker->detector(...);
worker->detector.load_image(...);
bboxes = worker->detector.infer_detector(...);// 2nd stage classification
worker->classifier(...);
worker->classifier.load_from_bbox(...);
score = worker->classifier.infer_classifier(...);

5. 可调试性

我们在设计框架的时候为了让开发效率提高，需要考虑在框架中的几个比较关键的位置设置 debug 信息，方便查看我们的模型是否设计有问题，比如说 YOLO 检测模型在经过 NMS 处理之后 bbox 数量还剩多少，ONNX 模型在经过 TensorRT 优化以后的网络结构是什么样子，模型各个 layer 所支持的输入 data layout 是 NCHW 还是 NHWC 等等。我们可以实现一个 logger 来方便我们管理这些，logger 可以通过传入的不同参数显示不同级别的日志信息，比如说如果我们在 main 中声明想打印有关 VERBOSE 信息，我们就可以打印在代码中所有以 LOGV() 显示的信息

OK，接下来我们就来分析下代码中是如何体现上面提到的几点内容的

3. 代码分析

3.1 main.cpp

我们先从 main.cpp 看起：

#include "trt_model.hpp"
#include "trt_logger.hpp"
#include "trt_worker.hpp"
#include "utils.hpp"using namespace std;int main(int argc, char const *argv[])
{/*这么实现目的在于让调用的整个过程精简化*/string onnxPath    = "models/onnx/resnet50.onnx";auto level         = logger::Level::INFO;auto params        = model::Params();params.img         = {224, 224, 3};params.num_cls     = 1000;params.task        = model::task_type::CLASSIFICATION;params.dev         = model::device::GPU;// 创建一个worker的实例, 在创建的时候就完成初始化auto worker   = thread::create_worker(onnxPath, level, params);// 根据worker中的task类型进行推理worker->inference("data/cat.png");worker->inference("data/gazelle.png");worker->inference("data/eagle.png");worker->inference("data/fox.png");worker->inference("data/tiny-cat.png");worker->inference("data/wolf.png");return 0;
}

我们在 main 函数首先设置了模型路径，日志级别以及推理参数的初始化，包括输入大小，分类的类别数、任务类型、前后处理处理方式（CPU/GPU），这么实现的目的在于让调用的整个过程精简化

接着调用 thread 命名空间下的 create_worker 函数创建了一个推理工作器 worker，在创建过程中工作器就完成了必要的初始化步骤，例如模型加载、推理引擎构建等，最后调用 worker 类的 inference 方法对传入的图像进行推理

这个 main.cpp 相比于上节案例代码以更加简洁的方式进行模型的推理，通过模块化的设计将日志、模型加载、推理工作器、推理任务等不同功能分离开来，主程序通过简单的几行代码就可以完成模型的推理过程

3.2 trt_worker.cpp

接下来我们看下 create_worker 函数是怎么做的：

#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

首先在 Worker 类的构造函数中我们根据提供的日志级别 level 创建了一个日志记录器 m_logger，接着根据任务类型创建相应的推理模型，这里主要针对的是分类模型，因此调用了 make_classifier 工厂函数，根据传入的 ONNX 模型路径、日志级别以及参数创建一个分类器实例并赋值给 m_classifier，这个分类器将用于加载图像并进行推理。

这个部分体现的是代码安全性的 RAII 部分，资源获取即初始化，拿到 Wokrer 类的实例对象时就已经完成了初始化，当然这里还体现了代码的可扩展性，可以通过不同的 task_type 创建不同的 trt_model 进行推理

在 Worker 类的 inference 推理函数中我们先检查 m_classifier 是否已经正确初始化（即指针是否为空），如果分类器存在，则执行接下来的推理操作。先调用分类器的 load_image 方法，加载指定路径的图像，这个步骤通常包括图像的预处理如缩放、归一化等，加载图像后调用 inference 方法进行推理，这一步将执行实际的模型推理过程，生成输出结果

create_worker 函数主要用来创建 Worker 类的实例对象，使用 std::make_shared 创建一个 Worker 实例，通过这种方式创建的 Worker 对象，其生命周期将由 shared_ptr 管理，确保在不再使用时自动释放资源

这种方法体现了代码的安全性，简化了对象的管理，通过 create_worker 函数返回一个 Worker 对象的 shared_ptr，可以在调用放轻松管理 Worker 的生命周期，而不需要显示的 new 和 delete 操作

这个 create_worker 函数以及 Worker 类的实现展示了一个简洁而灵活的设计，能够根据不同的任务（如分类、检测、分割等）创建合适的推理模型，并通过智能指针来管理对象的生命周期。这种设计方式非常适合需要处理多种模型推理任务的应用场景，同时保持代码的简洁和易维护性。

未来若要支持更多的任务类型，只需在 Worker 构造函数中添加相应的逻辑，并在 trt_worker.hpp 文件中扩展支持的任务模型类即可，这种设计的扩展性和可维护性都非常高。

3.3 trt_logger.cpp

接下来我们看 Worker 类的构造函数中 create_logger 创建日志器的实现：

#include "trt_logger.hpp"
#include "NvInfer.h"
#include <cstdlib>using namespace std;namespace logger {Level Logger::m_level = Level::INFO;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);}
}shared_ptr<Logger> create_logger(Level level) {return make_shared<Logger>(level);
}} // namespace logger

和 create_worker 一样 create_logger 函数也是使用 std::make_shared 创建一个 Logger 实例，并返回一个 shared_ptr，用于管理 Logger 对象的生命周期

这个 create_logger 函数和 Logger 类的实现展示了一个灵活且高效的日志记录系统，能够根据不同的日志级别动态调整日志输出，通过将日志级别映射到 TensorRT，可以很好地与 TensorRT 的日志系统集成，并确保只输出需要的日志信息。同时，通过使用智能指针管理 Logger 对象，代码变得更加安全和易于维护

3.4 trt_classifier.cpp

接下来我们看 Worker 类的构造函数中 make_classifier 创建分类器的实现：

#include "opencv2/imgproc.hpp"
#include "trt_model.hpp"
#include "utils.hpp" 
#include "trt_logger.hpp"#include "NvInfer.h"
#include "NvOnnxParser.h"
#include <string>#include "opencv2/core/core.hpp"
#include "opencv2/highgui/highgui.hpp"
#include "opencv2/opencv.hpp"
#include "imagenet_labels.hpp"
#include "trt_classifier.hpp"
#include "trt_preprocess.hpp"
#include "utils.hpp"using namespace std;
using namespace nvinfer1;namespace model{namespace classifier {/*classification model的初始化相关内容。包括设置input/output bindings, 分配host/device的memory等
*/
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];
}bool Classifier::preprocess_cpu() {/*Preprocess -- 获取mean, std*/float mean[]       = {0.406, 0.456, 0.485};float std[]        = {0.225, 0.224, 0.229};/*Preprocess -- 读取数据*/cv::Mat input_image;input_image = cv::imread(m_imagePath);if (input_image.data == nullptr) {LOGE("ERROR: Image file not founded! Program terminated"); return false;}/*Preprocess -- 测速*/m_timer->start_cpu();/*Preprocess -- resize(默认是bilinear interpolation)*/cv::resize(input_image, input_image, cv::Size(m_params->img.w, m_params->img.h), 0, 0, cv::INTER_LINEAR);/*Preprocess -- host端进行normalization和BGR2RGB, NHWC->NCHW*/int index;int offset_ch0 = m_imgArea * 0;int offset_ch1 = m_imgArea * 1;int offset_ch2 = m_imgArea * 2;for (int i = 0; i < m_inputDims.d[2]; i++) {for (int j = 0; j < m_inputDims.d[3]; j++) {index = i * m_inputDims.d[3] * m_inputDims.d[1] + j * m_inputDims.d[1];m_inputMemory[0][offset_ch2++] = (input_image.data[index + 0] / 255.0f - mean[0]) / std[0];m_inputMemory[0][offset_ch1++] = (input_image.data[index + 1] / 255.0f - mean[1]) / std[1];m_inputMemory[0][offset_ch0++] = (input_image.data[index + 2] / 255.0f - mean[2]) / std[2];}}/*Preprocess -- 将host的数据移动到device上*/CUDA_CHECK(cudaMemcpyAsync(m_inputMemory[1], m_inputMemory[0], m_inputSize, cudaMemcpyKind::cudaMemcpyHostToDevice, m_stream));m_timer->stop_cpu();m_timer->duration_cpu<timer::Timer::ms>("preprocess(CPU)");return true;
}bool Classifier::preprocess_gpu() {/*Preprocess -- 获取mean, std*/float mean[]       = {0.406, 0.456, 0.485};float std[]        = {0.225, 0.224, 0.229};/*Preprocess -- 读取数据*/cv::Mat input_image;input_image = cv::imread(m_imagePath);if (input_image.data == nullptr) {LOGE("ERROR: file not founded! Program terminated"); return false;}/*Preprocess -- 测速*/m_timer->start_gpu();/*Preprocess -- 使用GPU进行双线性插值, 并将结果返回到m_inputMemory中*/process::preprocess_resize_gpu(input_image, m_inputMemory[1],m_params->img.h, m_params->img.w, mean, std, process::tactics::GPU_BILINEAR);m_timer->stop_gpu();m_timer->duration_gpu("preprocess(GPU)");return true;
}bool Classifier::postprocess_cpu() {/*Postprocess -- 测速*/m_timer->start_cpu();/*Postprocess -- 将device上的数据移动到host上*/int output_size    = m_params->num_cls * sizeof(float);CUDA_CHECK(cudaMemcpyAsync(m_outputMemory[0], m_outputMemory[1], output_size, cudaMemcpyKind::cudaMemcpyDeviceToHost, m_stream));CUDA_CHECK(cudaStreamSynchronize(m_stream));/*Postprocess -- 寻找label*/ImageNetLabels labels;int pos = max_element(m_outputMemory[0], m_outputMemory[0] + m_params->num_cls) - m_outputMemory[0];float confidence = m_outputMemory[0][pos] * 100;m_timer->stop_cpu();m_timer->duration_cpu<timer::Timer::ms>("postprocess(CPU)");LOG("Inference result: %s", labels.imagenet_labelstring(pos).c_str());   LOG("Confidence is %.3f%%\n", confidence);   return true;
}bool Classifier::postprocess_gpu() {/*由于classification task的postprocess比较简单，所以CPU/GPU的处理这里用一样的对于像yolo这种detection model, postprocess会包含decode, nms这些处理。可以选择在CPU还是在GPU上跑*/return postprocess_cpu();}shared_ptr<Classifier> make_classifier(std::string onnx_path, logger::Level level, model::Params params)
{auto classifier = make_shared<Classifier>(onnx_path, level, params);classifier->init_model();return classifier;
}}; // namespace classifier}; // namespace model

make_classifier 函数用于创建一个 Classifier 对象，用于进行分类任务的推理。这个函数封装了 Classifier 的创建、初始化等一系列操作，确保 Classifier 对象在返回前已经准备好进行推理任务

Classifier::setup 函数主要是进行一系列的初始化工作，包括推理引擎和执行上下文的创建，输入和输出张量维度的获取，CUDA 流的创建，内存分配以及输入输出内存的绑定

Classifier::preprocess_cpu 函数则是我们前面 6.0 案例提到过的 CPU 进行图像预处理的代码

Classifier::postprocess_cpu 函数主要将 engine 推理结果从 device 上传输到 host 上，然后解析推理结果，找到概率最大的类别索引和其置信度

make_classifier 函数使用 std::make_shared 创建 Classifier 对象，并调用 init_model() 方法加载 ONNX 模型生成 engine，最后返回已经初始化好的 Classifier 对象

make_classifier 函数通过一系列步骤创建并初始化一个 Classifier 对象，Classifier 类的实现则处理了从模型加载、内存管理到推理前后的处理逻辑。通过将这些功能模块化，make_classifier 函数提供了一个易于使用且灵活的接口

3.5 trt_model.cpp

接下来我们来分析下 make_classifier 中 init_model() 方法的实现：

#include "trt_model.hpp"
#include "utils.hpp" 
#include "trt_logger.hpp"#include "NvInfer.h"
#include "NvOnnxParser.h"
#include <string>using namespace std;
using namespace nvinfer1;
using namespace nvonnxparser;namespace 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());}
}void Model::init_model() {/* 一个model的engine, context这些一旦创建好了，当多次调用这个模型的时候就没必要每次都初始化了*/if (m_context == nullptr){if (!fileExists(m_enginePath)){LOGV("%s not found. Building trt engine...", m_enginePath.c_str());build_engine();} else {LOGV("%s has been generated! loading trt engine...", m_enginePath.c_str());load_engine();}}
}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>);// 保存序列化后的enginesave_plan(*plan);// 根据runtime初始化engine, context, 以及memorysetup(plan->data(), plan->size());// 把优化前和优化后的各个层的信息打印出来LOGV("Before TensorRT optimization");print_network(*network, false);LOGV("After TensorRT optimization");print_network(*network, true);return true;
}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, 以及memorysetup(modelData.data(), modelData.size());return true;
}void Model::save_plan(IHostMemory& plan) {auto f = fopen(m_enginePath.c_str(), "wb");fwrite(plan.data(), 1, plan.size(), f);fclose(f);
}/* 可以根据情况选择是否在CPU上跑pre/postprocess对于一些edge设备，为了最大化GPU利用效率，我们可以考虑让CPU做一些pre/postprocess，让其执行与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;
}void Model::print_network(INetworkDefinition &network, bool optimized) {int inputCount = network.getNbInputs();int outputCount = network.getNbOutputs();string layer_info;for (int i = 0; i < inputCount; i++) {auto input = network.getInput(i);LOGV("Input info: %s:%s", input->getName(), printTensorShape(input).c_str());}for (int i = 0; i < outputCount; i++) {auto output = network.getOutput(i);LOGV("Output info: %s:%s", output->getName(), printTensorShape(output).c_str());}int layerCount = optimized ? m_engine->getNbLayers() : network.getNbLayers();LOGV("network has %d layers", layerCount);if (!optimized) {for (int i = 0; i < layerCount; i++) {char layer_info[1000];auto layer   = network.getLayer(i);auto input   = layer->getInput(0);int n = 0;if (input == nullptr){continue;}auto output  = layer->getOutput(0);LOGV("layer_info: %-40s:%-25s->%-25s[%s]", layer->getName(),printTensorShape(input).c_str(),printTensorShape(output).c_str(),getPrecision(layer->getPrecision()).c_str());}} else {auto inspector = shared_ptr<IEngineInspector>(m_engine->createEngineInspector());for (int i = 0; i < layerCount; i++) {LOGV("layer_info: %s", inspector->getLayerInformation(i, nvinfer1::LayerInformationFormat::kJSON));}}
}} // namespace model

init_model() 方法是 Model 类中用于初始化模型的重要方法，它负责确保模型已经准备好进行推理，该方法主要包括两个部分：检查模型是否已经加载或生成，以及在必要时构建或加载推理引擎

在 init_model() 方法中首先会检查 context 上下文是否为空，如果为空则说明模型尚未初始化，需要进行引擎构建或加载操作。接着检查引擎文件是否存在，如果不存在则需要调用 build_engine 方法构建并生成新的 TensorRT 引擎，并将其保存到文件中，如果存在则调用 load_engine 方法直接从文件中加载已生成的 TensorRT 引擎

build_engine 方法负责从 ONNX 模型文件构建 TensorRT 引擎，沿用的是我们之前的代码，这边博主就不再赘述了，同样 load_engine 也是之前的代码，大家可以简单再看看

init_model() 方法确保模型在首次使用时已经准备好进行推理，它通过调用 builder_engine() 或 load_engine 来创建或加载 TensorRT 引擎，并在成功构建或加载引擎后初始化推理环境。这种设计不仅简化了模型的使用，还通过缓存已经生成的引擎提高了模型加载的效率

3.6 inference部分

值得注意的是以上的分析只是创建一个 worker 实例时所做的事情，还没有真正做推理，下面我们就来看看 inference 时所做的事情：

worker->inference("data/cat.png");

首先我们调用 worker 实例的 inference 方法进行推理

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

在 inference 方法中会先调用 m_classifier 类的 load_image 方法加载图像：

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());LOG("Precision:  %s", getPrec(m_params->prec).c_str());}
}

本质上调用的是 Model 父类的 load_image 方法，将 image_path 赋值给成员变量 m_imagePath，并打印一些基本信息

接着会调用 m_classifier 类的 inference 方法进行推理：

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

当然这里也是调用的 Model 父类的 inference 方法，在 Model 类的 inference 方法中首先根据参数选择预处理是在 CPU 上完成还是在 GPU 上完成

我们先看 CPU 预处理函数：

bool Classifier::preprocess_cpu() {/*Preprocess -- 获取mean, std*/float mean[]       = {0.406, 0.456, 0.485};float std[]        = {0.225, 0.224, 0.229};/*Preprocess -- 读取数据*/cv::Mat input_image;input_image = cv::imread(m_imagePath);if (input_image.data == nullptr) {LOGE("ERROR: Image file not founded! Program terminated"); return false;}/*Preprocess -- 测速*/m_timer->start_cpu();/*Preprocess -- resize(默认是bilinear interpolation)*/cv::resize(input_image, input_image, cv::Size(m_params->img.w, m_params->img.h), 0, 0, cv::INTER_LINEAR);/*Preprocess -- host端进行normalization和BGR2RGB, NHWC->NCHW*/int index;int offset_ch0 = m_imgArea * 0;int offset_ch1 = m_imgArea * 1;int offset_ch2 = m_imgArea * 2;for (int i = 0; i < m_inputDims.d[2]; i++) {for (int j = 0; j < m_inputDims.d[3]; j++) {index = i * m_inputDims.d[3] * m_inputDims.d[1] + j * m_inputDims.d[1];m_inputMemory[0][offset_ch2++] = (input_image.data[index + 0] / 255.0f - mean[0]) / std[0];m_inputMemory[0][offset_ch1++] = (input_image.data[index + 1] / 255.0f - mean[1]) / std[1];m_inputMemory[0][offset_ch0++] = (input_image.data[index + 2] / 255.0f - mean[2]) / std[2];}}/*Preprocess -- 将host的数据移动到device上*/CUDA_CHECK(cudaMemcpyAsync(m_inputMemory[1], m_inputMemory[0], m_inputSize, cudaMemcpyKind::cudaMemcpyHostToDevice, m_stream));m_timer->stop_cpu();m_timer->duration_cpu<timer::Timer::ms>("preprocess(CPU)");return true;
}

我们前面分析过这部分内容就是 6.0 小节案例的代码，通过指针访问图像的每个像素完成预处理操作，下面我们重点看下 GPU 上的预处理过程：

bool Classifier::preprocess_gpu() {/*Preprocess -- 获取mean, std*/float mean[]       = {0.406, 0.456, 0.485};float std[]        = {0.225, 0.224, 0.229};/*Preprocess -- 读取数据*/cv::Mat input_image;input_image = cv::imread(m_imagePath);if (input_image.data == nullptr) {LOGE("ERROR: file not founded! Program terminated"); return false;}/*Preprocess -- 测速*/m_timer->start_gpu();/*Preprocess -- 使用GPU进行双线性插值, 并将结果返回到m_inputMemory中*/process::preprocess_resize_gpu(input_image, m_inputMemory[1],m_params->img.h, m_params->img.w, mean, std, process::tactics::GPU_BILINEAR);m_timer->stop_gpu();m_timer->duration_gpu("preprocess(GPU)");return true;
}

在 preprocess_gpu 方法中我们主要是调用了 process 命名空间下的 preprocess_resize_gpu 函数，其内容如下：

// 根据比例进行缩放 (GPU版本)
void preprocess_resize_gpu(cv::Mat &h_src, float* d_tar, const int& tar_h, const int& tar_w, float* h_mean, float* h_std, tactics tac) 
{float*   d_mean = nullptr;float*   d_std  = nullptr;uint8_t* d_src  = nullptr;int height   = h_src.rows;int width    = h_src.cols;int chan     = 3;int src_size  = height * width * chan * sizeof(uint8_t);int norm_size = 3 * sizeof(float);// 分配device上的src和mean, std的内存CUDA_CHECK(cudaMalloc(&d_src, src_size));CUDA_CHECK(cudaMalloc(&d_mean, norm_size));CUDA_CHECK(cudaMalloc(&d_std, norm_size));// 将数据拷贝到device上CUDA_CHECK(cudaMemcpy(d_src, h_src.data, src_size, cudaMemcpyHostToDevice));CUDA_CHECK(cudaMemcpy(d_mean, h_mean, norm_size, cudaMemcpyHostToDevice));CUDA_CHECK(cudaMemcpy(d_std, h_std, norm_size, cudaMemcpyHostToDevice));// device上处理resize, BGR2RGB的核函数resize_bilinear_gpu(d_tar, d_src, tar_w, tar_h, width, height, d_mean, d_std, tac);// host和device进行同步处理CUDA_CHECK(cudaDeviceSynchronize());CUDA_CHECK(cudaFree(d_std));CUDA_CHECK(cudaFree(d_mean));CUDA_CHECK(cudaFree(d_src));// 因为接下来会继续在gpu上进行处理，所以这里不用把结果返回到host
}

在 preprocess_resize_gpu 函数中在 device 上分配了 image、mean 和 std 的内存，并将 host 上的数据拷贝到 device 上，最终调用的是 resize_bilinear_gpu 即核函数的启动函数来处理，实现如下所示：

void resize_bilinear_gpu(float* d_tar, uint8_t* d_src, int tarW, int tarH, int srcW, int srcH, float* d_mean, float* d_std,process::tactics tac) 
{dim3 dimBlock(32, 32, 1);dim3 dimGrid(tarW / 32 + 1, tarH / 32 + 1, 1);//scaled resizefloat scaled_h = (float)srcH / tarH;float scaled_w = (float)srcW / tarW;float scale = (scaled_h > scaled_w ? scaled_h : scaled_w);switch (tac) {case process::tactics::GPU_NEAREST:nearest_BGR2RGB_nhwc2nchw_norm_kernel <<<dimGrid, dimBlock>>>(d_tar, d_src, tarW, tarH, srcW, srcH, scaled_w, scaled_h, d_mean, d_std);break;case process::tactics::GPU_NEAREST_CENTER:nearest_BGR2RGB_nhwc2nchw_norm_kernel <<<dimGrid, dimBlock>>>(d_tar, d_src, tarW, tarH, srcW, srcH, scale, scale, d_mean, d_std);break;case process::tactics::GPU_BILINEAR:bilinear_BGR2RGB_nhwc2nchw_norm_kernel <<<dimGrid, dimBlock>>> (d_tar, d_src, tarW, tarH, srcW, srcH, scaled_w, scaled_h, d_mean, d_std);break;case process::tactics::GPU_BILINEAR_CENTER:bilinear_BGR2RGB_nhwc2nchw_shift_norm_kernel <<<dimGrid, dimBlock>>> (d_tar, d_src, tarW, tarH, srcW, srcH, scale, scale, d_mean, d_std);break;default:LOGE("ERROR: Wrong GPU resize tactics selected. Program terminated");exit(1);}
}

resize_bilinear_gpu 是一个启动 CUDA 核函数的主函数，用于在 GPU 上执行图像的缩放操作，该函数根据不同的策略 tactics 选择不同的 CUDA 核函数来执行特定的图像缩放和预处理操作

这里我们选择的预处理策略是 GPU_BILINEAR，调用的是 bilinear_BGR2RGB_nhwc2nchw_norm_kernel 核函数，其内容如下：

__global__ void bilinear_BGR2RGB_nhwc2nchw_norm_kernel(float* tar, uint8_t* src, int tarW, int tarH, int srcW, int srcH, float scaled_w, float scaled_h,float* d_mean, float* d_std) 
{// bilinear interpolation -- resized之后的图tar上的坐标int x = blockIdx.x * blockDim.x + threadIdx.x;int y = blockIdx.y * blockDim.y + threadIdx.y;// // bilinear interpolation -- 计算x,y映射到原图时最近的4个坐标int src_y1 = floor((y + 0.5) * scaled_h - 0.5);int src_x1 = floor((x + 0.5) * scaled_w - 0.5);int src_y2 = src_y1 + 1;int src_x2 = src_x1 + 1;if (src_y1 < 0 || src_x1 < 0 || src_y2 > srcH || src_x2 > srcW) {// bilinear interpolation -- 对于越界的坐标不进行计算} else {// bilinear interpolation -- 计算原图上的坐标(浮点类型)在0~1之间的值float th   = ((y + 0.5) * scaled_h - 0.5) - src_y1;float tw   = ((x + 0.5) * scaled_w - 0.5) - src_x1;// bilinear interpolation -- 计算面积(这里建议自己手画一张图来理解一下)float a1_1 = (1.0 - tw) * (1.0 - th);  //右下float a1_2 = tw * (1.0 - th);          //左下float a2_1 = (1.0 - tw) * th;          //右上float a2_2 = tw * th;                  //左上// bilinear interpolation -- 计算4个坐标所对应的索引int srcIdx1_1 = (src_y1 * srcW + src_x1) * 3;  //左上int srcIdx1_2 = (src_y1 * srcW + src_x2) * 3;  //右上int srcIdx2_1 = (src_y2 * srcW + src_x1) * 3;  //左下int srcIdx2_2 = (src_y2 * srcW + src_x2) * 3;  //右下// bilinear interpolation -- 计算resized之后的图的索引int tarIdx    = y * tarW  + x;int tarArea   = tarW * tarH;// bilinear interpolation -- 实现bilinear interpolation的resize + BGR2RGB + NHWC2NCHW normalization// 注意，这里tar和src进行遍历的方式是不一样的tar[tarIdx + tarArea * 0] = (round((a1_1 * src[srcIdx1_1 + 2] + a1_2 * src[srcIdx1_2 + 2] +a2_1 * src[srcIdx2_1 + 2] +a2_2 * src[srcIdx2_2 + 2])) / 255.0f - d_mean[2]) / d_std[2];tar[tarIdx + tarArea * 1] = (round((a1_1 * src[srcIdx1_1 + 1] + a1_2 * src[srcIdx1_2 + 1] +a2_1 * src[srcIdx2_1 + 1] +a2_2 * src[srcIdx2_2 + 1])) / 255.0f - d_mean[1]) / d_std[1];tar[tarIdx + tarArea * 2] = (round((a1_1 * src[srcIdx1_1 + 0] + a1_2 * src[srcIdx1_2 + 0] +a2_1 * src[srcIdx2_1 + 0] +a2_2 * src[srcIdx2_2 + 0])) / 255.0f - d_mean[0]) / d_std[0];}
}

上述 CUDA 核函数主要是实现双线性插值，我们在第二章的时候讲过，只是这里额外多了一些预处理操作而已，例如除以 255，减均值除标准差等，大家感兴趣的可以看看：二. CUDA编程入门-双线性插值计算

OK，预处理之后我们会调用 enqueue_bindings 函数用作推理：

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

在该函数中就是调用 enqueueV2 执行推理并统计了推理所需时间

推理完成之后我们会执行后处理操作，值得注意的是由于 classification 任务的后处理比较简单因此是放在 CPU 上完成的，所以我们只要看 CPU 上的后处理即可，而像 Detection、Segmentation 等后处理相比而言复杂一点，我们可以放在 GPU 上加速处理

postprocess_cpu 函数内容如下：

bool Classifier::postprocess_cpu() {/*Postprocess -- 测速*/m_timer->start_cpu();/*Postprocess -- 将device上的数据移动到host上*/int output_size    = m_params->num_cls * sizeof(float);CUDA_CHECK(cudaMemcpyAsync(m_outputMemory[0], m_outputMemory[1], output_size, cudaMemcpyKind::cudaMemcpyDeviceToHost, m_stream));CUDA_CHECK(cudaStreamSynchronize(m_stream));/*Postprocess -- 寻找label*/ImageNetLabels labels;int pos = max_element(m_outputMemory[0], m_outputMemory[0] + m_params->num_cls) - m_outputMemory[0];float confidence = m_outputMemory[0][pos] * 100;m_timer->stop_cpu();m_timer->duration_cpu<timer::Timer::ms>("postprocess(CPU)");LOG("Result:     %s", labels.imagenet_labelstring(pos).c_str());   LOG("Confidence  %.3f%%\n", confidence);   return true;
}

后处理的过程比较简单，我们前面也分析过，主要是将推理后的 device 上的数据拷贝到 host 上，然后寻找最大概率的标签并计算置信度，最后将其打印出来

如果大家想看每张图片推理时的前处理、后处理的时间可以在 main.cpp 中将日志级别设置为 VERB：

int main(int argc, char const *argv[])
{...auto level         = logger::Level::VERB;...
}

再次执行后就可以在终端看到前后处理的时间了，如下图所示：

在这里插入图片描述

OK，以上就是整个推理过程的分析了，相对来说还是比较简单的

结语

本次课程我们优化了 6.1 分类器的代码，主要是从代码可复用性、可读性、安全性、可扩展性、可调试性五个部分出发，在代码中我们通过创建的 worker 实例来做推理，在创建时就已经完成了一系列初始化和模型构造工作，接着只需要调用 inference 接口推理即可

另外博主想说的是如果大家对 tensorRT_Pro 或者 CUDA-BEVFusion 这两个 repo 熟悉的话，会发现这里说的推理框架的设计思想都在这两个 repo 中有所体现，例如 RAII+接口模式、内存复用、多线程、智能指针、日志等等

OK，以上就是 6.2 小节案例的全部内容了，下节我们来学习 6.3 小节 INT8 量化，敬请期待😄