TensorRT 常用函数

摘要

• TensorRT 常用函数

nvinfer1::createInferRuntime() 函数 详解

nvinfer1::createInferRuntime() 函数详解

1. 作用

nvinfer1::createInferRuntime() 用于创建 TensorRT 的推理运行时（Infer Runtime）对象，该对象用于管理 TensorRT 推理相关的操作，比如反序列化引擎、执行推理等。

2. 语法

#include "NvInfer.h"

nvinfer1::IRuntime* createInferRuntime(nvinfer1::ILogger& logger);

• 参数
  • logger：一个 nvinfer1::ILogger 类型的对象，用于记录 TensorRT 运行时的日志信息。
• 返回值
  • 返回一个 nvinfer1::IRuntime* 指针（即推理运行时对象）。
  • 如果创建失败，返回 nullptr。

3. 相关说明

• createInferRuntime() 由 nvinfer1 命名空间提供，需要包含 TensorRT 头文件 NvInfer.h。
• 生成的 IRuntime 对象用于加载 TensorRT 序列化的模型（即 ICudaEngine），然后进行推理。
• 需要在使用完毕后手动释放 IRuntime，调用 IRuntime::destroy() 或 delete 释放资源。

4. 代码示例

#include <NvInfer.h>
#include <iostream>

class Logger : public nvinfer1::ILogger {
public:
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING) {
            std::cout << "[TensorRT] " << msg << std::endl;
        }
    }
};

int main() {
    Logger logger;
    
    // 创建推理运行时
    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
    if (!runtime) {
        std::cerr << "Failed to create TensorRT runtime!" << std::endl;
        return -1;
    }

    std::cout << "TensorRT runtime created successfully!" << std::endl;

    // 释放资源
    runtime->destroy();  // 或者 delete runtime;

    return 0;
}

5. 结合 deserializeCudaEngine() 进行模型加载

通常，我们使用 createInferRuntime() 来反序列化 TensorRT 生成的 .engine 文件：

std::ifstream file("model.engine", std::ios::binary);
file.seekg(0, std::ios::end);
size_t size = file.tellg();
file.seekg(0, std::ios::beg);
std::vector<char> engineData(size);
file.read(engineData.data(), size);
file.close();

// 反序列化 TensorRT 引擎
nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engineData.data(), size, nullptr);
if (!engine) {
    std::cerr << "Failed to load TensorRT engine!" << std::endl;
}

// 释放资源
engine->destroy();
runtime->destroy();

6. 注意事项

• createInferRuntime() 只能用于反序列化引擎（即推理阶段），不能用于创建和编译新的模型（应使用 createInferBuilder()）。
• IRuntime 不支持创建 ICudaEngine，它只能通过 deserializeCudaEngine() 加载已经序列化的 .engine 文件。
• IRuntime 必须在所有 TensorRT 资源（如 ICudaEngine）释放后再销毁，否则可能会导致内存泄漏或错误。

7. 总结

• createInferRuntime() 用于创建 TensorRT 推理运行时。
• 需要提供 ILogger 记录日志。
• 主要用于 deserializeCudaEngine() 反序列化 TensorRT 引擎。
• 需要手动释放资源，调用 destroy()。

这个函数是 TensorRT 进行推理的核心步骤之一，合理使用可以提升深度学习模型的推理性能。 🚀

C++ createInferBuilder() 函数 详解

在C++中，createInferBuilder() 函数通常是与NVIDIA TensorRT库相关的，用于创建一个 TensorRT 的推理构建器（InferBuilder）。TensorRT是一个用于深度学习推理的高性能推理引擎，由NVIDIA提供。

以下是对 createInferBuilder() 函数的详解：

nvinfer1::IBuilder* createInferBuilder(nvinfer1::ILogger& logger);

在这里，createInferBuilder() 函数返回一个 nvinfer1::IBuilder* 指针，它指向 TensorRT 推理构建器（InferBuilder）。让我们对参数和返回值进行详细解释：

• 参数：

  • nvinfer1::ILogger& logger：一个 TensorRT 提供的日志记录器（ILogger）引用。日志记录器用于在 TensorRT 运行时记录各种信息，包括警告、错误等。通过传递一个日志记录器引用，您可以捕获 TensorRT 的运行时消息。

• 返回值：

  • nvinfer1::IBuilder*：返回一个指向 TensorRT 推理构建器的指针。该指针可以用于配置和构建推理引擎。

使用示例：

#include <NvInfer.h>

int main() {
    // 创建一个简单的日志记录器
    nvinfer1::ILogger* logger = new nvinfer1::Logger(nvinfer1::ILogger::Severity::kINFO);

    // 使用 createInferBuilder() 创建 TensorRT 推理构建器
    nvinfer1::IBuilder* builder = createInferBuilder(*logger);

    // 在这里可以配置和构建推理引擎

    // 释放资源
    builder->destroy();
    delete logger;

    return 0;
}

在上面的示例中，首先创建了一个简单的日志记录器，然后使用 createInferBuilder() 函数创建了一个 TensorRT 推理构建器。接下来，您可以使用该构建器配置和构建推理引擎。最后，记得释放资源，包括销毁构建器和释放日志记录器。