TorchScript vs ONNX vs TFLite 对比：三大模型序列化格式的互操作性与部署生态分析

TorchScript vs ONNX vs TFLite 对比：三大模型序列化格式的互操作性与部署生态分析

一、模型格式是推理部署的咽喉要道

训练完成的 PyTorch 模型要部署到边缘设备，第一关就是格式转换。选择错误的中间格式可能导致算子不兼容、精度损失甚至完全无法部署。本文从格式设计理念、互操作性、推理后端生态、量化支持四个维度对比三大主流格式。

二、三种格式的设计哲学

三、互操作性实测

3.1 PyTorch → 三格式导出代码对比

""" PyTorch 模型导出三大格式完整示例 展示每种格式的导出 API、参数配置和常见陷阱 """ import torch import torchvision import numpy as np import os # ============ 加载预训练模型 ============ try: model = torchvision.models.resnet50(weights="IMAGENET1K_V1") model.eval() except Exception as e: print(f"[错误] 模型加载失败: {e}") print("请确认 torchvision 版本 >= 0.15") exit(1) dummy_input = torch.randn(1, 3, 224, 224) # ============ 1. TorchScript 导出 ============ def export_torchscript(model, dummy_input, output_path): """ TorchScript 导出 - 两种模式： - trace: 追踪执行路径，适合无控制流的模型 - script: 解析 Python 源码，支持 if/for，但对 Python 语法有限制 """ os.makedirs(os.path.dirname(output_path), exist_ok=True) # 优先使用 trace（更稳定，覆盖大多数 CNN 模型） try: traced = torch.jit.trace(model, dummy_input, strict=True) except RuntimeError as e: # strict=True 模式下若有数据依赖的控制流会失败 print(f"[警告] strict trace 失败: {e}") print("降级到 non-strict trace（可能遗漏部分计算路径）") traced = torch.jit.trace(model, dummy_input, strict=False) # 验证导出正确性 with torch.no_grad(): original_out = model(dummy_input) traced_out = traced(dummy_input) diff = torch.max(torch.abs(original_out - traced_out)).item() if diff > 1e-5: print(f"[警告] TorchScript 输出差异: {diff:.6f}，" f"可能存在未被追踪的控制流") traced.save(output_path) print(f"[TorchScript] 已保存: {output_path}") return traced # ============ 2. ONNX 导出 ============ def export_onnx(model, dummy_input, output_path): """ ONNX 导出 - 关键参数说明： - opset_version: 操作集版本，越高越好但不一定被所有后端支持 - dynamic_axes: 动态维度配置（batch size 可变等） - input/output_names: 命名对后端调试至关重要 """ try: torch.onnx.export( model, dummy_input, output_path, export_params=True, # 同时保存权重 opset_version=17, # 使用较新 opset，支持更多算子 do_constant_folding=True, # 折叠常量（减少算子数量） input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch_size"}, # batch 维度动态 "output": {0: "batch_size"}, } ) print(f"[ONNX] 已保存: {output_path}") except torch.onnx.ExportError as e: print(f"[错误] ONNX 导出失败: {e}") print("常见原因：模型中包含 ONNX 不支持的算子") print("解决方法：使用 torch.onnx.register_custom_op_symbolic 注册") return None # 使用 onnx 工具验证模型 try: import onnx onnx_model = onnx.load(output_path) onnx.checker.check_model(onnx_model) print("[ONNX] 模型验证通过") # 获取算子统计 op_types = [node.op_type for node in onnx_model.graph.node] from collections import Counter op_count = Counter(op_types) print(f"[ONNX] 算子统计: {dict(op_count)}") except ImportError: print("[警告] 未安装 onnx 包，跳过模型验证") except onnx.checker.ValidationError as e: print(f"[警告] ONNX 模型验证失败: {e}") # ============ 3. TFLite 导出（通过 ONNX 中转）============ def export_tflite_via_onnx(onnx_path, tflite_path): """ PyTorch → ONNX → TFLite 转换链路 注意：这条路存在算子兼容性损失，不是所有模型都能成功转换 """ try: import onnx from onnx_tf.backend import prepare import tensorflow as tf except ImportError: print("[错误] TFLite 导出需要: pip install onnx-tf tensorflow") return None # ONNX → TensorFlow try: onnx_model = onnx.load(onnx_path) tf_rep = prepare(onnx_model) tf_path = onnx_path.replace(".onnx", "_tf") tf_rep.export_graph(tf_path) print(f"[中间] TF SavedModel 已保存: {tf_path}") except Exception as e: print(f"[错误] ONNX→TF 转换失败: {e}") print("可能是算子不兼容，考虑使用 onnx2tf 替代方案") return None # TF → TFLite try: converter = tf.lite.TFLiteConverter.from_saved_model(tf_path) converter.optimizations = [tf.lite.Optimize.DEFAULT] # 默认使用动态范围量化 converter.target_spec.supported_types = [tf.float16] tflite_model = converter.convert() os.makedirs(os.path.dirname(tflite_path), exist_ok=True) with open(tflite_path, "wb") as f: f.write(tflite_model) # 打印模型大小对比 onnx_size = os.path.getsize(onnx_path) / (1024 * 1024) tflite_size = os.path.getsize(tflite_path) / (1024 * 1024) print(f"[TFLite] 已保存: {tflite_path}") print(f"[TFLite] 文件大小: {tflite_size:.1f} MB " f"(ONNX: {onnx_size:.1f} MB, " f"压缩率: {(1-tflite_size/onnx_size)*100:.1f}%)") return tflite_path except Exception as e: print(f"[错误] TFLite 转换失败: {e}") return None # ============ 执行导出 ============ if __name__ == "__main__": output_dir = "./exported_models" os.makedirs(output_dir, exist_ok=True) export_torchscript(model, dummy_input, os.path.join(output_dir, "resnet50.pt")) export_onnx(model, dummy_input, os.path.join(output_dir, "resnet50.onnx")) export_tflite_via_onnx( os.path.join(output_dir, "resnet50.onnx"), os.path.join(output_dir, "resnet50.tflite"))

四、部署后端兼容性矩阵

实测中常见的格式兼容性陷阱：

五、总结

格式选型决策树：

• 纯 PyTorch 生态部署（服务器 GPU / libtorch 边缘）→TorchScript，零转换损耗，完整算子支持，.pt文件直接加载推理。
• 跨框架部署 / 多后端支持→ONNX，一次导出对接 ONNX Runtime、TensorRT、OpenVINO、CoreML 等几乎所有主流后端，是互操作性的最优解。
• 移动端 / 超低功耗 MCU→TFLite，FlatBuffers 的零拷贝加载在内存受限设备上优势明显，量化工具链也最为成熟。
• 如果模型中使用大量自定义算子，建议优先评估 ONNX 的自定义算子注册机制，其次是 TorchScript（直接用 C++ 扩展），TFLite 的自定义算子门槛最高。

当前趋势：PyTorch 2.0 的 torch.compile + ExecuTorch 正在试图提供"从训练到部署一条龙"的体验，未来可能削弱 ONNX 的中间层角色，但 ONNX 的多后端互操作优势短期内无法被替代。