ONNX模型外部数据存储与形状推断的注意事项

概述

在使用ONNX模型时，开发者经常会遇到模型体积过大的问题。ONNX提供了将张量数据外部存储的功能，这能有效减小模型文件体积。然而，当模型使用外部数据存储时，可能会影响形状推断(infer_shapes)功能的正常工作。

问题现象

当开发者尝试使用infer_shapes_path方法对包含外部数据的ONNX模型进行形状推断时，发现输出形状信息未能正确推断。具体表现为：

1. 使用内部数据存储时，形状推断正常工作
2. 使用外部数据存储时，输出形状变为未定义状态
3. 使用infer_shapes方法(非路径版本)时，即使有外部数据也能正常工作

技术分析

深入分析ONNX的实现代码可以发现，形状推断功能在设计上存在一个重要的限制条件：它无法处理存储在外部的形状数据。这是因为：

1. 形状推断过程需要实时访问张量数据来计算输出形状
2. 当数据被外部化存储后，形状推断器无法在推断过程中动态加载这些数据
3. 特别是对于形状数据(通常是包含4个或更少元素的数组)，外部化存储会导致关键信息丢失

解决方案

针对这一问题，建议开发者采取以下最佳实践：

1. 合理设置外部化阈值：在调用convert_model_to_external_data时，应该设置一个合理的size_threshold参数，确保小型张量(特别是形状数据)保留在模型内部

2. 优先使用内部存储：对于关键形状数据，考虑强制保留在模型内部，即使它们的大小超过了外部化阈值

3. 选择适当的推断方法：对于小型模型，可以直接使用infer_shapes方法；对于大型模型，确保在外部化前完成形状推断

实现建议

以下是一个改进后的代码示例，展示了如何安全地处理外部数据和形状推断：

# 创建模型时保留小型张量
convert_model_to_external_data(
    onnx_model,
    size_threshold=1024,  # 设置合理的阈值
    location=data_path
)

# 先进行形状推断再外部化大型数据
if model_size < THRESHOLD:
    inferred_model = onnx.shape_inference.infer_shapes(raw_model)
else:
    # 确保形状相关张量保留在内部
    inferred_model = carefully_infer_with_external_data(raw_model)

总结

ONNX模型的外部数据存储是一个强大的功能，但在使用时需要考虑其对模型处理工具链的影响。特别是形状推断这种需要访问张量数据的操作，开发者应该特别注意数据外部化的范围和程度。通过合理设置外部化阈值和正确处理流程，可以既享受外部存储带来的好处，又不影响模型的后续处理。