1. 问题背景与现象描述
最近在将一个语音识别模型的encoder部分从ONNX格式转换为TensorRT引擎时，遇到了一个棘手的动态shape问题。模型在PyTorch训练和ONNX导出阶段都能正常工作，但在使用trtexec工具转换为TensorRT格式时，出现了以下关键错误：
code复制shape tensor must have build-time extent
这个错误发生在包含Range算子的节点处。初步分析表明，ONNX模型中的Range算子输出的tensor长度依赖于前一个Cast算子的输出值，而TensorRT要求所有shape tensor（用于决定其他tensor形状的tensor）在构建时（build-time）就必须知道其具体维度值。
2. 问题根因分析
2.1 TensorRT对动态shape的限制
TensorRT虽然支持动态shape，但其对"动态"的定义有严格限制：
- 动态batch size和动态维度是支持的
- 但所有决定张量形状的计算必须只依赖于输入的形状，而不能依赖于输入的具体值
在我们的案例中，Range算子的limit参数依赖于前一个算子的输出值，这违反了TensorRT的shape计算原则。具体表现为：
code复制/ReduceMax_output_0: () int64 13
/Range_output_0: (13,) int64 [ 0 1 2 3 4 5 6 7 8 9 10 11 12]
这里Range输出的长度13是由前序计算得到的值，而不是预定义的shape信息。
2.2 模型结构可视化分析
使用Netron.app可视化模型，发现问题的计算路径如下：
code复制x_lens → Cast → ReduceMax → Range → Unsqueeze → ...
关键问题节点：
x_lens：输入序列长度，动态一维tensorReduceMax：计算得到最大长度（如13）Range：生成0到12的序列，长度为13
这种数据依赖的形状计算正是TensorRT所禁止的。
3. 解决方案探索
3.1 初步尝试：用TopK替代ReduceMax
考虑到TopK算子可以输出shape tensor，尝试用TopK(k=1)替代ReduceMax：
python复制# 修改后的节点结构
x_lens → Cast → TopK(k=1) → Gather → Range → ...
但转换时出现新错误：
code复制Error[9]: ITopKLayer cannot be used to compute a shape tensor
原因是TopK的输出依然被视为数据tensor而非shape tensor。
3.2 根本解决：重构输入结构
分析模型整体结构后发现，x_lens经过一系列计算后最终生成一个mask tensor。于是决定直接将该mask作为模型输入，避免内部的动态shape计算：
python复制def solve_replace_dynamic_range_with_mask_input(onnx_path):
model = onnx.load(onnx_path)
graph = gs.import_onnx(model)
# 创建新的mask输入
x_len_mask = gs.Variable(name="x_len_mask", dtype=np.bool_, shape=['N', 'L'])
# 替换原有计算路径
for node in graph.nodes:
for i, input_t in enumerate(node.inputs):
if input_t.name == "/GreaterOrEqual_output_0":
node.inputs[i] = x_len_mask
# 清理旧输入输出
graph.inputs = [inp for inp in graph.inputs if inp.name != "x_lens"]
graph.outputs = [out for out in graph.outputs if out.name != "encoder_out_lens"]
graph.inputs.append(x_len_mask)
graph.cleanup()
graph.toposort()
onnx.save(gs.export_onnx(graph), "encoder_mask_input_solved.onnx")
3.3 新问题：Slice算子兼容性
修改输入后，转换过程变长，但出现了新错误：
code复制Error Code 2: /encoder/Slice_2 requires bool I/O but node can not be handled by Myelin. Dynamic shapes are not equal for slice params.
这是因为Slice算子的输入现在是动态mask。解决方案是用Range+Gather组合替代Slice：
- 用Shape算子获取动态长度
- 用Gather提取需要的部分
- 避免使用INF等TensorRT不支持的值
4. 环境升级与最终方案
4.1 TensorRT版本问题
在TensorRT 8.x环境下，上述修改后仍出现各种兼容性问题。升级到TensorRT 10.0 + CUDA 12.4后，大部分动态shape问题得到解决，但暴露出新的条件分支问题：
code复制IConditionalOutputLayer /encoder/1/encoder/encoder_pos/If_OutputLayer: dimensions not compatible for if-conditional outputs
4.2 条件分支输出不一致问题
模型中的If算子存在分支输出shape不一致：
- then_branch: 固定size的tensor
- else_branch: 动态size的tensor
可能的解决方案：
- 预设最大长度，两个分支都padding到该长度
- 使用mask标识有效部分
- 修改模型结构避免这种不一致
但由于改动较大，暂时保留此问题待后续优化。
5. 完整转换流程与参数
5.1 最终转换命令
bash复制trtexec \
--onnx=./onnx/encoder.onnx \
--saveEngine=./tensorrt/encoder.engine \
--minShapes=x:1x16x80,x_len_mask:1x4 \
--optShapes=x:8x512x80,x_len_mask:8x252 \
--maxShapes=x:16x384000x80,x_len_mask:16x191996 \
--verbose
5.2 关键参数说明
- minShapes: 最小输入shape，用于确定下限
- optShapes: 优化shape，引擎会针对此shape进行优化
- maxShapes: 最大输入shape，决定内存分配上限
- fp16/int8: 可添加--fp16或--int8启用精度优化
6. 经验总结与避坑指南
6.1 动态shape处理原则
- Shape计算纯静态化：所有shape计算只能基于输入shape，不能基于输入值
- 避免数据依赖的形状：如Range(length=data)这种模式要重构
- 统一分支输出：If/loop等控制流的各分支输出shape必须一致
6.2 实用调试技巧
- 使用polygraphy检查shape：
bash复制polygraphy run --onnxrt model.onnx --input-shapes 'x:[1,16,80]' 'x_len_mask:[1,8]'
- 可视化工具组合：
- Netron.app：查看模型结构
- ONNX Runtime：验证模型正确性
- TensorRT的verbose日志：定位错误层
- 渐进式修改：
- 每次只修改一个问题点
- 保留中间各版本模型
- 使用git管理修改历史
6.3 性能考量
-
动态shape会降低性能：相比固定shape，动态shape会：
- 增加引擎构建时间
- 可能阻止某些图优化
- 增加运行时内存需求
-
合理设置shape范围：
- min/opt/max shapes不宜相差过大
- 考虑实际应用场景的合理范围
7. 替代方案探讨
如果动态shape问题难以解决，可以考虑：
-
固定shape+padding：
- 设定最大输入长度
- 短输入padding到该长度
- 输出后mask掉padding部分
-
模型分片：
- 按输入长度分多个区间
- 每个区间导出单独引擎
- 运行时根据长度选择引擎
-
使用ONNX Runtime：
- 保留ONNX格式
- 利用ONNX Runtime的TensorRT EP
- 牺牲部分性能换取兼容性
这次转换过程虽然遇到了诸多挑战，但也积累了宝贵的动态shape处理经验。后续计划尝试TensorRT 10.0的新特性，并考虑重构模型结构从根本上避免动态shape问题。