五. TensorRT API的基本使用-load-model

前言

自动驾驶之心推出的 《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