0
点赞
收藏
分享

微信扫一扫

deepsort使用tensorrt加速部署1——加载Onnx并转换成engine代码解读并保存engine文件


代码见:
https://github.com/RichardoMrMu/yolov5-deepsort-tensorrt

代码整体的一个流程如下:
创建InferBuilder对象——创建Network对象——创建BuilderConfig对象——创建OptimizationProfile对象——创建OptimizationProfile对象——设置输入维度——创建解析器对象——解析ONNX模型文件——使用BuilderConfig对象构建TensorRT引擎——序列化引擎并保存到文件——创建执行上下文

代码部分在featuretensor.cpp文件中

void FeatureTensor::loadOnnx(std::string onnxPath)
{
    // 创建InferBuilder对象
    auto builder = createInferBuilder(*gLogger);
    assert(builder != nullptr);
    // 创建Network对象
    const auto explicitBatch = 1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH);
    auto network = builder->createNetworkV2(explicitBatch);
    assert(network != nullptr);
    // 创建BuilderConfig对象
    auto config = builder->createBuilderConfig();
    assert(config != nullptr);
    // 创建OptimizationProfile对象
    auto profile = builder->createOptimizationProfile();
    Dims dims = Dims4{1, 3, imgShape.height, imgShape.width};
    // 设置输入维度
    profile->setDimensions(inputName.c_str(),
                           OptProfileSelector::kMIN, Dims4{1, dims.d[1], dims.d[2], dims.d[3]});
    profile->setDimensions(inputName.c_str(),
                           OptProfileSelector::kOPT, Dims4{maxBatchSize, dims.d[1], dims.d[2], dims.d[3]});
    profile->setDimensions(inputName.c_str(),
                           OptProfileSelector::kMAX, Dims4{maxBatchSize, dims.d[1], dims.d[2], dims.d[3]});
    config->addOptimizationProfile(profile);
    // 创建解析器对象
    nvonnxparser::IParser *parser = nvonnxparser::createParser(*network, *gLogger);
    assert(parser != nullptr);
    // 解析ONNX模型文件
    auto parsed = parser->parseFromFile(onnxPath.c_str(), static_cast<int>(ILogger::Severity::kWARNING));
    assert(parsed);
    // 设置最大工作空间大小
    config->setMaxWorkspaceSize(1 << 20);
    // 使用BuilderConfig对象构建TensorRT引擎
    engine = builder->buildEngineWithConfig(*network, *config);
    assert(engine != nullptr);

    // 序列化引擎并保存到文件
    std::string enginePath = "./deepsort.engine";
    std::ofstream engineFile(enginePath, std::ios::binary);
    if (engineFile)
    {
        nvinfer1::IHostMemory *serializedEngine = engine->serialize();
        engineFile.write(static_cast<const char *>(serializedEngine->data()), serializedEngine->size());
        serializedEngine->destroy();
        engineFile.close();
    }
    else
    {
        std::cerr << "Failed to open engine file for writing: " << enginePath << std::endl;
    }
    // 创建执行上下文
    context = engine->createExecutionContext();
    assert(context != nullptr);
    initResource();
}

上面代码我新增了读取onnx帮保存为engine的代码,这样就不用每次重新生成了。但是换硬件平台的时候,还是需要重新生成。

单独保存engine的代码如下:

// 检查引擎是否已创建
    assert(engine != nullptr);

    // 序列化引擎并保存到文件
    std::string enginePath = "./deepsort.engine";
    std::ofstream engineFile(enginePath, std::ios::binary);
    if (engineFile) {
        nvinfer1::IHostMemory* serializedEngine = engine->serialize();
        engineFile.write(static_cast<const char*>(serializedEngine->data()), serializedEngine->size());
        serializedEngine->destroy();
        engineFile.close();
    } else {
        std::cerr << "Failed to open engine file for writing: " << enginePath << std::endl;
    }

deepsort使用tensorrt加速部署1——加载Onnx并转换成engine代码解读并保存engine文件_c++


举报

相关推荐

0 条评论