目录
- 前言
- 0. 简述
- 1. 案例运行
- 2. 代码分析
- 2.1 main.cpp
- 2.2 model.hpp
- 2.3 model.cpp
- 2.4 其它
- 总结
- 下载链接
- 参考
前言
自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考
本次课程我们来学习课程第五章—TensorRT API 的基本使用，一起来学习手动实现一个 build
课程大纲可以看下面的思维导图
0. 简述
本小节目标：手动实现 build 完成模型的序列化
今天我们来讲第五章节第二小节—5.2-load-model 这个案例，上个小节我们主要是通过官方 MNIST 案例让大家熟悉 TensorRT 的一些 API 的使用，这个小节我们主要是模仿官方案例来自己手写一个 build
下面我们开始本次课程的学习🤗
1. 案例运行
在正式开始课程之前，博主先带大家跑通 5.2-load-model 这个小节的案例🤗
源代码获取地址：https://github.com/kalfazed/tensorrt_starter.git
首先大家需要把 tensorrt_starter 这个项目给 clone 下来，指令如下：
git clone https://github.com/kalfazed/tensorrt_starter.git
也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2024/7/14 日，若有改动请参考最新）
整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV，大家可以参考 Ubuntu20.04软件安装大全 进行相应软件的安装，博主这里不再赘述
假设你的项目、环境准备完成，下面我们来一起运行 5.2 小节案例代码
开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter5-tensorrt-api-basics/5.2-load-model 小节中创建一个 models 文件夹，接着在 models 文件夹下创建 onnx 和 engine 文件夹，总共三个文件夹需要创建
创建完后 5.2 小节整个目录结构如下：
接着我们需要执行 python 文件创建一个 ONNX 模型，先进入到 5.2 小节中：
cd tensorrt_starter/chapter5-tensorrt-api-basics/5.2-load-model
执行如下指令：
python src/python/generate_onnx.py
Note：大家需要准备一个虚拟环境，安装好 torch、onnx、onnxsim 等第三方库
输出如下：
生成好的 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
输出如下：
接着执行：
./trt-infer
输出如下：
可以看到输出了很多日志信息，该案例主要是通过自定义 build 构建一个 engine 并保存到 models/engine/sample.engine 中，最后打印输入和输出的维度信息
如果大家能够看到上述输出结果，那就说明本小节案例已经跑通，下面我们就来看看具体的代码实现
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/sample.onnx");
if(!model.build()){
LOGE("ERROR: fail in building model");
return 0;
}
return 0;
}
通过传入 ONNX 模型文件路径创建一个 Model 实例，之后调用 build 函数构建 engine
2.2 model.hpp
我们来看 Model 类的定义：
class Model{
public:
Model(std::string onnxPath);
bool build();
private:
std::string mOnnxPath;
std::string mEnginePath;
nvinfer1::Dims mInputDims;
nvinfer1::Dims mOutputDims;
std::shared_ptr<nvinfer1::ICudaEngine> mEngine;
bool constructNetwork();
bool preprocess();
};
Model 类中公有方法主要是 build，而其它的比如 mEngine，constructNetwork，preprocess 等都是私有方法，没有必要暴露给用户的
2.3 model.cpp
接着我们来看 Model 的构造函数：
Model::Model(string onnxPath){
if (!fileExists(onnxPath)) {
LOGE("%s not found. Program terminated", onnxPath.c_str());
exit(1);
}
mOnnxPath = onnxPath;
mEnginePath = getEnginePath(mOnnxPath);
}
首先它会去检查传入的 onnxPath 文件是否存在，如果不存在则打印错误信息并退出，接着把 onnxPath 赋值给私有成员变量 mOnnxPath，通过 getEnginePath 函数拿到对应的 mEnginePath。另外这里的 LOGE 是通过宏定义实现的一个打印函数，它可以用来控制不同的输出日志等级
我们再来看 build 的实现，首先是检查 mEnginePath 是否存在：
if (fileExists(mEnginePath)){
LOG("%s has been generated!", mEnginePath.c_str());
return true;
} else {
LOG("%s not found. Building engine...", mEnginePath.c_str());
}
如果存在则不用再重新 build，如果不存在则需要通过下面的流程进行 build
首先我们实例化一个 logger：
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;
}
};
Logger logger;
我们上节课讲过在创建一个 builder 的时候需要绑定一个 logger，因此我们这里自己手动实现了一个 Logger 类，它继承自 nvinfer1::ILogger，在 Logger 类中我们必须自己手动来实现 log 虚函数。Severity 是一个枚举类，用于控制日志消息的等级，然后将不同的 str 附加不同颜色，如果 severity 级别小于或等于 kINFO，则会通过 cout 将带有前缀的 str 日志信息打印出来
创建完 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);
if (!parser->parseFromFile(mOnnxPath.c_str(), 1)){
LOGE("ERROR: failed to %s", mOnnxPath.c_str());
return false;
}
其实和上节课讲的流程一样，我们先创建一个 builder，然后通过 builder 创建 network、config，接着把 network 和 logger 丢到 nvonnxparser::createParser 函数中创建一个 parser
接着通过 config 设置了最大的 workspace size，其实 config 可以设置非常多的参数，包括 setCalibrationProfile 设置校准文件，setInt8Calibrator 设置校准器等等，这些都是跟模型创建相关的东西，大家自己可以看下
另外这些 API 的说明在官方文档中描述都比较详细，大家也可以参考：tensorrt/developer-guide
config 设置完成之后，通过 parserFromFile 函数将 onnx parser 到 network 里面去
上面这些都是准备工作，接着我们就可以来创建 engine：
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);
通过 builder->buildEngineWithConfig 把 network 和 config 丢进去创建 engine，之后把创建好的 network 做一个序列化保存到 plan 中去，其中 plan 是一个 IHostMemory 的指针，然后我们创建了一个 runtime 方便后续反序列化测试
接着我们把序列化好的 plan 文件通过 fwrite 写入保存到指定路径，方便下次加载使用
下面我们打印了模型的一些基本信息：
mEngine = shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(plan->data(), plan->size()));
mInputDims = network->getInput(0)->getDimensions();
mOutputDims = network->getOutput(0)->getDimensions();
LOG("Input dim is %s", printDims(mInputDims).c_str());
LOG("Output dim is %s", printDims(mOutputDims).c_str());
return true;
通过 runtime->deserializeCudaEngine 来反序列化拿到我们的 engine，其中的 mEngine 是 ICudaEnigne 的指针，是一个推理引擎，然后我们可以通过 network 将输入输出的一些维度信息打印出来
这里有一个小技巧，大家在学习 API 的时候可以通过一些名字大概猜测其主要实现的功能，比如 network 它其中的以 getXXX 为例的 API 一般来说都是去获取网络的一些信息，比如 getLayer、getName 等等，再比如 engine 也有类似于 getXXX 的 API，比如 getDeviceMemorySize、getNbOptimizationProfiles 等等
2.4 其它
在 src/python 文件夹下还有一个 generate_onnx.py 的脚本文件，其内容如下：
import torch
import torch.nn as nn
import torch.onnx
import onnxsim
import onnx
import os
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = nn.Linear(in_features=10, out_features=5, bias=False)
def forward(self, x):
x = self.linear(x)
return x
def setup_seed(seed):
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def export_norm_onnx():
current_path = os.path.dirname(__file__)
file = current_path + "/../../models/onnx/sample.onnx"
input = torch.rand(1, 10)
model = Model()
torch.onnx.export(
model = model,
args = (input,),
f = file,
input_names = ["input0"],
output_names = ["output0"],
opset_version = 15)
print("Finished normal onnx export")
# check the exported onnx model
model_onnx = onnx.load(file)
onnx.checker.check_model(model_onnx)
# use onnx-simplifier to simplify the onnx
print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")
model_onnx, check = onnxsim.simplify(model_onnx)
assert check, "assert check failed"
onnx.save(model_onnx, file)
def infer():
setup_seed(1)
model = Model()
input = torch.tensor([[0.0193, 0.2616, 0.7713, 0.3785, 0.9980, 0.9008, 0.4766, 0.1663, 0.8045, 0.6552]])
output = model(input)
print(input)
print(output)
if __name__ == "__main__":
export_norm_onnx()
infer()
它就是创建了一个非常简单的 ONNX 模型，其中包含一个 Linear 节点，如下所示：
总结
本次课程我们主要模仿官方案例自己手动实现了一个 builder，和官方流程类似，先创建一个 logger，然后创建 builder，通过 builder 创建 network、config，然后创建 parser，通过 parseFromFile 将 ONNX parser 到 network 中，接着创建完 engine，通过 buildSerializedNetwork 进行序列化生成 plan，并将 plan 保存下来，最后调用一个 API 来打印一些输入输出维度信息。总的来说，实现还是比较简单的，关于一些 API 的使用大家可以多尝试尝试
OK，以上就是 5.2 小节案例的全部内容了，下节我们来学习 5.3 小节自己构建一个 infer 来推理模型，敬请期待😄
下载链接
- tensorrt_starter源码
- 5.2-load-model案例文件
参考
- Ubuntu20.04软件安装大全
- https://github.com/kalfazed/tensorrt_starter.git
- https://docs.nvidia.com/deeplearning/tensorrt/developer-guide
