PyTorch学习总结(三)——ONNX

最新推荐文章于 2025-06-02 09:01:05 发布

manong_wxd

最新推荐文章于 2025-06-02 09:01:05 发布

阅读量5.7w

点赞数 13

CC 4.0 BY-SA版权

分类专栏： PyTorch

本文链接：https://blog.youkuaiyun.com/manong_wxd/article/details/78720236

本文介绍了ONNX的用途，它是AI模型的一种开源格式，促进不同框架间的互操作性。重点讲解了PyTorch中torch.onnx模块如何将模型导出到ONNX，并提供了AlexNet从PyTorch转到Caffe2的示例。讨论了导出的局限性，如基于跟踪的exporter对动态模型的支持不足，以及ONNX支持的操作符列表。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

这里写图片描述

1.什么是ONNX

Open Neural Network Exchange (ONNX)是开放生态系统的第一步，它使人工智能开发人员可以在项目的发展过程中选择合适的工具;ONNX为AI models提供了一种开源格式。它定义了一个可以扩展的计算图模型，同时也定义了内置操作符和标准数据类型。最初我们关注的是推理(评估)所需的能力。

Caffe2, PyTorch, Microsoft Cognitive Toolkit, Apache MXNet 和其他工具都在对ONNX进行支持。在不同的框架之间实现互操作性，并简化从研究到产品化的过程，将提高人工智能社区的创新速度。

2.torch.onnx

本文我们将主要介绍PyTorch中自带的torch.onnx模块。该模块包含将模型导出到ONNX IR格式的函数。这些模型可以被ONNX库加载，然后将它们转换成可在其他深度学习框架上运行的模型。

3 End-to-end AlexNet from PyTorch to Caffe2

这里有一个简单的脚本，它将torchvision中预训练的AlexNet导出为ONNX。它运行一个简单的推断，然后将生成的跟踪模型保存到alexnet.proto:

from torch.autograd import Variable
import torch.onnx
import torchvision

dummy_input = Variable(torch.randn(10, 3, 224, 224)).cuda()
model = torchvision.models.alexnet(pretrained=True).cuda()
torch.onnx.export(model, dummy_input, "alexnet.proto", verbose=True)

alexnet.proto是一个二进制的protobuf文件，它包含您导出的模型的网络结构和参数(在这里，模型是AlexNet)。关键参数verbose=True使exporter可以打印出一种人类可读的网络表示：

# All parameters are encoded explicitly as inputs.  By convention,
# learned parameters (ala nn.Module.state_dict) are first, and the
# actual inputs are last.
graph(%1 : Float(64, 3, 11, 11)
      %2 : Fl