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

  • 开发
  • 23

目录

前言

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

本次课程我们来学习课程第五章—TensorRT API 的基本使用，一起来学习官方 MNIST 案例

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

在这里插入图片描述

0. 简述

本小节目标：理解官方 MNIST 案例

这节课程开始我们进入第五章节—TensorRT API 的基本使用，这个章节偏实战，为大家准备了几个案例：

  • 5.1-mnist-sample
  • 5.2-load-model
  • 5.3-infer-model
  • 5.4-print-structure
  • 5.5-build-model
  • 5.6-build-sub-graph
  • 5.7-custom-basic-trt-plugin
  • 5.8-plugin-unit-test

第五章节准备的案例一共是八个，第一个 mnist-sample，这个是从 TensorRT 官方拿过来的一个案例，这里面韩君老师做了很多注释，大家可以多看看；5.2 小节主要教大家如何利用 TensorRT API 去 load model，如何读取一个 ONNX 模型并将其序列化；5.3 小节主要是推理部分，自己简单写一个推理框架，来完成模型的推理；5.4 主要是教大家如何利用 TensorRT API 将模型优化前后的结构打印出来；5.5 小节后就不再是利用 ONNX 去构建 engine，而是教大家利用 TensorRT C++ API 从零搭建一个 model 并序列化成 engine；5.6 下节同样也是利用 TensorRT C++ API 去搭建一个个小模块，然后组成一个网络；5.7 小节跟大家简单介绍下自己创建一个 plugin 的话应该怎么做；5.8 小节主要是利用 Python API 去快速搭建一个 plugin 的单元测试

今天这个小节我们主要来看第五章节第一小节—5.1-minist-sample 这个案例

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

1. 案例运行

在正式开始课程之前，博主先带大家跑通 5.1-minist-sample 这个小节的案例🤗

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

这个小节的案例其实来自于 TensorRT 官方，只不过韩君老师添加了一些注释，因此我们直接来看官方案例的运行

在开始之前需要你安装好 TensorRT，具体安装流程可以参考：Ubuntu20.04软件安装大全

我们找到安装好的 tensorRT 目录，博主安装的位置是 /opt/TensorRT-8.4.1.5，在这个文件夹下有一个 samples/sampleMNIST 文件夹，本小节的案例就是来自于这里

首先进入该文件夹，执行如下指令：

cd /opt/TensorRT-8.4.1.5/samples/sampleMNIST

执行编译指令：

make -j64

输出结果如下图所示：

在这里插入图片描述

执行运行指令：

./../../bin/sample_mnist

输出结果如下图所示：

在这里插入图片描述

可以看到输入是手写数字 3，模型预测的数字是 3，推理结果正确

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

2. 代码分析

2.1 main函数

我们先从 main 函数的实现看起：

int main(int argc, char** argv)
{
    samplesCommon::Args args;
    bool argsOK = samplesCommon::parseArgs(args, argc, argv);
    if (!argsOK)
    {
        sample::gLogError << "Invalid arguments" << std::endl;
        printHelpInfo();
        return EXIT_FAILURE;
    }
    if (args.help)
    {
        printHelpInfo();
        return EXIT_SUCCESS;
    }

    // 创建一个logger用来保存日志。
    // 这里需要注意一点，日志一般都是继承nvinfer1::ILogger来实现一个自定义的。
    // 由于ILogger有一些函数都是虚函数，所以我们需要自己设计
    auto sampleTest = sample::gLogger.defineTest(gSampleName, argc, argv);
    sample::gLogger.reportTestStart(sampleTest);

    // 创建sample对象，只暴露build和infer接口
    SampleOnnxMNIST sample(initializeSampleParams(args));

    sample::gLogInfo << "Building and running a GPU inference engine for Onnx MNIST" << std::endl;

    // 创建推理引擎
    if (!sample.build())
    {
        return sample::gLogger.reportFail(sampleTest);
    }

    // 推理
    if (!sample.infer())
    {
        return sample::gLogger.reportFail(sampleTest);
    }

    return sample::gLogger.reportPass(sampleTest);
}

前面都是一些参数的解析我们不用管，这里 gLogger 比较重要，TensorRT 在创建 engine 的时候需要绑定一个 logger 日志，我们需要自己实现一个 logger。不过值得注意的是 logger 的实现我们一般是通过继承 nvInfer1::ILogger 来实现一个自定义的，由于 ILogger 中有一些函数是虚函数，因此我们继承之后需要自己来设计实现这些函数，比如 void log() 函数

接着我们创建 sample 对象只暴露 build 和 infer 接口，通过 build 接口创建推理引擎，通过 infer 接口实现推理。那大家可能会好奇，为什么 sample 对象会只暴露这两个接口呢？它具体的实现为什么没有暴露呢？

这其实跟它采用的设计模式有关，这里采用的是接口模式，是一个封装模式，实现类与接口类分离的模式，我们只向用户暴露调用的接口，而具体的实现细节隐藏起来，让使用者只考虑具体的接口，而不必关心具体的实现，这样可以为代码提供更大的灵活性和可维护性。

我们在杜老师的课程 7.4.tensorRT高级(2)-使用RAII接口模式对代码进行有效封装 和韩君老师后续的课程 八. 实战：CUDA-BEVFusion部署分析-学习CUDA-BEVFusion推理框架设计模式 都有简单提到这种设计模式，大家感兴趣的可以看看

2.2 build接口

我们来看 build 接口，代码如下所示：

bool SampleOnnxMNIST::build()
{
    // 创建builder的时候需要传入一个logger来记录日志
    auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
    if (!builder)
    {
        return false;
    }

    // 在创建network的时候需要指定是implicit batch还是explicit batch
    // - implicit batch: network不明确的指定batch维度的大小, 值为0
    // - explicit batch: network明确指定batch维度的大小, 值为1
    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
    if (!network)
    {
        return false;
    }

    // IBuilderConfig是推理引擎相关的设置，比如fp16, int8, workspace size，dla这些都是在config里设置
    auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    if (!config)
    {
        return false;
    }

    // network的创建可以通过parser从onnx导出为network。注意不同layer在不同平台所对应的是不同的
    // 建议这里大家熟悉一下trt中的ILayer都有哪些。后面会用到
    auto parser
        = SampleUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
    if (!parser)
    {
        return false;
    }

    // 为网络设置config, 以及parse
    auto constructed = constructNetwork(builder, network, config, parser);
    if (!constructed)
    {
        return false;
    }

    // 指定profile的cuda stream(平时用的不多)
    auto profileStream = samplesCommon::makeCudaStream();
    if (!profileStream)
    {
        return false;
    }
    config->setProfileStream(*profileStream);

    // 通过builder来创建engine的过程，并将创建好的引擎序列化
    // 平时大家写的时候这里序列化一个引擎后会一般保存到文件里面，这个案例没有写出直接给放到一片memory中后面使用
    SampleUniquePtr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
    if (!plan)
    {
        return false;
    }

    // 其实从这里以后，一般都是infer的部分。大家在创建推理引擎的时候其实写到序列化后保存文件就好了
    // 创建一个runtime来负责推理
    SampleUniquePtr<IRuntime> runtime{createInferRuntime(sample::gLogger.getTRTLogger())};
    if (!runtime)
    {
        return false;
    }

    // 通过runtime来把序列化后的引擎给反序列化, 当作engine来使用
    mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
        runtime->deserializeCudaEngine(plan->data(), plan->size()), samplesCommon::InferDeleter());
    if (!mEngine)
    {
        return false;
    }

    ASSERT(network->getNbInputs() == 1);
    mInputDims = network->getInput(0)->getDimensions();
    ASSERT(mInputDims.nbDims == 4);

    ASSERT(network->getNbOutputs() == 1);
    mOutputDims = network->getOutput(0)->getDimensions();
    ASSERT(mOutputDims.nbDims == 2);

    return true;
}

创建网络即 build 的流程基本上如下所示，我们可以对照着一步步来看：

  • 1. 创建一个 builder
  • 2. 通过 builder 创建一个 network
  • 3. 通过 builder 创建一个 config
  • 4. 通过 config 创建一个 opt（当前案例未体现）
  • 5. 对 network 进行创建
    • 可以使用 parser 直接将 onnx 中各个 layer 转换为 trt 能够识别的 layer（当前案例使用的方式）
    • 也可以通过 trt 提供的 ILayer 相关的 API 自己从零搭建 network（后面会讲）
  • 6. 序列化引擎（当前案例未体现）
  • 7. Free resource（如果使用的是智能指针的话，可以省去这一步）

首先通过 nvinfer1::createInferBuilder 接口创建一个 builder：

auto builder = SampleUniquePtr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
if (!builder)
{
    return false;
}

前面的 SampleUniquePtr<nvinfer1::IBuilder> 大家可能不是很熟悉，它其实就是一个智能指针，大家可以想我们在创建一个对象后在程序结束时往往需要手动去进行 destroy 或者 free 释放资源，因此我们可以利用智能指针在程序结束时自动去释放资源

接着通过 builder 创建一个 network：

const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
auto network = SampleUniquePtr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
if (!network)
{
    return false;
}

在创建 network 的时候需要指定是 implicit batch 还是 explicit batch：

  • implicit batch：隐式 batch，network 不明确的指定 batch 维度的大小，值为 0
  • explicit batch：显式 batch，network 明确指定 batch 维度的大小，值为 1

创建完 network 之后接着创建 config：

auto config = SampleUniquePtr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
if (!config)
{
    return false;
}

nvinfer1::IBuilderConfig 是推理引擎相关的设置，比如 fp16，int8，workspace size，dla 这些都是在 config 里面设置的

接着创建完 config 之后我们通过 parser 从 onnx 导出为 network：

auto parser
    = SampleUniquePtr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
if (!parser)
{
    return false;
}

上述步骤做完之后我们就可以创建网络：

bool SampleOnnxMNIST::constructNetwork(
    SampleUniquePtr<nvinfer1::IBuilder>& builder,
    SampleUniquePtr<nvinfer1::INetworkDefinition>& network, 
    SampleUniquePtr<nvinfer1::IBuilderConfig>& config,
    SampleUniquePtr<nvonnxparser::IParser>& parser)
{
    auto parsed = parser->parseFromFile(locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(),
        static_cast<int>(sample::gLogger.getReportableSeverity()));
    if (!parsed)
    {
        return false;
    }

    if (mParams.fp16)
    {
        config->setFlag(BuilderFlag::kFP16);
    }
    if (mParams.int8)
    {
        config->setFlag(BuilderFlag::kINT8);
        samplesCommon::setAllDynamicRanges(network.get(), 127.0f, 127.0f);
    }

    samplesCommon::enableDLA(builder.get(), config.get(), mParams.dlaCore);

    return true;
}

auto constructed = constructNetwork(builder, network, config, parser);
if (!constructed)
{
    return false;
}

通过 parser->parseFromFile 接口从 ONNX 中来创建网络，如果我们不是使用的 parser 的话，需要自己通过 C++ API 一层层的搭建，这个我们后面会讲

接着我们 skip 掉 profile 的 cuda stream 的指定部分，这个我们平时用得不多

然后就到 engine 序列化部分：

SampleUniquePtr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
if (!plan)
{
    return false;
}

通过 builder->buildSerializedNetwork 接口把 network 和 config 丢进去序列化得到 engine，这个 engine 的类型是 IHostMemory 其实就是 Host 上的一块内存所存储的数据。值得注意的是我们在写 builder 的时候序列化一个引擎后一般会保存到文件里面，也就是会把这里的 plan 保存下来，后续直接加载推理即可，这个案例没有实现而是将其直接放到一片 memory 中后面使用

接着我们创建一个 runtime 对象来负责推理：

SampleUniquePtr<IRuntime> runtime{createInferRuntime(sample::gLogger.getTRTLogger())};
if (!runtime)
{
    return false;
}

通过 runtime 来把序列化后的引擎给反序列化，当作 engine 来使用：

mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
    runtime->deserializeCudaEngine(plan->data(), plan->size()), samplesCommon::InferDeleter());
if (!mEngine)
{
    return false;
}

以上就是 builder 的整个流程，大家可以对照着一步步来看，整体还是比较清晰的

2.3 infer接口

下面我们来看 infer 部分，首先我们需要创建一个 buffers：

samplesCommon::BufferManager buffers(mEngine);

这个 buffers 非常重要，它其实不属于 nvinfer1 里面的内容，只是 TensorRT 官方自己实现的。这个 BufferManager 类的对象 buffers 在创建的初期就已经帮我们把 engine 推理所需要的 host/device memory 已经分配好了。原本我们是需要自己计算 engine 的 input/output 的维度和大小，然后根据这些维度和大小进行 malloc 或者 cudaMalloc 内容分配的

具体的 BufferManager 类的实现如下：

BufferManager(
    std::shared_ptr<nvinfer1::ICudaEngine> engine, std::vector<int64_t> const& volumes, int32_t batchSize = 0)
    : mEngine(engine)
    , mBatchSize(batchSize)
{
    // Create host and device buffers
    for (int32_t i = 0; i < mEngine->getNbIOTensors(); i++)
    {
        auto const name = engine->getIOTensorName(i);
        mNames[name] = i;

        nvinfer1::DataType type = mEngine->getTensorDataType(name);

        std::unique_ptr<ManagedBuffer> manBuf{new ManagedBuffer()};
        manBuf->deviceBuffer = DeviceBuffer(volumes[i], type);
        manBuf->hostBuffer = HostBuffer(volumes[i], type);
        void* deviceBuffer = manBuf->deviceBuffer.data();
        mDeviceBindings.emplace_back(deviceBuffer);
        mManagedBuffers.emplace_back(std::move(manBuf));
    }
}

它其实就是通过一个 for 循环去变量 engine 中的输入输出，然后根据输入输出不同的维度去分配 device memory 和 host memory 最后 emplace_back 到 mManagedBuffers 中，具体细节大家可以参考：samples/common/buffers.h#L246

之后我们从 engine 中创建一个 context：

auto context = SampleUniquePtr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
{
    return false;
}

context 用来做推理，接着我们调用 processInput 函数做前处理：

bool SampleOnnxMNIST::processInput(const samplesCommon::BufferManager& buffers)
{
    const int inputH = mInputDims.d[2];
    const int inputW = mInputDims.d[3];

    // Read a random digit file
    srand(unsigned(time(nullptr)));
    std::vector<uint8_t> fileData(inputH * inputW);
    mNumber = rand() % 10;
    readPGMFile(locateFile(std::to_string(mNumber) + ".pgm", mParams.dataDirs), fileData.data(), inputH, inputW);

    // Print an ascii representation
    sample::gLogInfo << "Input:" << std::endl;
    for (int i = 0; i < inputH * inputW; i++)
    {
        sample::gLogInfo << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % inputW) ? "" : "\n");
    }
    sample::gLogInfo << std::endl;

    float* hostDataBuffer = static_cast<float*>(buffers.getHostBuffer(mParams.inputTensorNames[0]));
    for (int i = 0; i < inputH * inputW; i++)
    {
        hostDataBuffer[i] = 1.0 - float(fileData[i] / 255.0);
    }

    return true;
}

ASSERT(mParams.inputTensorNames.size() == 1);
if (!processInput(buffers))
{
    return false;
}

MNIST 前处理的实现主要包括：

  • 读取随机 digit 数据
  • 分配 buffers 中的 host 上的空间
  • 将数据转为浮点数

接着将预处理后的数据拷贝到 device 上，并通过 executeV2 执行推理：

buffers.copyInputToDevice();
bool status = context->executeV2(buffers.getDeviceBindings().data());
if (!status)
{
    return false;
}

其实 context 创建好之后，推理只需要使用 executeV2 或者 enqueueV2 就可以了，因为 context 与 engine 是绑定的，所以之后 trt 会自动根据创建好的 engine 来逐层进行 forward。

值得注意的是：

  • enqueue，enqueueV2 是异步推理，V2 代表 explicit batch
  • execute，executeV2 是同步推理，V2 代表 explicit batch

现在我们一般使用的都是 enqueueV2

之后将 device 上 forward 的数据拷贝到 host 上，然后再通过 verifyOutput 函数做后处理：

bool SampleOnnxMNIST::verifyOutput(const samplesCommon::BufferManager& buffers)
{
    const int outputSize = mOutputDims.d[1];
    float* output = static_cast<float*>(buffers.getHostBuffer(mParams.outputTensorNames[0]));
    float val{0.0f};
    int idx{0};

    // Calculate Softmax
    float sum{0.0f};
    for (int i = 0; i < outputSize; i++)
    {
        output[i] = exp(output[i]);
        sum += output[i];
    }

    sample::gLogInfo << "Output:" << std::endl;
    for (int i = 0; i < outputSize; i++)
    {
        output[i] /= sum;
        val = std::max(val, output[i]);
        if (val == output[i])
        {
            idx = i;
        }

        sample::gLogInfo << " Prob " << i << "  " << std::fixed << std::setw(5) << std::setprecision(4) << output[i]
                         << " "
                         << "Class " << i << ": " << std::string(int(std::floor(output[i] * 10 + 0.5f)), '*')
                         << std::endl;
    }
    sample::gLogInfo << std::endl;

    return idx == mNumber && val > 0.9f;
}

buffers.copyOutputToHost();
if (!verifyOutput(buffers))
{
    return false;
}

MNIST 后处理的实现：

  • 分配输出所需要的空间
  • 手动实现一个 cpu 版本的 softmax
  • 输出最大值以及所对应的 digit class

以上就是 infer 的整个流程，先预处理接着推理最后后处理，还是比较简单的

2.4 其他

这个 sample 提供的 infer 实现非常简单，主要在于 BufferManager 的实现，这个 BufferManager 是基于 RAII（Resource Acquisition Is Initialization）的设计思想建立的，方便我们在管理 CPU 和 GPU 上的 buffer 的使用，让整个代码变得简洁且可读性高，大家想了解 RAII 的同样可以看下杜老师的课程 7.4.tensorRT高级(2)-使用RAII接口模式对代码进行有效封装 和韩君老师后续的课程 八. 实战：CUDA-BEVFusion部署分析-学习CUDA-BEVFusion推理框架设计模式

推理实现的部分把反序列化的部分给省去了，直接从 context 开始的，context 其实就是上下文，用来创建一些空间来存储一些中间值，可以通过 engine 来创建。一个 engine 可以创建多个 context，用来负责多个不同的推理任务，另外 context 也可以复用，每次新的推理可以利用之前创建好的 context

总结

本次课程我们对官方的 MNIST 案例进行了学习，主要学习了 build 和 infer 两个接口，其中 build 的接口部分包括创建 builder、network、config 等东西，大家按照流程来就行，有点类似于搭积木，而 infer 部分主要是通过前面序列化好的 engine 反序列化，然后通过 engine 创建 context 去执行推理，值得注意的是预处理和后处理部分需要我们来实现。另外我们还简单了解了下 RAII 和接口模式，大家感兴趣的话可以看看相关链接

OK，以上就是 5.1 小节案例的全部内容了，下节我们来学习 5.2 小节利用 TensorRT C++ API 来创建模型，敬请期待😄

参考

阅读全部

相关推荐

  2. intellij idea中build project与build module以及rebuild module的区别与联系

    2024-07-14 11:02:02       28 阅读

  3. v-model的基本使用，v-model原理；v-model绑定；v-model的值绑定；v-model修饰符

    2024-07-14 11:02:02       40 阅读

  4. text-generation-inference使用

    2024-07-14 11:02:02       48 阅读

  5. infer。。。。

    2024-07-14 11:02:02       56 阅读

  7. 【自学记录5】【Pytorch2.0深度学习从零开始学 王晓华】第五章 基于Pytorch卷积层的MNIST分类实战

    2024-07-14 11:02:02       32 阅读

  8. SimLM: Can Language Models Infer Parameters of Physical Systems?

    2024-07-14 11:02:02       51 阅读

  10. 002-基于Pytorch的Fashion-MNIST图像分类

    2024-07-14 11:02:02       38 阅读

最近更新

  3. docker php8.1+nginx base 镜像 dockerfile 配置

    2024-07-14 11:02:02       66 阅读

  4. Could not load dynamic library ‘cudart64_100.dll‘

    2024-07-14 11:02:02       70 阅读

  9. 在Django里面运行非项目文件

    2024-07-14 11:02:02       57 阅读

  19. Python语言-面向对象

    2024-07-14 11:02:02       68 阅读

  20. 如何分清楚常见的 Git 分支管理策略Git Flow、GitHub Flow 和 GitLab Flow

    2024-07-14 11:02:02       62 阅读

热门阅读

  3. 解决npm install 安装报错记录贴

    2024-07-14 11:02:02       27 阅读

  9. 山洪灾害研究

    2024-07-14 11:02:02       25 阅读

  10. 小白C语言基础详解：函数

    2024-07-14 11:02:02       23 阅读

  11. 嵌入式C语言多线程编程最全面试题及编程实例（10万字长文）

    2024-07-14 11:02:02       31 阅读

  15. 【2024最新】C++扫描线算法介绍+实战例题

    2024-07-14 11:02:02       23 阅读

  16. 基于MacOS系统Sonoma 14.5的SSH服务禁止密码登录

    2024-07-14 11:02:02       23 阅读

  17. 【Druid 未授权访问漏洞】解决办法

    2024-07-14 11:02:02       23 阅读

  19. 电子版pdf格式标书怎么加盖公章？

    2024-07-14 11:02:02       28 阅读

  27. not enough information C#

    2024-07-14 11:02:02       21 阅读

  29. python 网络爬虫

    2024-07-14 11:02:02       24 阅读