使用 TensorRT 将 ONNX 转换为 Engine 代码注释

本文将 deepsort 的 .onnx 转 .engine 代码结合Kimi的K1.5模型进行了注释分析，上传自用并给有缘的朋友提供微末的帮助。

---

使用 TensorRT 将 ONNX 转换为 Engine 的完整步骤

1. 创建 builder、network 和 config
   • builder：创建 IBuilder 对象，用于构建 TensorRT 引擎。
   • network：创建 INetworkDefinition 对象，用于定义网络结构。
   • config：创建 IBuilderConfig 对象，用于配置构建过程中的参数（如精度、工作空间大小等）。
2. 设置输入形状
   • 使用 IOptimizationProfile 定义输入张量的动态范围（最小、最优、最大形状）。
   • 这一步是必要的，因为 TensorRT 需要了解输入张量的可能变化范围，以便优化引擎的性能。
   • 输出形状为什么没有设置？
     • 输出形状通常由网络结构和输入形状决定，不需要显式设置。
     • ONNX 模型中已经定义了网络的结构和层之间的关系，因此 TensorRT 可以根据输入形状自动推导出输出形状。
3. 使用 ONNX 解析器解析 ONNX 模型
   • 创建 nvonnxparser::IParser 对象，用于解析 ONNX 模型。
   • 调用 parser->parseFromFile() 方法，将 ONNX 模型加载到 TensorRT 的 INetworkDefinition 中。
4. 配置构建参数
   • 设置精度（如 FP16）和最大工作空间大小。
   • 这些配置会影响引擎的性能和资源使用。
5. 编译 TensorRT 引擎
   • 使用 builder->buildEngineWithConfig() 方法，根据网络定义和配置生成 TensorRT 引擎。
   • 这一步会进行优化和编译，生成高效的推理引擎。
6. 序列化并保存到磁盘
   • 调用 engine->serialize() 方法，将引擎序列化为二进制数据。
   • 将序列化后的数据保存到磁盘上的 .engine 文件中，供后续加载和使用。

示例代码

void DeepSortEngineGenerator::createEngine(std::string onnxPath, std::string enginePath) {
    /* 1. 创建 builder / network / config 三大件 */
    // 创建构建器 builder
    auto builder = createInferBuilder(*gLogger);
    assert(builder != nullptr);
    // 标志位：explicitBatch 其中的 kEXPLICIT_BATCH 表示启用显式批量大小模式
    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    // 创建网络 network
    auto network = builder->createNetworkV2(explicitBatch);// 调用 createNetworkV2 创建网络定义对象
    assert(network != nullptr);
    // 创建配置 config
    auto config = builder->createBuilderConfig();// 用于配置构建过程中的参数
    assert(config != nullptr);

    /* 2. 设置动态输入形状，固定输入名 "input"，形状 1×3×128×64~128×3×128×64 */
    // 创建优化配置文件 profile
    auto profile = builder->createOptimizationProfile();// 用于定义输入张量的动态范围
    Dims dims = Dims4{1, 3, IMG_HEIGHT, IMG_WIDTH};     // 1×3×128×64
    // 设置优化配置文件中输入张量的最小维度
    profile->setDimensions(INPUT_NAME.c_str(),// 最小维度为 1×3×128×64
                OptProfileSelector::kMIN, Dims4{1, dims.d[1], dims.d[2], dims.d[3]});
    // 设置优化配置文件中输入张量的最优维度
    profile->setDimensions(INPUT_NAME.c_str(),// 最优维度为 128×3×128×64
                OptProfileSelector::kOPT, Dims4{MAX_BATCH_SIZE, dims.d[1], dims.d[2], dims.d[3]});
    // 设置优化配置文件中输入张量的最大维度
    profile->setDimensions(INPUT_NAME.c_str(),// 最大维度为 128×3×128×64
                OptProfileSelector::kMAX, Dims4{MAX_BATCH_SIZE, dims.d[1], dims.d[2], dims.d[3]});
    // 将优化配置文件添加到构建配置中
    config->addOptimizationProfile(profile);

    /* 3. ONNX 解析 */
    // 创建 ONNX 解析器 parser
    nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, *gLogger);
    assert(parser != nullptr);
    // 使用解析器解析 ONNX 模型 ,parseFromFile 方法加载并解析 ONNX 文件, 设置日志严重性级别为警告（kWARNING）
    auto parsed = parser->parseFromFile(onnxPath.c_str(), static_cast<int>(ILogger::Severity::kWARNING));
    assert(parsed);

    /* 4. 精度与工作空间设置 */
    // 如果开启 FP16，则在配置中设置相应标志
    if (useFP16) config->setFlag(BuilderFlag::kFP16);
    // 设置构建过程中的最大工作空间大小，这里设置为 1MB
    config->setMaxWorkspaceSize(1 << 20);// 1 << 20是一种位运算表达式，表示将1向左移动20位，1,048,576字节，即 1MB

    /* 5. 编译引擎 */
    // 使用构建器和配置编译 TensorRT 引擎
    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);

    /* 6. 序列化并保存到磁盘 */
    // serialize 方法将引擎序列化为内存中的二进制数据
    IHostMemory* modelStream = engine->serialize();
    // 将序列化的数据复制到一个字符串中
    std::string serializeStr;
    // 用于将数据写入文件
    std::ofstream serializeOutputStream;
    // 使用 resize 方法调整字符串大小以匹配序列化数据的大小
    serializeStr.resize(modelStream->size());
    // 使用 memcpy 将数据从 modelStream 复制到字符串中
    memcpy((void*)serializeStr.data(), modelStream->data(), modelStream->size());
    // 打开目标文件路径，准备写入序列化数据
    serializeOutputStream.open(enginePath);
    // 将序列化的字符串写入文件
    serializeOutputStream << serializeStr;
    // 关闭文件流，确保所有数据都被正确写入
    serializeOutputStream.close();
}