SuperAPI AI 工具集

trt-onnx-quickstart

通用与其它

查看项目

更新于 10-16 07:24

技能介绍

用途:把“有 Hugging Face 模型 ID 或 ONNX 文件”带到“得到输出与原模型在容差内一致的 TensorRT 引擎”,对应 Import Workflows Guide 的 Path 1(ONNX → TensorRT)。适用于有 HF 模型 ID(如 google-bert/bert-base-uncased)想做 TRT 加速推理、有 .onnx 想生成带校验的 .plan 引擎、或 trtexec 报出陌生警告/错误时。LLM 文本生成应转 TensorRT-LLM;纯 PyTorch 用 trt-torch-quickstart;已有 .plan 跑 C++ 推理用 trt-cpp-runtime-quickstart;弱类型迁移用 trt-strong-typing-migration;扩散流水线需先拆分组件。

前提:NVIDIA GPU 与匹配 TensorRT 11.x 的驱动(TRT 11 需 CUDA 13.x,用 nvidia-smi 确认);Python ≥ 3.10;安装 TensorRT 11 并包含 trtexec(pip 包不含 trtexec,推荐 NGC 容器 nvcr.io/nvidia/tensorrt:<tag> 并以 --gpus all 运行,或用 .tar.gz/.deb)。验证:python3 -c "import tensorrt; print(tensorrt.version)" 与 trtexec --help。安装 pip install polygraphy onnx onnxruntime onnx_graphsurgeon onnxsim 'optimum-onnx',模型 ID 以 sentence-transformers/ 开头时另装 sentence-transformers;optimum-onnx 为维护中的包,旧 optimum 已不含 ONNX 集成。

流程:先确认模型受支持(扩散流水线需拆分);用 optimum-cli export onnx --model ... --task ... 导出,--task 选错会得到有效但错误的图,且后续校验无法发现;导出器打印的 max diff 告警属正常自检;受限模型需先 huggingface-cli login。可选清洗:onnxsim 与 polygraphy surgeon sanitize --fold-constants。用 polygraphy inspect model 读取输入名与动态轴,据此设置 --minShapes/--optShapes/--maxShapes。TRT 11 无 --fp16 等弱类型标志,先检测主版本,用 python3 -m modelopt.onnx.autocast 烘焙 FP16 再由 trtexec 构建;TRT 10 则传 --fp16。最后用 polygraphy run 做 ONNX Runtime 与 TRT 的数值一致性校验,注意为受限输入设置 --val-range、用 --check-error-stat mean 并设置 --seed。构建后可用 trtexec --loadEngine=... --shapes=... 验证运行。限制:plan 不跨计算能力移植,需在部署 GPU 重建;构建 OOM 可降低 --memPoolSize 或禁用昂贵 tactic;不支持的算子、INT64、缺失形状信息等错误需按提示处理,导出失败应转 trt-export-rewrite,不对上游库打补丁。费用与商业条款资料未提供。