4、模型加载与转换:使用trtexec加载ONNX、UFF、Caffe等模型,理解--onnx、--deploy、--model等参数

trtexec是TensorRT官方提供的命令行性能测试工具,它不仅是跑benchmark的利器,更是验证模型能否被TensorRT正确解析、转换和优化的第一道关口。在实际工程中,我们拿到的模型可能来自PyTorch导出的ONNX、TensorFlow冻结的UFF(现已基本废弃)、Caffe的prototxt/caffemodel,甚至是已经序列化好的TensorRT引擎文件(.plan/.engine)。trtexec通过不同的参数入口来区分这些格式,并触发对应的解析器(Parser)完成从框架模型到TensorRT网络定义的转换。理解这些加载参数,是后续所有性能调优的前提——如果模型加载阶段就出错,后续的精度校验、延迟测试、吞吐测试都无从谈起。

trtexec 模型加载与转换链路 ONNX 模型 --onnx=model.onnx PyTorch/ TF导出 UFF 模型 --uff=model.uff TF1.x 冻结图 Caffe 模型 --deploy + --model prototxt + caffemodel TensorRT 引擎 --loadEngine=model.plan 直接加载序列化引擎 TensorRT Parser ONNXParser UffParser CaffeParser Network Definition (INetworkDefinition) 统一中间表示 Builder 优化 FP16/INT8/FP32 层融合/内核选择 显存优化 --saveEngine 生成 .plan 文件 --uff 已废弃,推荐 ONNX;--loadEngine 跳过解析直接反序列化

上图展示了trtexec加载不同格式模型的完整链路。从左侧的四种输入格式出发,经过对应的Parser解析为TensorRT的Network Definition,再通过Builder进行优化,最终生成可执行的引擎。值得注意的是,--loadEngine路径绕过了Parser,直接加载已经优化好的序列化引擎,这是生产部署中最常用的方式。下面我们逐一拆解每个参数的使用方法、适用场景和注意事项。

4.1 核心参数总览:--onnx、--deploy、--model、--uff、--loadEngine

trtexec的参数设计遵循“一个模型入口对应一个参数”的原则。下表汇总了最常用的模型加载参数及其对应关系:

参数名 模型格式 典型文件后缀 配套参数 适用场景
--onnx ONNX .onnx --inputShapes(可选) PyTorch/TF2导出,当前主流格式
--deploy Caffe prototxt .prototxt --model(必填) 老项目Caffe模型迁移
--model Caffe caffemodel .caffemodel --deploy(必填) 与--deploy成对使用
--uff UFF .uff --inputShapes(必填) TF1.x冻结图,已废弃,不推荐新项目
--loadEngine TensorRT引擎 .plan / .engine 无 生产部署,跳过解析直接加载

这里需要特别强调:--deploy和--model必须成对出现,前者指向网络结构文件(prototxt),后者指向权重文件(caffemodel)。如果只给其中一个,trtexec会直接报错退出。而--uff参数虽然仍然保留在trtexec中,但NVIDIA官方已明确宣布UFF Parser在TensorRT 8.x之后不再维护,建议所有TF用户通过ONNX中间格式转换。

4.2 操作流程:从原始模型到TensorRT引擎的标准步骤

无论输入哪种格式,trtexec的加载流程都遵循“解析→构建→序列化→加载”四个阶段。以下是标准操作流程:

  1. 确定模型格式:检查文件后缀,确认是ONNX、Caffe还是UFF。如果是.plan文件,直接跳到第5步。
  2. 选择对应参数:ONNX用--onnx;Caffe用--deploy+--model;UFF用--uff。
  3. 指定输入维度(可选):对于动态batch或动态分辨率模型,使用--inputShapes指定实际推理时的shape,例如--inputShapes=input:1x3x224x224。
  4. 执行构建与保存:添加--saveEngine=output.plan将优化后的引擎保存到磁盘,避免每次启动都重新构建。
  5. 加载引擎测试:后续测试直接用--loadEngine=output.plan,跳过解析和构建阶段,大幅缩短启动时间。

一个典型的ONNX加载命令如下:

# 加载ONNX并保存引擎
trtexec --onnx=resnet50.onnx --saveEngine=resnet50.plan --fp16

# 直接加载引擎进行benchmark
trtexec --loadEngine=resnet50.plan --fp16 --shapes=input:1x3x224x224

对于Caffe模型,命令格式为:

trtexec --deploy=deploy.prototxt --model=resnet50.caffemodel --saveEngine=resnet50.plan

这里有个容易忽略的细节:Caffe的prototxt中如果包含input_dim字段,trtexec会直接读取;如果使用input_shape(旧版Caffe),则需要手动通过--inputShapes指定。

4.3 实战案例:ONNX模型加载失败排查与解决

下面通过一个真实案例来展示模型加载过程中可能遇到的问题及排查思路。某团队需要将一个PyTorch训练的YOLOv5s模型部署到TensorRT上,他们首先尝试了最直接的ONNX加载方式。

初始现象:执行trtexec --onnx=yolov5s.onnx --saveEngine=yolov5s.plan后,程序在解析阶段报错:

[E] Error[0]: Node (Mul_123) failed to parse.
[E] Error[1]: Assertion failed: inputs.size() == 2

分析过程:错误指向Mul节点,且断言输入数量应为2。检查ONNX模型发现,该节点是YOLOv5中常见的scale操作,其输入为tensor * scale_factor。问题在于PyTorch导出时,scale_factor被表示为一个常量张量,但ONNX Parser在解析时将其视为一个动态输入,导致输入数量不匹配。

解决步骤:

  1. 使用onnx-simplifier对模型进行简化,将常量折叠进权重:python -m onnxsim yolov5s.onnx yolov5s_sim.onnx。
  2. 重新执行trtexec --onnx=yolov5s_sim.onnx --saveEngine=yolov5s.plan,解析通过。
  3. 进一步添加--fp16参数,引擎构建成功,性能提升约1.8倍。

结果对比:简化前模型解析失败;简化后成功生成引擎,且FP16精度损失在0.3%以内(mAP@0.5)。

复盘经验:ONNX模型中的常量节点、Shape操作、动态Resize等是Parser的常见“雷区”。建议在加载前先用onnx-simplifier或polygraphy进行预处理,能规避大部分解析错误。

4.4 常见误区与总结复盘

在实际使用trtexec加载模型时,工程师们经常陷入以下几个误区:

  • 误区一:忽略输入shape的指定。对于动态batch模型,如果不通过--inputShapes指定具体维度,trtexec会默认使用ONNX中定义的动态维度,导致构建时间剧增且性能下降。正确做法是明确指定实际部署时的shape。
  • 误区二:混淆--deploy和--model的用途。有些用户将--model误用于ONNX文件,导致Parser无法识别。记住:--model只用于Caffe的权重文件。
  • 误区三:每次测试都重新构建引擎。构建引擎是一个耗时操作(大型模型可能需数分钟),且每次构建结果可能因GPU状态略有差异。正确做法是先用--saveEngine保存,后续用--loadEngine加载。
  • 误区四:忽视UFF的废弃状态。新项目如果还在用--uff,建议尽快迁移到ONNX。TensorRT 10.x中UFF Parser已被移除,继续使用会直接报错。

风险提示:当使用--loadEngine加载引擎时,必须确保引擎构建时的TensorRT版本、GPU架构、精度设置与当前运行环境完全一致。跨版本或跨GPU加载会导致无法预测的错误,甚至CUDA illegal memory access。生产环境中建议将引擎文件与TensorRT版本号、GPU型号一同打包记录。

总结来看,trtexec的模型加载参数虽然简单,但背后涉及Parser兼容性、模型预处理、引擎序列化等多个环节。掌握--onnx、--deploy+--model、--uff、--loadEngine的适用边界,理解它们与--saveEngine、--inputShapes的配合关系,是高效使用TensorRT的第一步。建议读者在正式benchmark之前,先花10分钟用trtexec完成一次完整的模型加载与引擎保存,确认无误后再进入性能测试阶段,这样可以避免后续大量无效的调试时间。