onnx之tensorrt推理过程

完整地串联到从用户提供ONNX文件开始的整个流程中。这个过程主要分为加载编译期和运行期两个阶段 全流程纵览从ONNX文件到推理结果下图清晰地展示了整个过程 阶段一模型加载与编译这是最关键的幕后工作阶段用户只提供了一个ONNX文件路径但引擎在背后做了大量优化。用户发起请求用户代码创建Ort::Session并传入ONNX文件路径和配置好的TensorRT EP选项 。模型加载与分割ONNX Runtime (ORT) 核心加载并解析ONNX文件形成一个计算图。然后ORT会根据TensorRT EP支持的算子列表将原始图进行分割Partition把TensorRT能够加速的部分子图标记出来准备交给TensorRT EP处理 。准备NodeComputeInfo对于分割出的子图中的每一个节点TensorRT EP都需要提供一份“说明书”也就是NodeComputeInfo结构体。这个结构体定义了该节点的创建、计算和释放行为。调用create_state_func(核心步骤)EP遍历子图节点为每个节点调用我们在代码中看到的create_state_func这个lambda表达式。创建状态对象这个函数会创建一个TensorrtFuncState对象。std::unique_ptrTensorrtFuncState p std::make_uniqueTensorrtFuncState();就是在这里执行的。这个对象就像一个“百宝箱”里面存放了该节点所需的一切核心资源引用如builder_(TensorRT构建器)、指向engines_和contexts_映射的指针。编译配置如fp16_enable_,int8_enable_,dla_core_等用户指定的选项 。输入输出信息如input_info_,output_info_描述了张量的名称和索引。缓存路径如cache_path_用于读写序列化的engine文件 。同步原语如tensorrt_mu_的指针一个互斥锁用于后续的线程安全。保存状态初始化完毕后通过*state p.release();这个动态分配的TensorrtFuncState对象的所有权被转移给了ONNX Runtime的运行时框架。框架会拿着这个指针等待后续推理时使用。TensorRT引擎构建/加载在create_state_func被调用后或者在其内部的某个环节TensorRT EP会利用TensorrtFuncState中保存的配置调用真正的TensorRT API来构建Build或反序列化Deserialize一个ICudaEngine。如果开启了engine缓存 (engine_cache_enable_)并且缓存文件存在它会直接从磁盘加载 。否则它会使用builder_和网络定义开始耗时的引擎优化和构建过程最后将构建好的engine也保存在与节点关联的状态中 。Session准备就绪当所有节点的NodeComputeInfo都处理完毕并且对应的TensorRT引擎都构建/加载完成后InferenceSession的初始化就完成了随时可以接受推理请求。⚙️ 阶段二模型推理用户发起推理用户代码调用session.Run()并传入实际的输入数据 。调用compute_funcONNX Runtime运行时框架将执行流转到TensorRT EP并传入之前保存的与该子图节点对应的FunctionState指针它实际上指向我们创建的TensorrtFuncState对象。恢复状态与执行推理compute_func(同样是lambda) 开始执行。恢复状态它做的第一件事就是TensorrtFuncState* trt_state reinterpret_castTensorrtFuncState*(state);将泛型的FunctionState转回具体的TensorrtFuncState指针从而可以访问之前保存的所有资源和配置。线程安全std::lock_guardstd::mutex lock(*(trt_state-tensorrt_mu_ptr));加锁保证在多线程环境下对TensorRT资源的访问是安全的。执行推理函数从状态对象中取出之前构建好的engine和context并从ONNX Runtime的OrtKernelContext中获取实际的输入张量和CUDA流 (stream)然后调用TensorRT的核心推理方法trt_context-enqueueV3(stream)。处理结果等待GPU计算完成后将输出张量的数据映射回ONNX Runtime的输出。返回结果推理结果沿着调用栈返回最终用户拿到模型的输出。 阶段三资源清理当InferenceSession被销毁时ONNX Runtime会调用每个节点的release_state_func。这个函数非常简单就是delete static_castTensorrtFuncState*(state);确保之前动态分配的TensorrtFuncState对象被正确释放避免内存泄漏。总结一下你提到的create_state_func,compute_func,release_state_func三个lambda共同定义了TensorRT EP上一个算子节点的完整生命周期。而TensorrtFuncState对象则是这个生命周期的数据核心它像一个智慧的“时间胶囊”在编译期被创建并装满各种配置和资源引用然后在运行期被打开用于指导高效、安全的推理计算。