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)
%