SuperAPI AI 工具集

trt-torch-quickstart

通用与其它

查看项目

更新于 10-16 02:34

技能介绍

用途:将 PyTorch nn.Module 通过 torch_tensorrt 编译为 TensorRT 引擎,无需先导出 ONNX。适合模型含难以干净导出 ONNX 的算子(自定义 autograd、动态控制流)、希望不支持子图自动回退到 PyTorch,或需要可序列化引擎用于 C++ 部署的场景。已有可用 ONNX 文件应改用 ONNX 路径;LLM 文本生成应走 TensorRT-LLM;已有 .plan 文件走 C++ runtime;弱类型网络迁移走强类型迁移技能。

前置条件:Torch-TensorRT 版本与 (torch, tensorrt-cu13, CUDA) 三元组强耦合,混用 -cu12 与 -cu13 会导入失败。TRT 11.0 需要 tensorrt-cu13 11.x、CUDA 13.x(驱动 R590+)、Python ≥3.10,torch 与 torch_tensorrt 按 pytorch/TensorRT releases 的对应版本固定。推荐使用带 PyTorch 的 NGC TensorRT 容器并以 --gpus all 运行。验证命令:python3 -c "import torch, torch_tensorrt, tensorrt; print(torch.version, torch_tensorrt.version, tensorrt.version)"。

主要步骤:模型需在 CUDA 上并处于 eval() 模式;AOT 路径用 torch_tensorrt.dynamo.compile 接收 torch.export.export 的 ExportedProgram,再用 torch_tensorrt.save(trt_gm, 路径, inputs=[...]) 序列化,不要用 submod.engine 取原始字节写文件。torch_tensorrt ≥2.12 默认 use_explicit_typing=True,精度由导出模型 dtype 决定,此时传 enabled_precisions={torch.float16} 会报 AssertionError;需要 FP16 就先对模型和示例张量 .half()。弱类型路径需显式设 use_explicit_typing=False 才生效。仅需 Python 内可调用、不要序列化引擎时,可用 torch.compile(model, backend="torch_tensorrt") 的 JIT 路径,首次调用时惰性编译,不能供 C++ 加载。

动态形状:用 torch_tensorrt.Input 指定 min_shape/opt_shape/max_shape,并与 torch.export.export 的 dynamic_shapes 保持一致,否则会在编译数分钟后才报 builder 错误;opt_shape 应设为最常见运行时形状;只把真正变化的维度标为动态。

验证与限制:FP16 下不建议仅用 assert_close 做逐元素比较,应按模型类型做语义校验(分类看 top-1/top-5 与平均 softmax 距离,嵌入看余弦相似度,检测看有意义子张量)。FP32 基线应满足 rtol=1e-4、atol=1e-4。C++ 加载不在本技能范围:.ep 是 torch.export 归档,需 output_format="torchscript" 或走 ONNX 路径生成 .plan。常见问题包括负维度报错、不支持算子告警、回退过多(可用 TORCH_LOGS="graph_breaks"、调整 min_block_size、排查形状相关控制流)、输出异常多为 dtype 问题、编译过久可开启 builder 日志。本技能不涉及自定义转换器、TensorRT-LLM 与 QAT 训练流程。