ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

TensorRT 11 实战:用 Torch-TensorRT 将 PyTorch 模型编译为 TensorRT 引擎(AOT/JIT 双路径与强类型默认行为)

TensorRT 11 实战:用 Torch-TensorRT 将 PyTorch 模型编译为 TensorRT 引擎(AOT/JIT 双路径与强类型默认行为) TensorRT 11 实战用 Torch-TensorRT 将 PyTorch 模型编译为 TensorRT 引擎AOT/JIT 双路径与强类型默认行为【免费下载链接】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 11.x 环境下借助 Torch-TensorRT 前端将 PyTorchnn.Module直接编译为 TensorRT 引擎全程无需 ONNX 中间格式。文章覆盖 AOT 序列化编译面向生产与 C 部署与 JIT 进程内编译面向纯 Python 推理两条路径并重点剖析 Torch-TensorRT ≥ 2.12 默认开启的强类型strong typing行为下 FP16 精度应如何正确表达、动态形状如何配置、引擎如何序列化/反序列化以及编译失败、过度回退、数值异常等高频问题的定位与排查方法。什么时候使用 Torch-TensorRTTorch-TensorRT 是 TensorRT 的 PyTorch 前端它把torch.export导出的计算图划分成若干子图交给 TensorRT 构建引擎不支持的部分自动回退fallback到 PyTorch 执行。本仓库中与之配套的技能文档.agents/skills/trt-torch-quickstart/SKILL.md给出了清晰的选型矩阵场景是否使用本路径手里是 PyTorchnn.Module希望不写 ONNX 中间文件直接得到 TensorRT 引擎是模型含无法干净导出 ONNX 的算子自定义 autograd、动态控制流是希望利用 Torch-TensorRT 对不支持子图自动回退 PyTorch 的能力是从 PyTorch 出发最终要得到可序列化引擎并在 C 侧部署是——走 AOT 路径再由 C 运行时技能加载已有可用的 ONNX 文件否——走 ONNX 解析路径配套技能trt-onnx-quickstartLLM 文本生成Llama、Mistral、Qwen 等否——应使用 TensorRT-LLM已有.plan文件只想在 C 侧运行否——走 C 运行时技能迁移旧的弱类型weakly-typedTRT 网络到强类型否——走强类型迁移技能.agents/skills/trt-strong-typing-migration/SKILL.md简单概括没有 ONNX、想从 PyTorch 直达引擎、且需要自动回退兜底时Torch-TensorRT 是首选涉及 LLM 生成任务请直接使用 TensorRT-LLM。环境准备版本固定矩阵是前提Torch-TensorRT 与(torch, tensorrt-cu13, CUDA)三元组强耦合版本错配会在import torch_tensorrt时直接抛出 ABI 错误。以 TensorRT 11.0 为例的推荐组合如下组件要求说明TensorRT11.xtensorrt-cu13wheel混用-cu12与-cu13wheel 会破坏 importCUDA 工具包13.x驱动需 R590 及以上Python≥ 3.10TRT 11 已放弃 3.9 及更早版本torch以所选torch_tensorrt版本的 release notes 为准由官方发布页给出的固定矩阵决定torch_tensorrt与 TRT 11.0 匹配的发布版本版本矩阵的唯一事实来源是官方 releases 页不要凭记忆编造版本号推荐的运行环境是 NGC 的 TensorRT/PyTorch 容器nvcr.io/nvidia/pytorch:tag或nvcr.io/nvidia/tensorrt:tag之上再 pip 安装torch_tensorrt启动时加--gpus all。装好后执行如下验证命令确认三个关键库的版本能够协同工作python3 -c import torch, torch_tensorrt, tensorrt; print(torch.__version__, torch_tensorrt.__version__, tensorrt.__version__)第一步加载并准备模型使用eval()模式下的torch.nn.ModuleTorch-TensorRT 会对其做追踪traceimport torch import torchvision.models as models model models.resnet50(weightsNone).eval().cuda() example torch.randn(1, 3, 224, 224, devicecuda)三个要点weightsNone跳过权重下载适合做形状与性能测试做精度验证时应加载真实权重。模型必须已位于 CUDA 上再编译Torch-TensorRT 不会替你搬移设备。.eval()很重要BatchNorm 与 Dropout 在训练模式下行为不同可能产出不同的引擎。第二步AOT 编译到 TensorRT 引擎两条编译路径按部署目标选择目标路径API序列化引擎生产、C 部署、复用AOTtorch_tensorrt.dynamo.compile→torch_tensorrt.save仅进程内 Python 可调用对象无需序列化引擎JITtorch.compile(backendtorch_tensorrt)AOT 路径使用torch_tensorrt.dynamo.compile返回可序列化的ExportedProgram。强类型是默认行为FP16 靠模型 dtype 表达在 torch_tensorrt ≥ 2.12 中强类型use_explicit_typingTrue是默认值。此时引擎精度由导出模型的 dtype 推断如果仍传enabled_precisions{torch.float16}会抛出AssertionError错误信息形如 enabled_precisions should not be used when use_explicit_typingTrue。要编译 FP16应先把模型与示例张量 cast 成 FP16 再导出import torch import torch_tensorrt model MyModel().eval().cuda().half() # cast 到 FP16 example torch.randn(1, 3, 224, 224, devicecuda, dtypetorch.float16) trt_gm torch_tensorrt.dynamo.compile( torch.export.export(model, (example,)), inputs[example], # 不传 enabled_precisions —— 精度来自模型 dtype truncate_doubleTrue, min_block_size1, ) # 序列化供后续 / C 加载 —— 必须用 torch_tensorrt.save而不是裸字节 torch_tensorrt.save(trt_gm, resnet50_trt.ep, inputs[example])不要通过submod.engine取出原始引擎字节再用open(...).write(...)写文件这绕过了 Torch-TensorRT 的元数据包装产物无法经torch_tensorrt.load(...)回读还会悄悄丢失处理部分回退子图的调度图。torch_tensorrt.save(trt_gm, path, inputs...)是唯一受支持的序列化路径。如果需要混合精度或希望覆盖模型 dtype可以显式关闭强类型use_explicit_typingFalse此时enabled_precisions恢复原有语义即弱类型路径在 TRT 11 中已标记弃用trt_gm torch_tensorrt.dynamo.compile( torch.export.export(model, (example,)), inputs[example], use_explicit_typingFalse, # 弱类型已弃用 enabled_precisions{torch.float16, torch.float32}, truncate_doubleTrue, )关键参数说明truncate_doubleTrue静默地把 FP64 常量降为 FP32。不开启时任何 FP64 算子都会强制产生一个分区边界。min_block_size决定至少多少个节点的子图才值得交给 TRT。1是激进设置默认值5可避免把收益盖不过 launch 开销的微小 TRT 子图交给引擎。workspace_sizeTRT builder 可用的 scratch 内存字节数。不设置时由 TRT 自行决定。JIT 替代方案纯 Python 推理如果只要一个可调用对象、不需要序列化引擎用torch.compile Torch-TensorRT 后端即可。在强类型默认下精度同样来自模型 dtype——要 FP16 就.half()模型而不是传enabled_precisionsmodel model.half() # 强类型默认下由模型 dtype 推断 FP16 example example.half() trt_model torch.compile(model, backendtorch_tensorrt) out trt_model(example) # 首次调用时惰性编译之后进程内走 TRT该方式跳过了torch.export得到的是torch.compile包装的可调用对象无法被 C 加载——需要 C 侧加载时请走上面的 AOT 路径。仓库源码佐证强类型构建的落地形态本仓库虽然没有直接包含 Torch-TensorRT 的 C/Python 绑定源码但强类型是 TRT 11 全链路的一致性行为可以在多个 Python 示例中找到直接证据strongly_type_autocast 示例 用builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))显式创建强类型网络然后从 ONNX 解析、build_engine_with_config构建、engine.serialize()落盘。network_api_pytorch_mnist 示例 同样以NetworkDefinitionCreationFlag.STRONGLY_TYPED创建网络且全程没有出现任何BuilderFlag.FP16之类的精度提示 flag——这正是强类型下精度写在图里、而非写在构建配置里的典型写法。trtexec 的 README Example 6 说明--stronglyTyped会以NetworkDefinitionCreationFlag::kSTRONGLY_TYPED创建网络张量数据类型由网络输入类型与算子类型规格推断且不允许再配合--int8、--best等精度类 flag。Torch-TensorRT 强类型路径与上述行为完全一致精度信息必须显式存在于图中例如.half()之后导出的模型天然带 FP16 dtype 语义构建配置不再承担精度提示职责。第三步动态形状用torch_tensorrt.Input包装输入并给出min_shape/opt_shape/max_shapedynamic_input torch_tensorrt.Input( min_shape(1, 3, 224, 224), opt_shape(8, 3, 224, 224), max_shape(32, 3, 224, 224), dtypetorch.float16, ) trt_gm torch_tensorrt.dynamo.compile( torch.export.export(model, (example,), dynamic_shapes{x: {0: torch.export.Dim(batch, min1, max32)}}), inputs[dynamic_input], # 强类型下FP16 来自输入/模型 dtype而不是 enabled_precisions )三个高频坑torch.export.export的dynamic_shapes参数必须与Input的范围一致。不一致时不会在导出阶段报错而是编译进行到 23 分钟时 builder 才报错。opt_shape是 TRT 调优内核所针对的形状应设为最常见的运行时形状而不是中间值。只把运行时确实变化的维度标为动态多余的动态维度会损害性能。第四步运行推理并验证对 FP32 而言裸张量对比assert_close够用但对FP16 会产生误导——即便下游行为完全一致个别离群位置也可能超出容差。应改用与模型类型匹配的语义检查with torch.inference_mode(): torch_out model(example) trt_out trt_gm(example)图像分类——比较 top-1或 top-5类别是否一致外加 softmax 输出的平均距离assert torch.equal(torch_out.argmax(-1), trt_out.argmax(-1)), top-1 mismatch mean_prob_diff (torch.softmax(torch_out.float(), -1) - torch.softmax(trt_out.float(), -1)).abs().mean() assert mean_prob_diff 1e-3句子/词元嵌入——按 token 计算余弦相似度再对整个 batch 取均值cos torch.nn.functional.cosine_similarity(torch_out.flatten(0, -2), trt_out.flatten(0, -2), dim-1) assert cos.mean() 0.999检测/回归头——对有意义的子张量如 bbox 坐标做有界的裸张量近似对比对类别 logits 做 softmax 检查。FP32 冒烟——任何模型下FP32不.half()的引擎都应通过torch.testing.assert_close(trt_out, torch_out, rtol1e-4, atol1e-4)。如果这一步失败问题出在编译本身而不是数值精度。加载序列化引擎Python 侧import torch_tensorrt loaded torch_tensorrt.load(resnet50_trt.ep).module() out loaded(example)C 侧本技能不覆盖 C 加载。torch_tensorrt.save默认产出的.ep文件是面向 Pythontorch_tensorrt.load(...)的torch.export归档不能直接被 TensorRT C 运行时IRuntime/deserializeCudaEngine消费——后者期望的是序列化的 TensorRT 引擎.plan。两条受支持的 C 路线保存时指定output_formattorchscript用 libtorch Torch-TensorRT C 运行时部署 TorchScript 模块若只需要裸 TensorRT 引擎则改走 ONNX 路径trt-onnx-quickstart构建再用 C 运行时技能现代IRuntimeenqueueV3setTensorAddress加载得到的.plan。常见问题导出时报RuntimeError: Trying to create tensor with negative dimension——模型存在依赖形状的控制流分支torch.export无法追踪。要么用torch.cond/torch.where重写要么退回irdynamomin_block_size1让不支持的区域留在 PyTorch 中运行。Unsupported operator警告——Torch-TensorRT 会在这些算子周围做图分区把相应节点留在 PyTorch 执行。设置require_full_compilationTrue可把这类回退升级为报错开发期用它找出究竟哪些算子实际发生了回退。ImportError: cannot import name XYZ from torch_tensorrt——版本不匹配。核对版本固定矩阵torch_tensorrt的 Python API 在小版本之间会变动。诊断过度的 PyTorch 回退如果编译成功但模型大部分回退到了 PyTorch只有少数子图跑在 TRT 上按顺序尝试三个旋钮先确认回退了什么。用TORCH_LOGSgraph_breaks配合require_full_compilationTrue把回退升级成能点名肇事算子的报错。没有这一步你根本不知道 autotuner 拒绝了什么。TORCH_LOGSgraph_breaks python3 your_compile.py重新审视min_block_size。这是一个真实的权衡不只是个旋钮。默认min_block_size5只有子图节点数 ≥ 5 时才交给 TRT。微小的 TRT 子图在内核 launch 开销上的花费超过其收益。min_block_size1激进——所有受支持的算子都进 TRT哪怕单节点子图。适合观察理论上支持什么但因每个子图都有 launch 税推理时通常更慢。建议生产环境保持默认5仅在诊断期降到1以观察最大 TRT 覆盖。查找形状相关控制流。torch.export只追踪if x.shape[0] 0:这类代码的一个分支用torch.where、torch.cond重写或把条件提到模型之外。trace 错误会点名肇事算子。三步之后回退依旧那就是真正不支持的算子——需要自行编写 converter见 Torch-TensorRT 上游文档或改走 ONNX 路径。引擎能构建但输出是乱码——几乎都是 dtype 问题。依次检查模型/输入 dtype 与你的管线一致——强类型下精度跟随模型 dtypeFP16 用.half()enabled_precisions仅在弱类型use_explicit_typingFalse下生效。模型含 FP64 常量时位置编码常见开启truncate_doubleTrue。导出前模型处于.eval()状态。编译超过 10 分钟——打开 builder 日志观察import torch_tensorrt.logging as ttlog ttlog.set_reportable_log_level(ttlog.Level.Info)通常是 autotuner 在为某个重 GEMM/卷积探索内核变体。要限制探索范围可把workspace_size调小scratch 越少 → 候选越少。数值调试如果assert_close超出容差约 5 倍以上先回到 FP32 复跑——去掉模型的.half()或弱类型下用enabled_precisions{torch.float32}。若与 PyTorch 对齐问题就是某个算子的 FP16 累加漂移通常是 softmax 或 LayerNorm。务实的修法是把出问题的子模块包成从编译中显式排除的模块min_block_size帮不上忙必须显式排除具体排除 API 见 Torch-TensorRT 官方 lowering 指南。先在 FP32 冒烟再引入 FP16 并二分定位——几乎总能揪出单个肇事算子。FP8 场景确认校准数据分布确实覆盖了运行时分布。FP8 对离群点没有任何余量。更完整的弱类型 vs 强类型讨论参见.agents/skills/trt-strong-typing-migration/SKILL.md。TRT 11.0 中 Torch-TensorRT 默认产出强类型引擎——类型从导出的程序中推断但你可以覆盖这一行为。强类型的仓库级佐证AutoCast 混合精度工作流强类型要求精度写在图里这在纯 ONNX 场景下同样成立。本仓库的 strongly_type_autocast 示例 给出了完整参考先用 ModelOpt 的 AutoCast 把 FP32 ONNX 转成 FP32-FP16 混合精度模型AutoCast 自动挑选需要保留 FP32 的节点并自动注入 Cast 算子再用NetworkDefinitionCreationFlag.STRONGLY_TYPED构建引擎。其中convert_to_mixed_precision的核心参数见 sample.pylow_precision_typefp16目标低精度类型fp16或bf16nodes_to_exclude保持 FP32 的节点名正则列表示例中排除Plus214配合op_types_to_exclude[MatMul]data_max/init_max可转换节点的最大绝对 I/O 值 / 初始化器最大绝对值示例取 4.0keep_io_typesTrue是否保留图输入/输出类型下游期望 FP32 I/O 时建议开启calibration_data参考 runner 的输入数据文件路径示例使用8.npz。示例的完整流水线STAGE 1→2→3是ONNX Runtime 在原始 FP32 模型上跑推理生成 gold 参考 → AutoCast 转换后再次用 ONNX Runtime 验证并以np.allclose(..., rtol5e-3, atol5e-3)对比 → 强类型构建 TRT 引擎、TRT 运行时推理、与 gold 参考对比。该流程同样印证先确保转换后模型在 ONNX Runtime 侧达标再交给 TRT否则任何引擎侧的工作都无法挽回图级转换错误。本路径不涵盖什么不涉及编写自定义 Torch-TensorRT converter——请参考上游 Torch-TensorRT 文档。不针对 TensorRT-LLM / LLM 推理——请使用 TRT-LLM 的示例流程。不涉及 QAT量化感知训练——Torch-TensorRT 消费 QAT 产出的模型但训练循环在上游。参考资源上游 Torch-TensorRThttps://pytorch.org/TensorRT/版本固定矩阵官方 pytorch/TensorRT releases 页本仓库内的强类型迁移指南.agents/skills/trt-strong-typing-migration/SKILL.md强类型 AutoCast 混合精度完整示例samples/python/strongly_type_autocast/强类型 Python builder 示例samples/python/network_api_pytorch_mnist/sample.pytrtexec 强类型用法samples/trtexec/README.md Example 6【免费下载链接】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),仅供参考
返回列表