【TensorRT 多输入】trtexec 如何对多输入的 ONNX 模型进行序列化（转化）？

在将包含多个输入的 ONNX 模型转换为 TensorRT 引擎时，可以使用 trtexec 工具。以下是基于讨论中的总结步骤：

1. 基本用法

要使用 trtexec 转换多输入的 ONNX 模型，需要使用 --shapes 参数指定每个输入的形状。形状可以是静态的或动态的。例如：

trtexec --onnx=model.onnx --saveEngine=model.trt --shapes=input1:shape1,input2:shape2,...

2. 动态形状

对于动态形状，可以使用 --minShapes、--optShapes 和 --maxShapes 参数来指定输入的最小、最优和最大形状。例如：

trtexec --onnx=model.onnx --saveEngine=model.trt --minShapes=input1:min_shape1,input2:min_shape2,... --optShapes=input1:opt_shape1,input2:opt_shape2,... --maxShapes=input1:max_shape1,input2:max_shape2,...

动态维度通常用 -1 表示。

3. 具体示例

下面是一个包含动态输入的实际示例命令：

trtexec --onnx=encoder-temp_rnnt.onnx \
--saveEngine=encoder-trtexec.trt --minShapes=audio_signal:[1,80,1],length:[1] \
--optShapes=audio_signal:[1,80,1000],length:[1] \
--maxShapes=audio_signal:[1,80,10000],length:[1]

对于另一个具有多输入的模型：

trtexec --onnx=decoder_joint-temp_rnnt.onnx \
--saveEngine=decoder-trtexec.trt \
--shapes=encoder_outputs:[1,1024,1],targets:[1,1],target_length:[1],input-states-1:[2,1,640],input-states-2:[2,1,640]

个人的写法如下，可供参考，在 log 中可以详细地看到问题、或者模型转化的经过：

trtexec --onnx=model_1.onnx \
        --saveEngine=model_1.plan \
        --minShapes=pixel_seq:4x1,xy_seq:4x1x2,text:4x50 \
        --optShapes=pixel_seq:8x360,xy_seq:8x360x2,text:8x50 \
        --maxShapes=pixel_seq:16x562,xy_seq:16x562x2,text:16x50 \
        --verbose \
> trtexec-1.log 2>&1

4. 错误处理

如果遇到错误 stoi（StringToInteger Error），可能是因为形状格式不正确。例如，将 [1,1024,1] 改为 1x1024x1：

trtexec --onnx=decoder_joint-temp_rnnt.onnx \
--saveEngine=decoder-trtexec.trt \
--shapes=encoder_outputs:1x1024x1,targets:1x1,target_length:1,input-states-1:2x1x640,input-states-2:2x1x640

同时确保 ONNX 模型的 batch 维度是动态的。

总结

通过正确指定输入的形状和动态范围，可以使用 trtexec 将多输入的 ONNX 模型成功转换为 TensorRT 引擎。在遇到问题时，检查形状格式和模型的批处理维度是解决问题的关键。

参考 TensorRT 官方 issue: https://github.com/NVIDIA/TensorRT/issues/2324