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

最新推荐文章于 2025-11-22 09:31:53 发布

原创

最新推荐文章于 2025-11-22 09:31:53 发布 · 1.1w 阅读

24 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #人工智能 #TensorRT #nvidia

本文详细阐述了如何在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