我们基于tensorrt-cpp-api这个仓库介绍。这个仓库的代码是一个非常不错的tensorrt的cpp api实现,可基于此开发自己的项目。

我们从src/main.cpp开始按顺序说明。

一、首先是声明我们创建tensorrt model的参数。

 // Specify our GPU inference configuration options
    Options options;
    // Specify what precision to use for inference
    // FP16 is approximately twice as fast as FP32.
    options.precision = Precision::FP16;
    // If using INT8 precision, must specify path to directory containing
    // calibration data.
    options.calibrationDataDirectoryPath = "";
    // Specify the batch size to optimize for.
    options.optBatchSize = 1;
    // Specify the maximum batch size we plan on running.
    options.maxBatchSize = 1;
    // Specify the directory where you want the model engine model file saved.
    options.engineFileDir = ".";

这里Options是在src/engine.h中声明的一个结构体


// 网络的选项配置
struct Options {
    // 用于 GPU 推理的精度模式
    // 可选值:FP32(单精度浮点)、FP16(半精度浮点)、INT8(8位量化)
    Precision precision = Precision::FP16;

    // 如果选择 INT8 精度,必须提供校准数据集的目录路径
    std::string calibrationDataDirectoryPath;

    // 用于计算 INT8 校准数据的批量大小
    // 应设置为 GPU 能支持的最大批量大小
    int32_t calibrationBatchSize = 128;

    // 优化的批量大小
    // 表示引擎在构建时会针对该批量大小进行优化
    int32_t optBatchSize = 1;

    // 最大允许的批量大小
    // 表示引擎支持的最大批量大小
    int32_t maxBatchSize = 16;

    // GPU 设备索引
    // 用于指定在哪个 GPU 上运行推理
    int deviceIndex = 0;

    // 引擎文件保存的目录
    // 用于存储 TensorRT 引擎文件
    std::string engineFileDir = ".";

    // 最大允许的输入宽度
    // 默认为 -1,表示期望固定的输入大小
    int32_t maxInputWidth = -1;

    // 最小允许的输入宽度
    // 默认为 -1,表示期望固定的输入大小
    int32_t minInputWidth = -1;

    // 优化的输入宽度
    // 默认为 -1,表示期望固定的输入大小
    int32_t optInputWidth = -1;
};

二、构建tensorrt推理的engine

Engine<float> engine(options);

这里Engine是一个模板类,继承字IEngine类。IEngine 是一个抽象接口类,定义了 TensorRT 推理引擎的核心功能。 它包括构建和加载网络、运行推理、获取输入和输出张量维度等功能。 通过模板参数 T,可以支持不同的数据类型(如 float 或 int)。 具体的实现需要在派生类中完成。

IEngine类的实现在include/interfaces/IEngine.h中。

template <typename T>
class IEngine {
public:
    // 虚析构函数,确保派生类的资源能够正确释放
    virtual ~IEngine() = default;

    // 构建并加载 ONNX 模型到 TensorRT 引擎文件
    // 参数:
    // - onnxModelPath: ONNX 模型文件的路径
    // - subVals: 输入数据的减法均值,用于归一化
    // - divVals: 输入数据的除法均值,用于归一化
    // - normalize: 是否对输入数据进行归一化
    // 返回值:
    // - 如果成功构建并加载网络,返回 true;否则返回 false
    virtual bool buildLoadNetwork(std::string onnxModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
                                  const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) = 0;

    // 从磁盘加载 TensorRT 引擎文件到内存
    // 参数:
    // - trtModelPath: TensorRT 引擎文件的路径
    // - subVals: 输入数据的减法均值,用于归一化
    // - divVals: 输入数据的除法均值,用于归一化
    // - normalize: 是否对输入数据进行归一化
    // 返回值:
    // - 如果成功加载网络,返回 true;否则返回 false
    virtual bool loadNetwork(std::string trtModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
                             const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) = 0;

    // 运行推理
    // 参数:
    // - inputs: 输入数据,格式为 [input][batch][cv::cuda::GpuMat]
    // - featureVectors: 输出数据,格式为 [batch][output][feature_vector]
    // 返回值:
    // - 如果推理成功,返回 true;否则返回 false
    virtual bool runInference(const std::vector<std::vector<cv::cuda::GpuMat>> &inputs, 
                              std::vector<std::vector<std::vector<T>>> &featureVectors) = 0;

    // 获取输入张量的维度
    // 返回值:
    // - 输入张量的维度列表,类型为 nvinfer1::Dims3
    virtual const std::vector<nvinfer1::Dims3> &getInputDims() const = 0;

    // 获取输出张量的维度
    // 返回值:
    // - 输出张量的维度列表,类型为 nvinfer1::Dims
    virtual const std::vector<nvinfer1::Dims> &getOutputDims() const = 0;
};

再来看Engine类的实现。

template <typename T>
class Engine : public IEngine<T> {
public:
    // 构造函数,初始化引擎选项
    Engine(const Options &options);

    // 析构函数,释放 GPU 缓冲区
    ~Engine();

    // 构建并加载 ONNX 模型到 TensorRT 引擎文件
    // 将模型缓存到磁盘以避免重复构建,并加载到内存中
    // 默认情况下,输入数据会被归一化到 [0.f, 1.f]
    // 如果需要 [-1.f, 1.f] 的归一化,可以通过 subVals 和 divVals 参数设置
    bool buildLoadNetwork(std::string onnxModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
                          const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) override;

    // 从磁盘加载 TensorRT 引擎文件到内存
    // 默认情况下,输入数据会被归一化到 [0.f, 1.f]
    // 如果需要 [-1.f, 1.f] 的归一化,可以通过 subVals 和 divVals 参数设置
    bool loadNetwork(std::string trtModelPath, const std::array<float, 3> &subVals = {0.f, 0.f, 0.f},
                     const std::array<float, 3> &divVals = {1.f, 1.f, 1.f}, bool normalize = true) override;

    // 运行推理
    // 输入格式:[input][batch][cv::cuda::GpuMat]
    // 输出格式:[batch][output][feature_vector]
    bool runInference(const std::vector<std::vector<cv::cuda::GpuMat>> &inputs, std::vector<std::vector<std::vector<T>>> &featureVectors) override;

    // 工具函数:调整图像大小并保持宽高比,通过在右侧或底部填充来实现
    // 适用于需要将检测坐标映射回原始图像的场景(例如 YOLO 模型)
    static cv::cuda::GpuMat resizeKeepAspectRatioPadRightBottom(const cv::cuda::GpuMat &input, size_t height, size_t width,
                                                                const cv::Scalar &bgcolor = cv::Scalar(0, 0, 0));

    // 获取输入张量的维度
    [[nodiscard]] const std::vector<nvinfer1::Dims3> &getInputDims() const override { return m_inputDims; };

    // 获取输出张量的维度
    [[nodiscard]] const std::vector<nvinfer1::Dims> &getOutputDims() const override { return m_outputDims; };

    // 工具函数:将三维嵌套的输出数组转换为二维数组
    // 适用于批量大小为 1 且有多个输出特征向量的情况
    static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<std::vector<T>> &output);

    // 工具函数:将三维嵌套的输出数组转换为一维数组
    // 适用于批量大小为 1 且只有一个输出特征向量的情况
    static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<T> &output);

    // 工具函数:将输入从 NHWC 格式转换为 NCHW 格式,并应用归一化和均值减法
    static cv::cuda::GpuMat blobFromGpuMats(const std::vector<cv::cuda::GpuMat> &batchInput, const std::array<float, 3> &subVals,
                                            const std::array<float, 3> &divVals, bool normalize, bool swapRB = false);

private:
    // 构建网络
    bool build(std::string onnxModelPath, const std::array<float, 3> &subVals, const std::array<float, 3> &divVals, bool normalize);

    // 将引擎选项序列化为字符串
    std::string serializeEngineOptions(const Options &options, const std::string &onnxModelPath);

    // 获取设备名称列表
    void getDeviceNames(std::vector<std::string> &deviceNames);

    // 清除 GPU 缓冲区
    void clearGpuBuffers();

    // 输入数据的归一化、缩放和均值减法参数
    std::array<float, 3> m_subVals{};
    std::array<float, 3> m_divVals{};
    bool m_normalize;

    // 存储输入和输出 GPU 缓冲区的指针
    std::vector<void *> m_buffers;
    std::vector<uint32_t> m_outputLengths{};
    std::vector<nvinfer1::Dims3> m_inputDims;
    std::vector<nvinfer1::Dims> m_outputDims;
    std::vector<std::string> m_IOTensorNames;
    int32_t m_inputBatchSize;

    // TensorRT 运行时和推理上下文
    std::unique_ptr<nvinfer1::IRuntime> m_runtime = nullptr;
    std::unique_ptr<Int8EntropyCalibrator2> m_calibrator = nullptr;
    std::unique_ptr<nvinfer1::ICudaEngine> m_engine = nullptr;
    std::unique_ptr<nvinfer1::IExecutionContext> m_context = nullptr;

    // 引擎选项
    const Options m_options;

    // TensorRT 日志器
    Logger m_logger;
};

// 构造函数:初始化引擎选项
template <typename T>
Engine<T>::Engine(const Options &options) : m_options(options) {}

// 析构函数:清除 GPU 缓冲区
template <typename T>
Engine<T>::~Engine() { clearGpuBuffers(); }

相较于抽象接口类IEngine,Engine类添加了resizeKeepAspectRatioPadRightBottom、transformOutput、blobFromGpuMats这几个static的public函数。

    // 工具函数:调整图像大小并保持宽高比,通过在右侧或底部填充来实现
    // 适用于需要将检测坐标映射回原始图像的场景(例如 YOLO 模型)
    static cv::cuda::GpuMat resizeKeepAspectRatioPadRightBottom(const cv::cuda::GpuMat &input, size_t height, size_t width,
                                                                const cv::Scalar &bgcolor = cv::Scalar(0, 0, 0));
    // 工具函数:将三维嵌套的输出数组转换为二维数组
    // 适用于批量大小为 1 且有多个输出特征向量的情况
    static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<std::vector<T>> &output);

    // 工具函数:将三维嵌套的输出数组转换为一维数组
    // 适用于批量大小为 1 且只有一个输出特征向量的情况
    static void transformOutput(std::vector<std::vector<std::vector<T>>> &input, std::vector<T> &output);

    // 工具函数:将输入从 NHWC 格式转换为 NCHW 格式,并应用归一化和均值减法
    static cv::cuda::GpuMat blobFromGpuMats(const std::vector<cv::cuda::GpuMat> &batchInput, const std::array<float, 3> &subVals,
                                            const std::array<float, 3> &divVals, bool normalize, bool swapRB = false);

在类中,static 修饰的成员函数是 静态成员函数,它与普通成员函数有以下区别:

(1)静态成员函数不依赖于类的实例 静态成员函数属于类本身,而不是类的某个实例。 调用静态成员函数时,不需要创建类的对象,可以直接通过类名调用。

(2)静态成员函数不能访问非静态成员变量 静态成员函数没有 this 指针,因此无法访问类的非静态成员变量或非静态成员函数。 静态成员函数只能访问类的静态成员变量或调用其他静态成员函数。

在这里为什么使用 static?

(1)工具函数的设计 transformOutput 是一个工具函数,用于将三维嵌套的输出数组转换为二维数组。 它的功能与类的实例无关,因此设计为静态函数更合理。 这样可以直接通过类名调用,而不需要创建类的对象。

(2)提高代码的可读性和效率 将与类实例无关的函数声明为静态函数,可以明确表示该函数不依赖于类的状态。 静态函数的调用效率通常比非静态函数更高,因为它不需要隐式传递 this 指针。

第二个修改是对于getInputDims和getOutputDims函数的处理。在抽象接口类IEngine中,其实现方式是

virtual const std::vector<nvinfer1::Dims3> &getInputDims() const = 0;

在Engine类中,其声明变成了

    // 获取输入张量的维度
    [[nodiscard]] const std::vector<nvinfer1::Dims3> &getInputDims() const override { return m_inputDims; };

[[nodiscard]] 是 C++17 引入的一个属性,用于提示调用者不要忽略函数的返回值。 如果调用者忽略了带有 [[nodiscard]] 属性的函数的返回值,编译器会发出警告。

在这段代码中,getInputDims() 返回的是输入张量的维度信息。如果调用者忽略了这个返回值,可能会导致程序逻辑错误。 使用 [[nodiscard]] 可以提醒开发者注意返回值的重要性。

例如engine.getInputDims(); // 如果没有使用返回值,编译器会发出警告。

override 是 C++11 引入的关键字,用于显式声明一个函数是从基类继承并重写的虚函数。 如果函数没有正确地重写基类中的虚函数(例如函数签名不匹配),编译器会报错。

它可以防止由于函数签名错误而导致的意外行为。 在这段代码中,getInputDims() 是从基类 IEngine 中继承并重写的虚函数。

第一个 const: 表示返回的引用是常量,调用者不能修改返回的 std::vector。 第二个 const: 表示该成员函数不会修改类的成员变量,保证函数是只读的。

需要注意的是,在src/engine.h中,使用了模版类的实现,这有点要注意:


////////////////////////////// 忽略之前的所有代码

// 构造函数:初始化引擎选项
template <typename T>
Engine<T>::Engine(const Options &options) : m_options(options) {}

// 析构函数:清除 GPU 缓冲区
template <typename T>
Engine<T>::~Engine() { clearGpuBuffers(); }

// Include inline implementations
#include "engine/EngineRunInference.inl"
#include "engine/EngineUtilities.inl"
#include "engine/EngineBuildLoadNetwork.inl"

在文件的最后,才include了这几个头文件。

EngineRunInference.inl: 包含 runInference 函数的实现。

EngineUtilities.inl: 包含工具函数(如 transformOutput)的实现。

EngineBuildLoadNetwork.inl: 包含 buildLoadNetwork 和 loadNetwork 函数的实现。

将 #include 头文件放在文件末尾的原因是为了包含 内联实现文件(.inl 文件),这种设计通常用于模板类或需要将实现与声明分离的情况下。

  1. .inl 文件的作用 .inl 文件通常用于存放模板类或函数的实现。 在 C++ 中,模板类或模板函数的实现必须在编译时可见,因此不能像普通类那样将实现放在 .cpp 文件中。 为了保持代码的清晰和模块化,开发者会将模板类的实现从头文件中分离出来,放入 .inl 文件中。
  2. 为什么将 .inl 文件放在文件末尾? 2.1 避免重复定义 如果在头文件的开头包含 .inl 文件,可能会导致重复定义的问题,尤其是在头文件被多次包含时。 将 .inl 文件放在头文件末尾,可以确保模板类的声明先被处理,然后再处理实现部分。

2.2 确保依赖的声明已完成 .inl 文件中的实现可能依赖于头文件中声明的类或函数。 将 .inl 文件放在头文件末尾,可以确保所有必要的声明都已完成,避免编译错误。

2.3 提高代码的可读性 将 .inl 文件放在头文件末尾,可以将类的声明和实现逻辑分开,增强代码的可读性和可维护性。 开发者可以在头文件中快速查看类的接口,而不需要直接阅读实现细节。 3. 为什么不直接放在 .cpp 文件中? 模板类或模板函数的实现不能放在 .cpp 文件中,因为模板的具体类型是在编译时实例化的,而 .cpp 文件是在链接阶段处理的。 如果将模板实现放在 .cpp 文件中,编译器在实例化模板时将无法找到实现,导致链接错误。

三、tensorrt推理Engine的具体实现

3.1 buildLoadNetwork实现

buildLoadNetwork在include/engine/EngineBuildLoadNetwork.inl中实现。

template <typename T>
bool Engine<T>::buildLoadNetwork(std::string onnxModelPath, const std::array<float, 3> &subVals, const std::array<float, 3> &divVals,
                                 bool normalize) {
    const auto engineName = serializeEngineOptions(m_options, onnxModelPath);
    const auto engineDir = std::filesystem::path(m_options.engineFileDir);
    std::filesystem::path enginePath = engineDir / engineName;
    spdlog::info("Searching for engine file with name: {}", enginePath.string());

    if (Util::doesFileExist(enginePath)) {
        spdlog::info("Engine found, not regenerating...");
    } else {
        if (!Util::doesFileExist(onnxModelPath)) {
            auto msg = "Could not find ONNX model at path: " + onnxModelPath;
            spdlog::error(msg);
            throw std::runtime_error(msg);
        }

        spdlog::info("Engine not found, generating. This could take a while...");
        if (!std::filesystem::exists(engineDir)) {
            std::filesystem::create_directories(engineDir);
            spdlog::info("Created directory: {}", engineDir.string());
        }

        auto ret = build(onnxModelPath, subVals, divVals, normalize);
        if (!ret) {
            return false;
        }
    }

    return loadNetwork(enginePath, subVals, divVals, normalize);
}

其中,最重要的是build函数的实现。build的实现比较长,该函数的作用是从一个 ONNX 模型文件构建 TensorRT 引擎。 它会解析 ONNX 模型文件,设置优化配置(如动态批量大小、动态输入宽度、精度模式等),并最终生成一个序列化的 TensorRT 引擎文件。

其完整实现如下代码,后面会对其中重要代码展开解释,一并写在代码注释上。

3.2 build函数实现(最重要的函数)

template <typename T>
bool Engine<T>::build(std::string onnxModelPath, const std::array<float, 3> &subVals, const std::array<float, 3> &divVals, bool normalize) {
    // Create our engine builder.
    //创建 TensorRT 构建器和网络

    // nvinfer1::createInferBuilder:
    // 创建一个 TensorRT 构建器(IBuilder),用于构建和优化网络。
    // m_logger 是一个自定义的日志器,用于记录 TensorRT 的日志信息。


    auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(m_logger));
    if (!builder) {
        return false;
    }

    // Define an explicit batch size and then create the network (implicit batch
    // size is deprecated). More info here:
    // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#explicit-implicit-batch
    
    // builder->createNetworkV2:
    // 创建一个网络定义(INetworkDefinition),用于描述神经网络的结构。
    // 使用 kEXPLICIT_BATCH 标志表示显式批量大小(TensorRT 7.0 之后推荐使用)。
    auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
    if (!network) {
        return false;
    }

    // Create a parser for reading the onnx file.
    //  创建 ONNX 解析器并解析模型

    // nvonnxparser::createParser:
    // 创建一个 ONNX 解析器(IParser),用于将 ONNX 模型解析为 TensorRT 的网络定义。

    auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, m_logger));
    if (!parser) {
        return false;
    }

    // We are going to first read the onnx file into memory, then pass that buffer
    // to the parser. Had our onnx model file been encrypted, this approach would
    // allow us to first decrypt the buffer.

    // 读取 ONNX 文件:
    // 将 ONNX 模型文件读入内存缓冲区(buffer),然后传递给解析器。
    // 这种方式允许对模型文件进行预处理(如解密)。
    std::ifstream file(onnxModelPath, std::ios::binary | std::ios::ate);
    std::streamsize size = file.tellg();
    file.seekg(0, std::ios::beg);

    std::vector<char> buffer(size);
    if (!file.read(buffer.data(), size)) {
        auto msg = "Error, unable to read engine file";
        spdlog::error(msg);
        throw std::runtime_error(msg);
    }

    // parser->parse:
    // 解析 ONNX 模型文件,将其转换为 TensorRT 的网络定义。
    // ONNX 是一种常见的模型格式,TensorRT 提供了专门的解析器来支持 ONNX 模型。
    // 通过将文件读入内存,可以灵活处理加密或压缩的模型文件。
    // Parse the buffer we read into memory.
    auto parsed = parser->parse(buffer.data(), buffer.size());
    if (!parsed) {
        return false;
    }

    // Ensure that all the inputs have the same batch size
 
    const auto numInputs = network->getNbInputs();
    if (numInputs < 1) {
        auto msg = "Error, model needs at least 1 input!";
        spdlog::error(msg);
        throw std::runtime_error(msg);
    }

    // 检查输入的批量大小
    // TensorRT 要求所有输入的批量大小必须一致。
    const auto input0Batch = network->getInput(0)->getDimensions().d[0];
    for (int32_t i = 1; i < numInputs; ++i) {
        if (network->getInput(i)->getDimensions().d[0] != input0Batch) {
            auto msg = "Error, the model has multiple inputs, each with differing batch sizes!";
            spdlog::error(msg);
            throw std::runtime_error(msg);
        }
    }

    // Check to see if the model supports dynamic batch size or not
    bool doesSupportDynamicBatch = false;
    if (input0Batch == -1) {
        doesSupportDynamicBatch = true;
        spdlog::info("Model supports dynamic batch size");
    } else {
        spdlog::info("Model only supports fixed batch size of {}", input0Batch);
        // If the model supports a fixed batch size, ensure that the maxBatchSize
        // and optBatchSize were set correctly.
        if (m_options.optBatchSize != input0Batch || m_options.maxBatchSize != input0Batch) {
            auto msg = "Error, model only supports a fixed batch size of " + std::to_string(input0Batch) +
                       ". Must set Options.optBatchSize and Options.maxBatchSize to 1";
            spdlog::error(msg);
            throw std::runtime_error(msg);
        }
    }

    const auto input3Batch = network->getInput(0)->getDimensions().d[3];
    bool doesSupportDynamicWidth = false;
    if (input3Batch == -1) {
        doesSupportDynamicWidth = true;
        spdlog::info("Model supports dynamic width. Using Options.maxInputWidth, Options.minInputWidth, and Options.optInputWidth to set the input width.");

        // Check that the values of maxInputWidth, minInputWidth, and optInputWidth are valid
        if (m_options.maxInputWidth < m_options.minInputWidth || m_options.maxInputWidth < m_options.optInputWidth ||
            m_options.minInputWidth > m_options.optInputWidth
            || m_options.maxInputWidth < 1 || m_options.minInputWidth < 1 || m_options.optInputWidth < 1) {
            auto msg = "Error, invalid values for Options.maxInputWidth, Options.minInputWidth, and Options.optInputWidth";
            spdlog::error(msg);
            throw std::runtime_error(msg);
        }
    }

    //设置优化配置

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

    //设置优化配置,包括动态批量大小和动态输入宽度。
    // 使用 IOptimizationProfile 定义输入的最小、优化和最大维度。
    // 为什么要这么写?
    // TensorRT 允许为动态输入设置优化配置,以便在推理时根据实际输入调整性能。
    // 通过设置优化配置,可以在性能和灵活性之间取得平衡。
    // Register a single optimization profile
    nvinfer1::IOptimizationProfile *optProfile = builder->createOptimizationProfile();
    for (int32_t i = 0; i < numInputs; ++i) {
        // Must specify dimensions for all the inputs the model expects.
        const auto input = network->getInput(i);
        const auto inputName = input->getName();
        const auto inputDims = input->getDimensions();
        int32_t inputC = inputDims.d[1];
        int32_t inputH = inputDims.d[2];
        int32_t inputW = inputDims.d[3];
    
        int32_t minInputWidth = std::max(m_options.minInputWidth, inputW);

        // Specify the optimization profile`
        if (doesSupportDynamicBatch) {
            optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, inputC, inputH, minInputWidth));
        } else {
            optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMIN,
                                      nvinfer1::Dims4(m_options.optBatchSize, inputC, inputH, minInputWidth));
        }

        if (doesSupportDynamicWidth) {
            optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kOPT,
                                      nvinfer1::Dims4(m_options.optBatchSize, inputC, inputH, m_options.optInputWidth));
            optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMAX,
                                      nvinfer1::Dims4(m_options.maxBatchSize, inputC, inputH, m_options.maxInputWidth));
        } else {
            optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kOPT,
                                    nvinfer1::Dims4(m_options.optBatchSize, inputC, inputH, inputW));
            optProfile->setDimensions(inputName, nvinfer1::OptProfileSelector::kMAX,
                                    nvinfer1::Dims4(m_options.maxBatchSize, inputC, inputH, inputW));
        }
    }
    config->addOptimizationProfile(optProfile);

    // Set the precision level
    const auto engineName = serializeEngineOptions(m_options, onnxModelPath);
    if (m_options.precision == Precision::FP16) {
        // Ensure the GPU supports FP16 inference
        if (!builder->platformHasFastFp16()) {
            auto msg = "Error: GPU does not support FP16 precision";
            spdlog::error(msg);
            throw std::runtime_error(msg);
        }
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    } else if (m_options.precision == Precision::INT8) {
        if (numInputs > 1) {
            auto msg = "Error, this implementation currently only supports INT8 "
                       "quantization for single input models";
            spdlog::error(msg);
            throw std::runtime_error(msg);
        }

        // Ensure the GPU supports INT8 Quantization
        if (!builder->platformHasFastInt8()) {
            auto msg = "Error: GPU does not support INT8 precision";
            spdlog::error(msg);
            throw std::runtime_error(msg);
        }

        // Ensure the user has provided path to calibration data directory
        if (m_options.calibrationDataDirectoryPath.empty()) {
            auto msg = "Error: If INT8 precision is selected, must provide path to "
                       "calibration data directory to Engine::build method";
            throw std::runtime_error(msg);
        }

        config->setFlag((nvinfer1::BuilderFlag::kINT8));

        const auto input = network->getInput(0);
        const auto inputName = input->getName();
        const auto inputDims = input->getDimensions();
        const auto calibrationFileName = engineName + ".calibration";

        m_calibrator = std::make_unique<Int8EntropyCalibrator2>(m_options.calibrationBatchSize, inputDims.d[3], inputDims.d[2],
                                                                m_options.calibrationDataDirectoryPath, calibrationFileName, inputName,
                                                                subVals, divVals, normalize);
        config->setInt8Calibrator(m_calibrator.get());
    }

    // CUDA stream used for profiling by the builder.
    cudaStream_t profileStream;
    Util::checkCudaErrorCode(cudaStreamCreate(&profileStream));
    config->setProfileStream(profileStream);

    // Build the engine
    // If this call fails, it is suggested to increase the logger verbosity to
    // kVERBOSE and try rebuilding the engine. Doing so will provide you with more
    // information on why exactly it is failing.

    // 构建引擎并保存到磁盘
    // 使用 builder->buildSerializedNetwork 构建序列化的 TensorRT 引擎。
    // 将引擎保存到磁盘,以便后续加载和使用。
    // 为什么要这么写?
    // 构建引擎是一个耗时操作,将引擎保存到磁盘可以避免重复构建。
    // 序列化的引擎文件可以直接加载到内存中,节省时间。
    std::unique_ptr<nvinfer1::IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
    if (!plan) {
        return false;
    }

    // Write the engine to disk
    const auto enginePath = std::filesystem::path(m_options.engineFileDir) / engineName;
    std::ofstream outfile(enginePath, std::ofstream::binary);
    outfile.write(reinterpret_cast<const char *>(plan->data()), plan->size());
    spdlog::info("Success, saved engine to {}", enginePath.string());

    Util::checkCudaErrorCode(cudaStreamDestroy(profileStream));
    return true;
}

函数的作用

解析 ONNX 模型文件。

设置优化配置(动态批量大小、动态输入宽度等)。

设置精度模式(FP16 或 INT8)。

构建 TensorRT 引擎并保存到磁盘。

涉及的 TensorRT API

点击jetson orin nano super AI模型部署之路(五)tensorrt C++ api介绍查看全文

更多推荐