如何调用一个具有动态维度的tensorrt engine

本文详细阐述了如何在C++中使用TensorRT处理具有动态维度的模型,涉及如何通过`IExecutionContext`的`setBindingDimensions`方法设置输入维度,并提供了从引擎文件加载、维度检查到推理执行的完整流程示例。

调用具体动态维度的模型engine时如果没有指定维度,会导致报类似这样的错误:

[TRT] Parameter check failed at: engine.cpp::resolveslots::1227, condition: allInputDimensionsSpecified(routine)

在python代码里,在调用engine推理前做这样的设置即可:

context.set_binding_shape(0, (BATCH, 3, INPUT_H, INPUT_W))

在C++代码里改如何设置,很少有文章提及,关于如何调用有动态维度的模型,一般都是举的python代码的例子,我查了一下TensorRT的头文件NvInferRuntime.h里的代码才知道,C++代码里应该调用IExecutionContext类型的实例的setBindingDimensions(int bindingIndex, Dims dimensions)方法。

总体思路是:拿到一个对维度未知的模型engine文件后,首先读入文件内容并做deserialize获得engine:

ARNet::ARNet(std::string engine_file,
                         std::string shape_file, std::string input_name,
                         std::string output_name)
    : mEngine(nullptr) {
  samplesCommon::OnnxSampleParams params;
  params.inputTensorNames.push_back(input_name.c_str());
  params.outputTensorNames.push_back(output_name.c_str());
  params.int8 = false;
  params.fp16 = true;
  mParams = params;
  std::string se_path = engine_file;  //"arnet_b1_fp16.engine";
  if (access(se_path.c_str(), 4) != -1) {
    nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(gLogger);
    std::ifstream fin(se_path);
    std::string cached_engine = "";
    while (fin.peek() != EOF) {
      std::stringstream buffer;
      buffer << fin.rdbuf();
      cached_engine.append(buffer.str());
    }
    fin.close();
    mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(
        runtime->deserializeCudaEngine(cached_engine.data(),
                                       cached_engine.size(), nullptr),
        samplesCommon::InferDeleter());
    if(!mEngine){
      std::cout <<"Deserialize from "<< se_path <<" failed!!! rebuild enigne ..." << std::endl;
      build();
      return;
    }
...

然后调用getBindingDimensions()查看engine的输入输出维度(如果知道维度就不用):

for (int i = 0; i < mEngine->getNbBindings(); i++)
{
     nvinfer1::Dims dims = mEngine->getBindingDimensions(i);
     printf("index %d, dims: (");
     for (int d = 0; d < dims.nbDims; d ++)
     {
         if (d < dims.nbDims -1) 
             printf("%d,", dims.d[d]);
         else
             printf("%d", dims.d[d]);
     }
     printf(")\n");
}

在调用context->executeV2()做推理前把维度值为-1的动态维度值替换成具体的维度并调用context->setBindingDimensions()设置具体维度,然后在数据填入input buffer准备好后调用contex

评论 5
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Arnold-FY-Chen

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值