PyTorch转ONNX格式，torch.export实现跨平台部署

本文详解了如何利用 PyTorch 2.3+ 的 `torch.export` 与 `torch.onnx.dynamo_export` 实现安全、可靠的模型跨平台部署，明确指出 `torch.export` 本身并不直接生成 ONNX 文件，而是产出中间格式 `ExportedProgram`，必须经由 `dynamo_export` 或第三方工具二次转换；文章深入剖析了导出成功的核心前提——模型需可追踪、无副作用、输入严格限定为 Tensor/tuple/dict、动态尺寸须显式声明，并系统梳理了常见报错（如不支持算子、输入类型错误、动态形状推断失败）的根因与修复方案，同时强调导出后务必通过 ONNX 校验、Runtime 推理比对及目标平台适配（如 OPSET 降级、BN 融合）三步验证，才能真正保障 ONNX 模型在 TensorRT、Core ML 等异构环境中的可用性。

torch.export 保存 ONNX 的前提条件是什么

torch.export 不是导出 ONNX 的函数，它导出的是 PyTorch 的原生 ExportedProgram。想得到 ONNX 文件，必须再调用 torch.onnx.dynamo_export 或第三方工具（如 onnxscript）做二次转换。直接调用 torch.export 后得到的 ExportedProgram 对象不能被 ONNX 运行时加载。

• torch.export 要求模型必须是「可追踪」且「无副作用」的：不能有 print()、time.sleep()、random.choice() 等运行时依赖
• 输入必须是 torch.Tensor 或嵌套的 tuple/dict，不能含 Python 基本类型（如 int、str）作为参数传入前向逻辑
• 模型中若使用了 torch.nn.functional.interpolate 等动态尺寸操作，需显式提供 dynamic_shapes，否则导出会失败或生成静态 shape 的 ONNX

torch.onnx.dynamo_export 是当前最稳的 ONNX 导出路径

PyTorch 2.3+ 推荐用 torch.onnx.dynamo_export 替代旧的 torch.onnx.export，它基于 TorchDynamo，能更好处理带控制流、高阶函数和自定义算子的模型。

• 支持 torch.compile 后的模型，也支持原始 nn.Module
• 必须传入示例输入（args），且该输入 shape 和 dtype 要覆盖所有可能分支——比如 if 分支里用了不同 size 的 tensor，就得选一个能触发该分支的输入
• 默认不导出权重，只导出计算图；如需权重内联，加参数 export_options=ExportOptions(dynamic_shapes=True)
• 示例：

  import torch
  import torch.onnx
  

model = MyModel().eval() x = torch.randn(1, 3, 224, 224) onnx_program = torch.onnx.dynamo_export(model, x) onnx_program.save("model.onnx")

常见报错：Unsupported op 或 Missing input shape

这类错误基本都源于模型写法越界或导出配置不匹配：

• Unsupported op: aten._local_scalar_dense.default：说明用了 .item() 或 float(tensor) 把 tensor 转成 Python 标量，必须改用 tensor.squeeze().detach().cpu().numpy() 等后处理方式，别在 forward 里做
• Input is not a tensor or tuple of tensors：检查传给 dynamo_export 的 args 是否包含 None、list 或 str ——只接受 Tensor、tuple、dict
• Dynamic shape inference failed：模型里用了 x.shape[0] 做索引或循环次数，但没声明该维度为 dynamic；需用 torch.export.Dim 显式标记，例如：

  batch = torch.export.Dim("batch")
  dynamic_shapes = {"x": {0: batch}}
  torch.onnx.dynamo_export(model, x, dynamic_shapes=dynamic_shapes)
  

导出后 ONNX 文件跑不通？先验证这三件事

ONNX 导出成功不等于能跑通，尤其跨框架部署时：

• 用 onnx.checker.check_model 验证文件结构是否合法，很多“静默导出成功但实际损坏”的情况靠这个能提前发现
• 用 onnxruntime.InferenceSession 加载并跑一次推理，看输出 shape 和数值是否与 PyTorch 一致；注意 ort_session.run 返回的是 list，不是 dict，键名不保留
• 如果目标平台是 TensorRT 或 Core ML，别直接部署原始 ONNX：它们对 OP set 版本敏感，建议用 onnx.version_converter 降级到 OPSET 15 或 16，并手动 fuse BatchNorm

PyTorch 的 ONNX 导出链路现在很依赖 dynamo 的表达能力，而 dynamo 对 control flow 和自定义 C++ 扩展的支持仍有限。如果你的模型里有大量 for 循环或 torch.vmap，导出前最好先用 torch.compile 跑一遍，确认它能正常编译——导不出，大概率是模型本身没过 dynamo 的语义检查。

今天关于《PyTorch转ONNX格式，torch.export实现跨平台部署》的内容就介绍到这里了，是不是学起来一目了然！想要了解更多关于的内容请关注golang学习网公众号！