4. 模型格式解析
在地平线模型转换流程中,第一步并不是直接调用转换工具,而是先回答一个更基础的问题:手里的模型到底是什么格式,它能不能被工具链正确读取?很多转换失败、精度异常、算子不支持的问题,根源都在模型格式与导出规范上。本节围绕常见深度学习模型格式、ONNX 的中间层作用、模型导出规范以及兼容性要求展开,帮助你在转换前把好第一道关。
一、常见深度学习模型格式全景
深度学习模型格式可以按“训练框架原生格式”和“通用交换格式”两大类理解。训练框架原生格式与框架强绑定,通常包含计算图、权重、优化器状态甚至训练配置;通用交换格式则用于跨框架、跨工具链传递模型。
| 格式 | 来源/框架 | 主要特点 | 在地平线转换中的角色 |
|---|---|---|---|
| PyTorch (.pt/.pth) | PyTorch | 动态图为主,保存 state_dict 或完整模型,灵活但结构不固定 | 通常需先导出为 ONNX |
| TensorFlow SavedModel | TensorFlow | 包含计算图与变量,目录结构固定,适合部署 | 可经 tf2onnx 转 ONNX |
| Keras (.h5/.keras) | Keras | 层结构清晰,权重与结构可分离保存 | 可转 ONNX 或 SavedModel |
| Caffe (.prototxt + .caffemodel) | Caffe | 结构文本与权重分离,早期部署常用 | 部分工具链可直接读取 |
| ONNX (.onnx) | 开放标准 | 跨框架计算图 + 权重,算子集版本化 | 地平线工具链的核心输入格式 |
| TFLite (.tflite) | TensorFlow Lite | 面向移动端,量化友好,算子受限 | 一般需先转回 ONNX 或重新导出 |
可以看到,ONNX 处在“训练框架”与“部署工具链”之间的枢纽位置。地平线模型转换工具链通常以 ONNX 作为标准输入,因此理解 ONNX 的作用,是掌握整个转换流程的关键。
二、ONNX 的作用:为什么它是转换流程的“通用语言”
ONNX(Open Neural Network Exchange)是一种开放的模型表示标准,它用统一的计算图描述模型:节点表示算子,边表示张量,权重以初始值形式附着在图中。它的核心价值在于解耦——训练框架负责产出 ONNX,部署工具链负责消费 ONNX,双方不必互相适配。
ONNX 模型内部包含几个关键要素:
- Graph(计算图):由 Node、Input、Output、Initializer 组成,描述数据从输入到输出的完整流向。
- Opset(算子集版本):决定算子语义与可用范围,例如 opset 11、opset 13、opset 17。版本越高,支持的算子越丰富,但工具链兼容性可能下降。
- IR Version(中间表示版本):ONNX 自身的格式版本,与 opset 相互配合。
- Metadata:可携带作者、框架、输入输出名称等辅助信息。
三、模型导出规范:从训练框架到 ONNX
导出 ONNX 的质量直接决定后续转换的成败。以下以 PyTorch 为例,说明导出时的关键规范。
import torch
import torch.onnx
model = MyModel()
model.load_state_dict(torch.load("model.pth"))
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model,
dummy_input,
"model.onnx",
export_params=True,
opset_version=11,
do_constant_folding=True,
input_names=["input"],
output_names=["output"],
dynamic_axes=None
)
导出时需要注意以下规范:
- 固定输入维度:地平线工具链通常要求静态 shape,导出时应避免 dynamic_axes,或明确固定 batch、height、width。
- 使用 eval 模式:必须调用 model.eval(),否则 BatchNorm、Dropout 等层会带入训练态行为,导致推理结果异常。
- 选择合适的 opset:建议优先使用工具链明确支持的 opset 版本,常见为 11 或 13,不要盲目追高。
- 命名规范:input_names、output_names 应清晰可读,便于后续在转换配置中引用。
- 常量折叠:开启 do_constant_folding 可减少冗余节点,但需确认折叠后不引入不支持的算子。
- 导出后校验:使用 onnx.checker 检查模型合法性,并用 onnxruntime 做一次推理比对。
python -m onnx.checker model.onnx
python -m onnxruntime.tools.onnx_model_compare model.onnx
四、兼容性要求与常见误区
兼容性问题通常集中在算子、版本和结构三个层面。下面这张关系图展示了从训练框架到地平线工具链的格式流转与兼容性检查点。
结合上图,兼容性要求可以归纳为三条主线:
- 算子兼容:ONNX 中的算子必须在地平线工具链支持列表内。遇到不支持的算子,需要在导出前替换为等价结构,或使用工具链提供的自定义算子映射。
- 版本兼容:opset 与 IR version 需要落在工具链支持区间。版本过高可能无法解析,版本过低可能缺少必要算子语义。
- 结构兼容:输入输出必须为静态 shape,避免动态 batch 或动态分辨率;同时要确保没有训练专用节点残留。
常见误区二:忽略 opset 版本。同一个模型用 opset 17 导出可能包含工具链不支持的算子,改用 opset 11 反而更顺利。
常见误区三:导出时忘记 model.eval(),导致推理结果与训练精度对不上,却误以为是量化问题。
五、案例:一次典型的格式兼容排查
某次转换任务中,PyTorch 模型导出 ONNX 后,工具链报出“Unsupported operator: GridSample”。排查过程如下:
- 用 onnx.checker 确认模型本身合法,排除文件损坏。
- 用 Netron 打开 ONNX,定位到 GridSample 节点,确认其来自模型中的可变形卷积模块。
- 查阅工具链算子支持列表,确认当前版本不支持 GridSample。
- 将可变形卷积替换为普通卷积加插值近似,重新导出 ONNX,opset 保持 11。
- 再次转换,算子检查通过,后续量化与编译顺利完成。
这个案例说明:格式解析不是“看一眼后缀名”,而是要深入到计算图内部,逐节点确认算子、版本与结构是否符合工具链要求。把这一步做扎实,后面的量化与编译才能稳定推进。