TensorRT 的 Polygraphy 推理运行器:TrtRunner 原理与实战指南 📅 发布时间:2026/9/15 19:03:33 👁 浏览次数: TensorRT 的 Polygraphy 推理运行器TrtRunner 原理与实战指南【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT本文围绕 NVIDIA TensorRT 开源仓库中 Polygraphy 的TrtRunner推理运行器展开讲解如何基于 TensorRT 引擎/执行上下文运行推理、管理输入输出缓冲区、控制动态 shape、优化 profile 与显存分配策略并深入源码说明其内部实现与 CLI 用法。读完本文你将掌握TrtRunner的完整 API、参数语义、底层调用链以及如何把它用于原型验证、测试与调试而非生产部署。说明本文核心依据为仓库文档 runner.rstModule: polygraphy.backend.trt及其指向的 runner.py 实现并结合同目录下的测试用例 test_runner.py、CLI 参数定义 runner.pyargs 与官方示例 build_and_run.py。一、TrtRunner 是什么在 Polygraphy 中Runner运行器是对“在某后端上执行推理”这一行为的统一抽象基类定义在 backend/base/runner.py 的BaseRunner中它负责管理激活/停用activate/deactivate、输入元数据查询get_input_metadata、推理infer以及推理耗时统计last_inference_time。TrtRunner是其中的 TensorRT 后端实现位于 backend/trt/runner.py它接收一个TensorRT 引擎trt.ICudaEngine、执行上下文trt.IExecutionContext或一个返回上述对象的可调用对象callable / 惰性加载器引擎加载完成后由 Polygraphy 自动创建执行上下文并为推理过程中的输入/输出张量分配与管理设备内存它同时承担set_profile()优化 profile 切换、显存分配策略allocation_strategy与权重流式传输weight_streaming_budget/weight_streaming_percent等高级能力的管理。需要特别强调的是源码 docstring 明确指出runners are not designed for production deployment and should generally be used only for prototyping, testing, and debugging.运行器并非为生产部署设计通常只应用于原型开发、测试与调试。这是理解TrtRunner定位的第一原则它是 Polygraphy 生态里“把 TensorRT 推理跑起来、跑得对、跑得快”的实验室工具而不是交付给线上服务的推理框架。1.1 与 BaseRunner 的关系从源码结构看TrtRunner继承BaseRunner只覆写了四个“实现”方法基类方法TrtRunner 覆写职责activate_impl()创建/校验上下文、分配缓冲区激活解析引擎、建上下文、准备设备缓冲get_input_metadata_impl()委托给trt_util.get_metadata_from_engine()返回输入张量名、shape、dtypeinfer_impl()完整推理流水线设置 shape/地址、执行execute_async_v3、回收输出deactivate_impl()释放所有 GPU 资源停用释放缓冲、流与上下文这种“模板方法”设计让BaseRunner.infer()负责公共逻辑激活检查、check_inputs校验、参数透传而 TensorRT 特异的细节全部收敛在TrtRunner内。二、TrtRunner 完整 API 与参数语义2.1 构造函数from polygraphy.backend.trt import TrtRunner runner TrtRunner( engine, # trt.ICudaEngine / trt.IExecutionContext / callable - 上述之一 nameNone, # 可读名称前缀实际名称会追加“序号时间戳” optimization_profileNone, # 激活时设置的优化 profile 索引 allocation_strategyNone, # static / profile / runtime weight_streaming_budgetNone, # 权重流式传输的显存字节预算 weight_streaming_percentNone, # 权重流式传输的百分比 )各参数详细语义如下均来自源码 docstringengine接受trt.ICudaEngine、trt.IExecutionContext或返回两者的可调用对象。传引擎runner 激活时自动调用engine.create_execution_context()创建上下文传上下文直接复用用户提供的上下文此时若同时指定allocation_strategyrunner 会发出警告要求用户确保上下文本身采用了相同策略传可调用对象激活时通过util.invoke_if_callable()惰性求值这正是 Polygraphy 惯用的“加载器链”写法如EngineFromNetwork(...)返回的 loader。name默认名称为trt-runner-N序号-日期-时间。命名规则定义在基类 backend/base/runner.py 中若name未指定则基于prefixTrtRunner 使用trt-runner加全局计数与时间戳生成。optimization_profile引擎激活时要设置的优化 profile 索引None时默认使用第 0 个 profile激活后可通过set_profile()动态切换仅当引擎以多个优化 profile构建时才有意义TensorRT 8.0 及以上版本通过该 runner 的 CUDA 流异步设置 profileset_optimization_profile_async老版本则回退到同步属性active_optimization_profile。allocation_strategy控制执行上下文的设备内存激活值/scratch 内存分配方式可选值取值含义实现要点见_infer_implstatic默认预分配一块足够覆盖所有 profile 任意输入尺寸的内存engine.create_execution_context()由 TRT 自动管理profile按当前 profile 的 max shapes 分配所需内存create_execution_context(USER_MANAGED)用get_device_memory_size_for_profile(_v2)计算并set_device_memoryruntime按当前输入形状分配所需内存update_device_memory_size_for_shapes()后设置实现细节当策略为profile/runtime时Polygraphy 以trt.ExecutionContextAllocationStrategy.USER_MANAGED创建上下文随后通过cuda.DeviceArray.raw()分配context_memory_buffer再调用context.set_device_memory(...)交给 TRT。因此“用户管理内存”后Polygraphy 必须在每次infer()前按需重算内存大小get_device_memory_size_for_profile_v2或update_device_memory_size_for_shapes并resize缓冲——这正是allocation_strategy影响性能与显存占用的核心机制。weight_streaming_budget字节预算取值含义None或-2运行时禁用权重流式传输-1由 TensorRT 自动决定预算 0TensorRT 允许用于权重的最大 GPU 内存字节数weight_streaming_percent百分比取值含义None或100%运行时禁用权重流式传输0 ~ 100流式传输权重的百分比0表示流式传输最多数量的权重注意两者互斥同时指定时 runner 会警告并优先采用字节值见_set_weight_streaming_budget。预算的底层换算因 TensorRT 版本而异V1 路径TRT 10.1百分比换算公式为budget (1 - percent/100) * (max_budget - min_budget) min_budget其中min_budget engine.minimum_weight_streaming_budget、max_budget engine.streamable_weights_sizeV2 路径TRT 10.1 或 TensorRT-RTX-1使用engine.get_weight_streaming_automatic_budget()百分比直接按streamable_weights_size比例计算无论哪条路径最终都会写回engine.weight_streaming_budget/engine.weight_streaming_budget_v2并校验设置是否生效失败则critical报错。前提条件引擎必须以启用权重流式传输的方式构建构建期CreateConfig(weight_streamingTrue)否则运行期设置预算没有意义。2.2 核心方法activate()/deactivate()激活/停用一般通过with TrtRunner(...) as runner:上下文管理器自动完成重复激活会告警停用会释放所有设备缓冲、context_memory_buffer与 CUDA 流。set_profile(index)运行时切换优化 profile要求 runner 已激活TRT 8.0 异步设置。get_input_metadata()返回TensorMetadata输入名、shape、dtype必须在激活后调用shape 中动态维度以None表示。infer(feed_dict, check_inputsTrue, copy_outputs_to_hostTrue, return_raw_buffersFalse)feed_dict输入名 →NumPy 数组 / PolygraphyDeviceView/ PyTorch tensor传 PyTorch tensor 时输出也以 PyTorch tensor 返回输入已在 GPU 上则不额外拷贝copy_outputs_to_hostFalse时返回 GPU 上的DeviceView或 PyTorch GPU tensor默认输出缓冲会被 runner 复用多次推理结果如需保存必须copy.deepcopy(outputs)每次infer()后可从runner.last_inference_time()读取本次推理耗时秒。inference_time属性最近一次infer()的耗时由infer_impl内部time.time()计时写入。三、一次 infer() 的内部调用链TrtRunner把一次推理组织成清晰的流水线以下结合 runner.py 源码逐段拆解3.1 激活阶段activate_impl惰性求值engine_or_context若拿到ICudaEngine按allocation_strategy创建执行上下文static直接create_execution_context()profile/runtime走USER_MANAGED初始化device_input_buffers、host_output_buffers、cuda.Stream()与output_allocator若指定了optimization_profile调用set_profile()。3.2 推理阶段_infer_impl输入处理遍历所有INPUT张量将feed_dict包装成FormattedArrayshape 张量is_shape_inference_io要求位于主机内存普通张量通过trt_util._get_array_on_gpu()见 backend/trt/util.py拷贝/复用 GPU 缓冲HWC 格式适配若张量格式为HWC用get_chw_shape_from_hwc()还原 CHW 语义 shape增量更新仅当context.get_tensor_shape(name) ! array_shape时才调用set_input_shape()仅当地址变化时才set_tensor_address()避免不必要的开销源码注释明确提到“否则就是在做多余工作”Debug ListenerTRT 支持时通过set_all_tensors_debug_state(True)与set_debug_listener()挂上DebugTensorWriter_make_debug_listener推理后可额外取回中间张量fp8/int4/fp4/bfloat16 等不支持的数据类型会告警跳过输出分配器为每个OUTPUT张量设置IOutputAllocator_make_output_allocator支持 NumPy 与 PyTorch 两种底层缓冲set_use_torch根据 feed_dict 是否含 torch tensor 自动切换显存分配allocation_strategy为profile/runtime时按 §2.1 重新计算并设置设备内存执行context.execute_async_v3(self.stream.ptr)失败即critical报错输出回收从output_allocator.buffers取原始字节缓冲结合notify_shape记录的 shape 与DataType生成视图copy_outputs_to_hostTrue时通过_get_array_on_cpu()拷回主机return_raw_buffersTrue或向量化格式时返回FormattedArray同步与收尾self.stream.synchronize()若 Debug Listener 有输出则并入output_buffers。3.3 停用阶段deactivate_impl释放所有设备输入缓冲、context_memory_buffer与 CUDA 流并删除引擎/上下文等内部状态引用避免 GPU 内存泄漏。3.4 输入校验BaseRunner.infer()默认在每次推理前调用check_inputs()比对feed_dict与引擎输入元数据。测试 test_runner.py 覆盖了错误名称Extra inputs in.../The following inputs were not found、错误 dtypeunexpected dtype.与错误 shapeincompatible shape.三类校验失败路径说明不匹配的 feed_dict 会被 Polygraphy 提前拦截。四、CLI 用法--trt 及其子参数Polygraphy 命令行工具通过RunnerSelectArgstools/args/backend/runner_select.py注册各后端 runner 选项TensorRT 对应的参数组定义在 tools/args/backend/trt/runner.py其中get_name_opt_impl()返回(TensorRT, trt)因此命令行使用--trt选择 TensorRT runnerpolygraphy run model.onnx --trt \ --optimization-profile 1 \ --allocation-strategy profile \ --weight-streaming-budget 2GCLI 参数类型说明对应 TrtRunner 参数--trtflag选择 TensorRT runner可多次指定与--onnxrt等组成对比组选择 runner 本身--optimization-profileint推理使用的优化 profile 索引optimization_profile--allocation-strategystatic/profile/runtime激活内存分配方式allocation_strategy--weight-streaming-budgetstr权重流式预算支持单位后缀与百分比weight_streaming_budget/weight_streaming_percent--weight-streaming-budget的解析逻辑见parse_impl以%结尾 → 解析为百分比并设置weight_streaming_percent范围[0, 100]否则按args_util.parse_num_bytes()解析字节数支持G/M/K等单位后缀要求-2、-1或 0最后add_to_script_impl生成TrtRunner(...)调用脚本连同TrtLoadEngineArgs的引擎加载链一起注入。典型对比验证用法同时指定--trt --onnxrtPolygraphy 会对同一组输入分别用 TensorRT 与 ONNX-Runtime 推理并比较结果。五、实战从引擎构建到推理验证5.1 Python API 最小示例以下改编自官方示例 build_and_run.py完整覆盖“ONNX → 网络 → FP16 引擎 → 保存 → TrtRunner 推理”链路import numpy as np from polygraphy.backend.trt import ( CreateConfig, EngineFromNetwork, NetworkFromOnnxPath, SaveEngine, TrtRunner, ) # 惰性加载器链ONNX - TensorRT Network - FP16 引擎 - 序列化保存 build_engine EngineFromNetwork( NetworkFromOnnxPath(identity.onnx), configCreateConfig(fp16True) ) build_engine SaveEngine(build_engine, pathidentity.engine) # 上下文管理器自动激活/停用防止内存泄漏 with TrtRunner(build_engine) as runner: inp_data np.ones(shape(1, 1, 2, 2), dtypenp.float32) outputs runner.infer(feed_dict{x: inp_data}) assert np.array_equal(outputs[y], inp_data) # identity 模型输出等于输入 print(fInference succeeded! 耗时: {runner.last_inference_time():.4f}s)要点build_engine是可调用对象而非引擎本身TrtRunner激活时才求值runner 拥有输出缓冲且会复用多次推理需copy.deepcopy(outputs)也可直接传已加载的引擎或engine.create_execution_context见 load_and_run.py。5.2 动态 shape 与多 profile从测试 test_runner.py 的test_multiple_profiles可以提炼出多 profile 引擎的标准用法from polygraphy.backend.trt import CreateConfig, EngineFromNetwork, NetworkFromOnnxBytes, Profile, TrtRunner profiles [ Profile().add(X, (1, 2, 1, 1), (1, 2, 1, 1), (1, 2, 1, 1)), # min, opt, max Profile().add(X, (1, 2, 1, 1), (1, 2, 2, 2), (1, 2, 4, 4)), Profile().add(X, (1, 2, 4, 4), (1, 2, 8, 8), (1, 2, 16, 16)), ] engine engine_from_network(NetworkFromOnnxBytes(model.loader), CreateConfig(profilesprofiles)) # 方式一构造时指定 with TrtRunner(engine, optimization_profileindex) as runner: ... # 方式二激活后动态切换 with TrtRunner(engine) as runner: runner.set_profile(index) assert runner.context.active_optimization_profile index5.3 异构输入PyTorch tensor 与 DeviceViewPyTorch tensorfeed_dict 中含 torch tensor 时输出也以 torch tensor 返回test_torch_tensors用例验证了 CPU/CUDA 两种 device 与copy_outputs_to_host两种组合DeviceView / DeviceArray已驻留 GPU 的数据零拷贝直接参与推理test_device_views混合输入同一 runner 可交替使用 NumPy、DeviceArray 与 torch tensor 推理test_subsequent_infers_with_different_input_types例外shape 张量is_shape_inference_io必须位于主机内存传入 GPU 缓冲会触发it must reside in host memory报错test_cannot_use_device_view_shape_tensor。5.4 权重流式传输测试test_weight_streaming展示了预算参数的边界用法kwargs {weight_streaming_budget: None, weight_streaming_percent: None} if 0 budget 1: kwargs[weight_streaming_percent] budget * 100 # 百分比 else: kwargs[weight_streaming_budget] int(budget) # 字节数 with TrtRunner(engine, optimization_profile0, **kwargs) as runner: model.check_runner(runner)配合构建期CreateConfig(weight_streamingTrue)测试中通过NetworkFromOnnxBytes(..., strongly_typedTrue)与CreateConfig(weight_streamingTrue)构造可流式引擎即可在推理侧按字节或百分比控制驻留 GPU 的权重规模用于大模型显存受限场景的实验。5.5 数据相关输出 shapezero-copy 场景test_data_dependent_shapesnonzero模型证明TrtRunner支持输出 shape 依赖输入数据内容的引擎IOutputAllocator会在运行时按实际输出大小reallocate_outputPolygraphy 再依据notify_shape回调的形状返回正确视图无需预先知道输出尺寸。六、多线程与生命周期注意事项共享引擎的多 runnertest_multithreaded_runners_from_engine显示同一引擎可创建多个独立TrtRunner各自持有独立上下文与缓冲可安全地在不同线程中并发infer()必须成对 activate/deactivate基类在__del__中会警告“激活但从未停用”可能造成内存泄漏建议始终使用with上下文管理器inference_time 语义BaseRunner.last_inference_time()依赖infer_impl设置inference_time属性若未设置会触发内部错误提示因此自定义 runner 必须遵守该约定TrtRunner已在infer_impl中实现。七、小结与延伸阅读TrtRunner是 Polygraphy 与 TensorRT 之间的“推理执行层”向上统一了BaseRunner的激活/推理/停用生命周期向下封装了执行上下文创建、动态 shape 设置、显存分配策略、优化 profile 切换、权重流式传输、Debug Listener 与输出分配器等 TensorRT 底层机制。无论是单引擎快速验证、多后端结果对比还是动态 shape 与显存优化的实验它都是值得优先使用的入口。可继续深入阅读的仓库材料源码实现backend/trt/runner.py、基类 backend/base/runner.py、TRT 工具函数 backend/trt/util.py测试用例tests/backend/trt/test_runner.pyCLI 参数tools/args/backend/trt/runner.py、tools/args/backend/runner_select.py官方示例examples/api/00_inference_with_tensorrt/、examples/api/01_comparing_frameworks/example.py、examples/api/02_validating_on_a_dataset/example.py相关文档docs/backend/trt/toc.rst、docs/backend/trt/loader.rst、docs/backend/trt/config.rst【免费下载链接】TensorRTNVIDIA® TensorRT™ is an SDK for high-performance deep learning inference on NVIDIA GPUs. This repository contains the open source components of TensorRT.项目地址: https://gitcode.com/GitHub_Trending/tens/TensorRT创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考