TensorRT学习笔记之ONNX转engine

#include <iostream>
#include "NvInfer.h"
#include "NvOnnxParser.h"
 
// 函数定义，用于创建TensorRT的引擎
nvinfer1::ICudaEngine* createEngine(const std::string& onnxModelFile) {
    // 创建builder
    nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(nvinfer1::ILogger::Severity::kWARNING);
    nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
 
    // 创建解析器
    auto parser = nvonnxparser::createParser(*builder, ::gLogger.getTRTLogger());
    const auto explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    nvinfer1::INetworkDefinition* network = builder->createNetworkV2(explicitBatch);
 
    // 解析ONNX模型文件
    parser->parseFromFile(onnxModelFile.c_str(), static_cast<int>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH));
 
    // 序列化网络
    builder->setMaxBatchSize(32);
    config->setMaxWorkspaceSize(1 << 20);
    nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
 
    // 清理资源
    parser->destroy();
    network->destroy();
    config->destroy();
    builder->destroy();
 
    return engine;
}
 
int main() {
    const std::string onnxModelFile = "model.onnx"; // 替换为你的ONNX模型文件路径
    nvinfer1::ICudaEngine* engine = createEngine(onnxModelFile);
 
    if (engine == nullptr) {
        std::cerr << "Engine could not be created." << std::endl;
        return -1;
    }
 
    // 保存引擎到文件
    std::ofstream planFile("model.plan", std::ios::binary);
    if (!planFile) {
        std::cerr << "Could not open plan output file for writing." << std::endl;
        return -1;
    }
    engine->serialize(planFile);
    planFile.close();
 
    // 释放资源
    engine->destroy();
    std::cout << "Engine has been created and saved to disk." << std::endl;
 
    return 0;
}

这段代码展示了如何使用TensorRT的C++ API将ONNX模型转换为TensorRT引擎，并将引擎序列化到一个文件中。代码中包含了创建TensorRT构建器、解析器和网络的步骤，并设置了构建引擎所需的配置。最后，引擎被序列化并保存到了文件系统中，同时释放了所有被创建的TensorRT资源。