前言
自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考
本次课程我们来学习课程第六章—部署分类器，一起来学习部署一个简单的分类器
课程大纲可以看下面的思维导图
0. 简述
本小节目标：学习部署一个简单的分类器
这个小节是一个初步分类器的实现
下面我们开始本次课程的学习🤗
1. 案例运行
在正式开始课程之前，博主先带大家跑通 6.1-deploy-classification 这个小节的案例🤗
首先大家需要把 tensorrt_starter 这个项目给 clone 下来，指令如下：
git clone https://github.com/kalfazed/tensorrt_starter.git
也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2024/7/14 日，若有改动请参考最新）
整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV，大家可以参考 Ubuntu20.04软件安装大全 进行相应软件的安装，博主这里不再赘述
假设你的项目、环境准备完成，下面我们一起来运行下 6.1-deploy-classification 小节案例代码
开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter6-deploy-classification-and-inference-design/6.1-deploy-classification 小节中创建一个 models 文件夹，接着在 models 文件夹下创建一个 onnx 和 engine 文件夹，总共三个文件夹需要创建
创建完后 6.1 小节整个目录结构如下：
接着我们需要执行 python 文件创建一个 ONNX 模型，先进入到 6.1 小节中：
cd tensorrt_starter/chapter6-deploy-classification-and-inference-design/6.1-deploy-classification
执行如下指令：
python src/python/export_pretained.py -d ./models/onnx/
Note：大家需要准备一个虚拟环境，安装好 torch、onnx、onnxsim 等第三方库
输出如下：
生成好的 reset50.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. 代码分析
2.1 main.cpp
我们先从 main.cpp 看起：
#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;
}
首先 Model 构造函数需要传入 ONNX 模型路径以及指定要生成的 engine 的精度，接着通过 build 接口构建 engine，并调用 infer 接口推理对每一张图片进行分类预测
2.2 model.cpp
我们先看 build 接口：
bool Model::build(){
if (fileExists(mEnginePath)){
LOG("%s has been generated!", mEnginePath.c_str());
return true;
} else {
LOG("%s not found. Building engine...", mEnginePath.c_str());
}
Logger logger;
auto builder = make_unique<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
auto network = make_unique<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));
auto config = make_unique<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
auto parser = make_unique<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));
config->setMaxWorkspaceSize(1<<28);
config->setProfilingVerbosity(nvinfer1::ProfilingVerbosity::kDETAILED);
if (!parser->parseFromFile(mOnnxPath.c_str(), 1)){
LOGE("ERROR: failed to %s", mOnnxPath.c_str());
return false;
}
if (builder->platformHasFastFp16() && mPrecision == nvinfer1::DataType::kHALF) {
config->setFlag(nvinfer1::BuilderFlag::kFP16);
config->setFlag(nvinfer1::BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
}
auto engine = make_unique<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config));
auto plan = builder->buildSerializedNetwork(*network, *config);
auto runtime = make_unique<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));
auto f = fopen(mEnginePath.c_str(), "wb");
fwrite(plan->data(), 1, plan->size(), f);
fclose(f);
// 如果想要观察模型优化前后的架构变化，可以取消注释
mEngine = shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(plan->data(), plan->size()));
mInputDims = network->getInput(0)->getDimensions();
mOutputDims = network->getOutput(0)->getDimensions();
int inputCount = network->getNbInputs();
int outputCount = network->getNbOutputs();
string layer_info;
LOGV("Before TensorRT optimization");
print_network(*network, false);
LOGV("");
LOGV("After TensorRT optimization");
print_network(*network, true);
LOGV("Finished building engine");
return true;
};
那这个 build 其实和前面的没有什么不同，不同的点在于 FP16 精度的指定：
if (builder->platformHasFastFp16() && mPrecision == nvinfer1::DataType::kHALF) {
config->setFlag(nvinfer1::BuilderFlag::kFP16);
config->setFlag(nvinfer1::BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
}
首先我们会做一个判断，检查当前的硬件平台是否支持 FP16 运算并检查我们设置的模型精度是否是 FP16，如果两个条件都满足，我们通过 config 的 setFalg 方法告诉 TensorRT 构建 engine 时启用 FP16 精度，并尽可能的优先保持指定的 FP16 精度要求，更多细节大家可以参考 https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/namespacenvinfer1
下面我们重点来看下 infer 接口：
bool Model::infer(string imagePath){
if (!fileExists(mEnginePath)) {
LOGE("ERROR: %s not found", mEnginePath.c_str());
return false;
}
vector<unsigned char> modelData;
modelData = loadFile(mEnginePath);
Timer timer;
Logger logger;
auto runtime = make_unique<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));
auto engine = make_unique<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(modelData.data(), modelData.size()));
auto context = make_unique<nvinfer1::IExecutionContext>(engine->createExecutionContext());
}
infer 前面的部分和我们之前说的没有什么区别，创建 runtime，通过 runtime 创建 engine，通过 engine 创建 context
我们接着看模型的前处理部分：
auto input_dims = context->getBindingDimensions(0);
auto output_dims = context->getBindingDimensions(1);
cudaStream_t stream;
CUDA_CHECK(cudaStreamCreate(&stream));
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);
/*
为了让trt推理和pytorch的推理结果一致，我们需要对其pytorch所用的mean和std
这里面opencv读取的图片是BGR格式，所以mean和std也按照BGR的顺序存储
可以参考pytorch官方提供的前处理方案: https://pytorch.org/hub/pytorch_vision_resnet/
*/
float mean[] = {0.406, 0.456, 0.485};
float std[] = {0.225, 0.224, 0.229};
/*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 -- 测速*/
timer.start_cpu();
/*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());
}
/*Preprocess -- resize(默认是bilinear interpolation)*/
cv::resize(input_image, input_image, cv::Size(input_width, input_height));
/*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];
}
}
/*Preprocess -- 将host的数据移动到device上*/
CUDA_CHECK(cudaMemcpyAsync(input_device, input_host, input_size, cudaMemcpyKind::cudaMemcpyHostToDevice, stream));
timer.stop_cpu();
timer.duration_cpu<Timer::ms>("preprocess(resize + norm + bgr2rgb + hwc2chw + H2D)");
/*Inference -- 测速*/
timer.start_cpu();
首先我们通过调用 getBindingDimensions 获取推理输入和输出的张量维度，接着通过 input_dims 和 output_dims 获取输入宽高以及输出的类别数并计算输入输出所需要的内容空间大小
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));
接着我们使用 cudaMalloc 函数为输入和输出张量在 device 上分配内存，使用 cudaMallocHost 函数为输入和输出张量在 host 上分配内存，这些内存将用于存储和传递模型推理所需的数据
cv::Mat input_image;
input_image = cv::imread(imagePath);
if (input_image.data == nullptr) {
LOGE("file not founded! Program terminated");
return false;
}
cv::resize(input_image, input_image, cv::Size(input_width, input_height));
然后我们使用 opencv 来读取输入图像并将图像调整为模型所需的输入尺寸
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];
}
}
接着我们使用上节课提到的 CPU 端的预处理代码对输入图像进行预处理操作，主要包括 bgr2rgb、normalization、hwc2chw 以及减均值除以标准差
Note：这里需要大家注意为了让 tensorrt 推理和 pytorch 的推理结果一致，我们需要保证它们做的前处理一模一样，也就是说 pytorch 中做了哪些图像预处理操作在 tensorrt 中也需要做同样的操作，所以大家如果发现部署的模型在 tensorrt 和 pytorch 上精度对不齐的时候，可以考虑从模型的前处理入手，是不是没有做 bgr2rgb 呢，是不是没有除以 255 呢，是不是没有减均值除以标准差呢
pytorch 官方在实现 resnet 时所做的前处理操作大家可以参考：https://pytorch.org/hub/pytorch_vision_resnet/
CUDA_CHECK(cudaMemcpyAsync(input_device, input_host, input_size, cudaMemcpyKind::cudaMemcpyHostToDevice, stream));
最后使用 cudaMemcpyAsync 将处理好的输入数据从主机内存 host 上复制到设备内存 device 上，传输是在指定的 CUAD 流中异步进行的
以上就是整个分类模型的前处理操作，下面我们来看推理部分
float* bindings[] = {input_device, output_device};
if (!context->enqueueV2((void**)bindings, stream, nullptr)){
LOG("Error happens during DNN inference part, program terminated");
return false;
}
首先定义一个 bindings 数组，将输入和输出张量的 device 指针传递给 TensorRT 执行上下文，调用 enqueueV2 函数在执行的 CUDA 流中异步执行推理操作，推理过程将使用预先加载的模型和输入数据在 GPU 上计算输出结果
以上就是整个分类模型的推理操作，比较简单，下面我们来看后处理部分
CUDA_CHECK(cudaMemcpyAsync(output_host, output_device, output_size, cudaMemcpyKind::cudaMemcpyDeviceToHost, stream));
CUDA_CHECK(cudaStreamSynchronize(stream));
首先使用 cudaMemcpyAsync 将推理结果从 device 内存复制回 host 内存，cudaStreamSynchronize 函数确保所有异步操作在这里完成
ImageNetLabels labels;
int pos = max_element(output_host, output_host + num_classes) - output_host;
float confidence = output_host[pos] * 100;
然后我们从推理中寻找分类结果，max_element 函数用于找到推理结果中概率最大的类别索引，并计算对应的置信度，利用 ImageNetLabels 类获取对应的标签名
LOG("Inference result: %s, Confidence is %.3f%%\n", labels.imagenet_labelstring(pos).c_str(), confidence);
CUDA_CHECK(cudaFree(output_device));
CUDA_CHECK(cudaFree(input_device));
CUDA_CHECK(cudaStreamDestroy(stream));
最后释放资源并输出结果
3. 补充说明
这个小节是一个根据第五章节的代码的改编，构建出来的一个 classification 的推理框架，可以实现一系列 preprocess + enqueue + postprocess 的推理。这里大家可以简单的参考下实现的方法，但是建议大家如果自己从零构建推理框架的话不要这么写，主要是因为目前这个框架有非常多的缺陷导致有很多潜在性的问题，主要是因为我们在设计初期并没有考虑太多，这里罗列几点
1. 代码看起来比较乱，主要有以下几个原因：
- 整体上没有使用任何 C++ 设计模式，导致代码复用不好、可读性差、灵活性也比较低，对以后的代码扩展不是很友好
- 比如说如果想让这个框架支持 detection 或者 segmentation 的话需要怎么办？
- 再比如说如果想要让框架做成 multi-stage，multi-task 模型的话需要怎么办？
- 再比如说如果想要这个框架既支持 image 输入，也支持 3D point cloud 输入需要怎么办？
- 封装没有做好，导致有一些没有必要的信息和操作暴露在外面，让代码的可读性差
- 比如说我们是否需要从 main 函数中考虑如果 build/infer 失败应该怎么办？
- 再比如说我们在创建一个 engine 的时候，是否可以只从 main 中提供一系列参数，其余的信息不暴露？
- 内存复用没有做好，导致出现一些额外的开辟销毁内存的开销
- 比如说多张图片依次推理的时候，是否需要每次都 cudaMalloc 或者 cudaMallocHost 以及 cudaStreamCreate 呢？
- 再比如说我们是否可以在 model 的构造函数初始化的时候，就把这个 model 所需要的 CPU 和 GPU 内存分配好以后就不变了呢？
2. 其次仍然还有很多功能没有写全：
- INT8 量化的 calibrator 的实现
- trt plugin 的实现
- CPU 和 GPU overlap 推理的实现
- 多 batch 或者动态 batch 的实现
- CPU 端 threa 级别的并行处理的实现
- …
当然还有很多需要扩展的地方，但是我们可以根据目标出现的一些问题，考虑一些解决方案，具体细节我们下节课程来讲
结语
本次课程我们学习了一个简单的分类器模型的部署，在 build 阶段主要是通过 onnxparser 构建 engine，在 infer 阶段主要包括前处理、推理以及后处理三部分，其中前处理需要和 pytorch 保持一致。另外这个小节的代码存在很多问题，比如代码看起来比较乱，没有封装，没有使用 C++ 设计模式等等，下节课我们就来解决这些问题
OK，以上就是 6.1 小节案例的全部内容了，下节我们来学习 6.2 小节优化分类器代码，敬请期待😄
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
