CUDA：engine与context（附示例代码

在 TensorRT 中，ICudaEngine 和 IExecutionContext 之间的关系是层次性和功能性的：

1. ICudaEngine（代表由engine变量表示）：

   • ICudaEngine 是模型的优化版本，它是通过解析和优化训练好的深度学习模型得到的。它包含了执行推理所需的所有操作和参数。
   • 您可以将 ICudaEngine 视为一个准备好的、优化过的模型，它知道如何进行高效的计算，但尚未准备好直接执行这些计算。

2. IExecutionContext（由context变量表示）：

   • IExecutionContext 是从 ICudaEngine 创建的，用于执行推理。它是执行环境的实例，可以进行实际的推理计算。
   • 您可以将 IExecutionContext 理解为对特定输入执行模型的“会话”。它保持了模型推理期间的状态，例如输入和输出的内存。

总的来说，ICudaEngine 提供了执行模型所需的所有信息和资源，而 IExecutionContext 则利用这些信息来进行实际的推理运算。在实际使用中，通常首先从 .plan 文件加载或生成 ICudaEngine，然后从这个引擎创建一个或多个 IExecutionContext，用于处理实际的推理请求。这种设计允许同一个引擎（模型）被用于多个并行的推理任务，每个任务使用各自的执行上下文。

示例代码：

从已有的.plan加载

如果有plan文件存在的情况下

初始化engine

ICudaEngine *engine = nullptr;

读取.plan文件

将plan文件读到engineFile中

std::ifstream engineFile(trtFile, std::ios::binary);
long int      fsize = 0;

engineFile.seekg(0, engineFile.end);
fsize = engineFile.tellg();
engineFile.seekg(0, engineFile.beg);
std::vector<char> engineString(fsize);
engineFile.read(engineString.data(), fsize);
if (engineString.size() == 0)
{
    std::cout << "Failed getting serialized engine!" << std::endl;
    return;
}
std::cout << "Succeeded getting serialized engine!" << std::endl;

创建context（在engine创建之后

IExecutionContext *context = engine->createExecutionContext();

获取网络结构信息，并设置输入张量，打印查看

long unsigned int        nIO     = engine->getNbIOTensors();    // 获取输入输出张量总数
long unsigned int        nInput  = 0;                           // 存储输入张量的数量
long unsigned int        nOutput = 0;                           // 存储输出张量的数量
std::vector<std::string> vTensorName(nIO);                      // 存储每个张量的名称

for (int i = 0; i < nIO; ++i)
{
    vTensorName[i] = std::string(engine->getIOTensorName(i));   //获取张量名称
    //打印调试
    std::cout << i<<"号tensor的name为：" << vTensorName[i] << std::endl;
    nInput += int(engine->getTensorIOMode(vTensorName[i].c_str()) == TensorIOMode::kINPUT);
    nOutput += int(engine->getTensorIOMode(vTensorName[i].c_str()) == TensorIOMode::kOUTPUT);
}


context->setInputShape(vTensorName[0].c_str(), Dims32 {3, {3, 4, 5}});

for (int i = 0; i < nIO; ++i)
{
    //engine与context设置输入的张量信息，并打印查看
    std::cout << std::string(i < nInput ? "Input [" : "Output[");
    std::cout << i << std::string("]-> ");
    std::cout << dataTypeToString(engine->getTensorDataType(vTensorName[i].c_str())) << std::string(" ");
    std::cout << shapeToString(engine->getTensorShape(vTensorName[i].c_str())) << std::string(" ");
    std::cout << shapeToString(context->getTensorShape(vTensorName[i].c_str())) << std::string(" ");
    std::cout << vTensorName[i] << std::endl;
}

context->getTensorShape 获取的是当前推理过程中的张量形状，

而 engine->getTensorShape 获取的是构建引擎时设置的原始张量形状。

参考链接：TensorRT学习笔记1（build、engine、context介绍） - 知乎