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

本文深入解析了PyTorch模型通过torch.export迈向ONNX跨平台部署的关键路径与实战陷阱：明确指出torch.export本身不生成ONNX，而是产出中间格式ExportedProgram，必须借助torch.onnx.dynamo_export（PyTorch 2.3+官方推荐）或第三方工具完成二次转换；详述了成功导出的硬性前提——模型需可追踪、无副作用，输入严格限定为Tensor/tuple/dict，动态尺寸必须显式声明，并直击常见报错根源（如.item()滥用、非张量输入、shape推断失败）；更进一步提醒读者，导出完成只是起点，还需通过ONNX校验、ORT推理比对及目标平台适配（如OPSET降级、BN融合）确保真正可用——这是一份兼顾原理透彻性与工程落地性的ONNX迁移必读指南。

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 的语义检查。

以上就是本文的全部内容了，是否有顺利帮助你解决问题？若是能给你带来学习上的帮助，请大家多多支持golang学习网！更多关于文章的相关知识，也可关注golang学习网公众号。