ONNX-Simplifier中自定义算子形状信息保留的技术探讨

背景介绍

ONNX-Simplifier是一个用于简化ONNX模型的开源工具，它通过优化计算图和执行形状推断来简化模型结构。然而，在处理包含自定义算子(Custom OP)的模型时，ONNX-Simplifier会遇到形状推断的挑战。

问题核心

当模型包含未在ONNX Runtime中注册的自定义算子时，ONNX-Simplifier无法对这些算子及其下游节点执行形状推断。这是因为形状推断需要详细的C++实现来定义这些自定义算子的行为特性。

技术挑战

1. 形状推断依赖：ONNX-Simplifier依赖形状推断来优化模型，但自定义算子缺乏必要的形状推断实现
2. 用户友好性：要求用户编写C++代码来实现自定义算子的形状推断增加了使用门槛
3. 静态形状处理：在许多情况下，输入形状是静态的，用户可以轻松确定自定义算子的输出形状

现有解决方案

目前，用户可以采用以下两种方法处理这个问题：

1. 后处理方式：先使用ONNX-Simplifier简化模型，然后调用onnx.shape_inference.infer_shapes()进行形状推断
2. C++实现：为自定义算子编写详细的C++形状推断代码

潜在改进方向

从技术实现角度，ONNX-Simplifier可以考虑以下改进：

1. 保留原始值信息：对于无法推断形状的自定义算子节点，保留输入图中的value_info
2. 用户指定形状：提供接口允许用户显式指定自定义算子的输出形状
3. 混合推断策略：对可推断部分执行自动推断，对自定义算子部分保留原始信息

技术实现建议

1. 值信息传递机制：在简化过程中建立原始值信息的保留机制
2. 自定义算子标记：识别并特殊处理包含自定义算子的节点
3. 形状推断流程优化：调整形状推断流程，使其能够处理部分未知的情况

对用户的影响

这种改进将显著降低用户使用门槛，特别是：

1. 减少C++知识要求：用户无需掌握C++即可处理自定义算子
2. 简化工作流程：避免额外的后处理步骤
3. 提高可靠性：对于静态形状场景，结果更加可控

总结

ONNX-Simplifier在处理自定义算子时的形状推断限制是一个值得关注的技术问题。通过改进值信息保留机制，可以在不增加用户负担的情况下提高工具的使用范围和可靠性。这种改进对于推动ONNX生态的发展，特别是支持更多自定义算子场景具有重要意义。