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

前言

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

本次课程我们来学习课程第五章—TensorRT API 的基本使用，一起来学习手动实现一个 infer

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

0. 简述

本小节目标：手动实现 infer 完成模型的推理

今天我们来讲第五章节第三小节—5.3-load-model 这个案例，上个小节我们自己手动实现了一个 build，这个小节我们主要是自己手动实现一个 infer 完成模型的推理过程

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

1. 案例运行

在正式开始课程之前，博主先带大家跑通 5.3-infer-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.3 小节案例代码

开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter5-tensorrt-api-basics/5.3-infer-model 小节中创建一个 models 文件夹，接着在 models 文件夹下创建 onnx 和 engine 文件夹，总共三个文件夹需要创建

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

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

cd tensorrt_starter/chapter5-tensorrt-api-basics/5.3-infer-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

输出如下：

我们这里读取一个模型，然后将输入数据打印出来，通过模型的推理拿到对应的输出，将输出数据也打印出来，我们的模型是使用的上节课的 onnx，就是一个 Linear 层

我们在执行下 python 的推理，对比二者是否相同：

可以看到 C++ 和 Python 的推理结果都是一致的，说明我们的 infer 是没有问题的

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

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("fail in building model");
        return 0;
    }
    if(!model.infer()){
        LOGE("fail in infering model");
        return 0;
    }
    return 0;
}

与上节 build 案例相比 main 函数中多了一个 infer 的接口，这个也是仿照着官方 MNIST 案例去写的

2.2 model.cpp

我们来看 infer 接口中具体做了些什么操作，其实 infer 需要做的事情主要包括以下几个部分：

• 1. 读取 model，创建 runtime，engine，context
• 2. 把数据进行 host-> device 传输
• 3. 使用 context 推理
• 4. 把数据进行 device->host 传输

首先我们来看第一部分：

/* 1. 读取model => 创建runtime, engine, context */
if (!fileExists(mEnginePath)) {
    LOGE("ERROR: %s not found", mEnginePath.c_str());
    return false;
}

/* 反序列化从文件中读取的数据以unsigned char的vector保存*/
vector<unsigned char> modelData;
modelData = loadFile(mEnginePath);

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

auto input_dims   = context->getBindingDimensions(0);
auto output_dims  = context->getBindingDimensions(1);

LOG("input dim shape is:  %s", printDims(input_dims).c_str());
LOG("output dim shape is: %s", printDims(output_dims).c_str());

我们先查看文件是否存在，接着通过 loadFile 函数加载文件中的数据保存到 modelData 中，然后实例化一个 logger，再通过 nvinfer1::createInferRuntime 实例化一个 runtime 对象，接着从 runtime 反序列化创建 engine，再从 engine 创建 context，到此为止准备工作就已经完成了

接着我们通过 context 的 getBindingDimensions 这个 API 获取到输入和输出的 dims，接着将其打印出来

Note：在 TensorRT 中，Binding 是指模型输入和输出张量与内存的关联，在实际使用中，我们需要确保为每个输入和输出绑定正确分配内存，并通过索引管理这些绑定。

下一部分开始做数据传递：

/* 2. host->device的数据传递 */
cudaStream_t stream;
cudaStreamCreate(&stream);

/* host memory上的数据*/
float input_host[] = {0.0193, 0.2616, 0.7713, 0.3785, 0.9980, 0.9008, 0.4766, 0.1663, 0.8045, 0.6552};
float output_host[5];

/* device memory上的数据*/
float* input_device = nullptr;
float* weight_device = nullptr;
float* output_device = nullptr;

int input_size = 10;
int output_size = 5;

/* 分配空间, 并传送数据从host到device*/
cudaMalloc(&input_device, sizeof(input_host));
cudaMalloc(&output_device, sizeof(output_host));
cudaMemcpyAsync(input_device, input_host, sizeof(input_host), cudaMemcpyKind::cudaMemcpyHostToDevice, stream);

和之前 MNIST 案例一样，我们先定义一个 stream，接着定义下 input 数据，这里是为了方便测试所以与 python 中的数据保持一致，定义 output 变量，接着我们在 device 上定义一些变量，包括 input_device、weight_device、output_device，然后通过 cudaMalloc 分配空间，通过 cudaMemcpyAsync 将 host 上的 input 数据拷贝到 device 上

接着第三步开始推理：

/* 3. 模型推理, 最后做同步处理 */
float* bindings[] = {input_device, output_device};
bool success = context->enqueueV2((void**)bindings, stream, nullptr);

我们将输入和输出绑定起来，然后送到模型中去执行推理

最后我们做数据传递，将 device 上的推理结果传递到 host 上：

/* 4. device->host的数据传递 */
cudaMemcpyAsync(output_host, output_device, sizeof(output_host), cudaMemcpyKind::cudaMemcpyDeviceToHost, stream);
cudaStreamSynchronize(stream);

LOG("input data is:  %s", printTensor(input_host, input_size).c_str());
LOG("output data is: %s", printTensor(output_host, output_size).c_str());
LOG("finished inference");

数据传递通过 stream 来做一个通过，接着将得到的推理结果打印出来

整个 infer 比较简单，因为没有涉及到预处理和后处理，大家简单了解下就行

2.3 其它

在 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)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0., std=1.)
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.wdight, 1)
                nn.init.constant_(m.bias, 0)
    
    def forward(self, x):
        x = self.linear(x)
        return x

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True

def export_norm_onnx(input, model):
    current_path = os.path.dirname(__file__)
    file = current_path + "/../../models/onnx/sample.onnx"
    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 eval(input, model):
    output = model(input)
    print("from infer------")
    print(input)
    print(output)

if __name__ == "__main__":
    setup_seed(1)

    input = torch.tensor([[0.0193, 0.2616, 0.7713, 0.3785, 0.9980, 0.9008, 0.4766, 0.1663, 0.8045, 0.6552]])
    model = Model()

    export_norm_onnx(input, model)
    eval(input, model)

它就是创建了一个非常简单的 ONNX 模型，其中包含一个 Linear 节点，如下所示：

然后准备了一些输入数据进行验证测试

总结

本次课程我们主要在 5.2 小节的案例上增加了 infer 接口完成了模型的推理，整个过程还是比较简单的，首先读取 model 创建 runtime、engine、context，接着把输入数据从 host 上传输到 device 上，然后使用 context 进行推理，最后把推理好的数据从 device 上传输到 host 上打印显示。这里并没有涉及到复杂的预处理和后处理所以整个过程比较清晰，大家了解下就行

OK，以上就是 5.3 小节案例的全部内容了，下节我们来学习 5.4 小节来打印观察下经过 TensorRT 优化前后模型的结构

下载链接

• tensorrt_starter源码
• 5.3-infer-model案例文件

参考

• Ubuntu20.04软件安装大全
• https://github.com/kalfazed/tensorrt_starter.git