什么是ONNX
onnx定义了一种可扩展的计算图模型\一系列内置的运算(op)和标准数据类型.每一个计算流图都定义为由节点组成的列表,并构建有向无环图.其中,每一个节点都有一个或者多个输入与输出,每一个节点都称为一个一个OP.这相当于一种通用的计算图,不同深度学习框架构建的计算图都能转换成它.
Caffe2, PyTorch, Microsoft Cognitive Toolkit, Apache MXNet 和其他工具都在对ONNX进行支持。在不同的框架之间实现互操作性,并简化从研究到产品化的过程,将提高人工智能社区的创新速度。
torch.onnx转出onnx模型
本文将介绍pytorch中自带的torch.onnx模块,该模块包含将pytorch模型导出到onnx格式. 以alexnet为pytorch转onnx为例子
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.onnx", verbose=True)
导出的alexnet.onnx包含模型的网络结构和参数.
torch.onnx.export参数的详细介绍
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None)
- model:要被导出的模型
- args(参数的集合):模型的输入,例如,这种model(*args)方式是对模型的有效调用。任何非Variable参数都将硬编码到导出的模型中;任何Variable参数都将成为导出的模型的输入,并按照他们在args中出现的顺序输入。如果args是一个Variable,这等价于用包含这个Variable的1-ary元组调用它。(注意:现在不支持向模型传递关键字参数。)
-f:一个类文件的对象(必须实现文件描述符的返回)或一个包含文件名字符串。一个二进制Protobuf将会写入这个文件中。 - export_params(bool,default True)-如果指定,所有参数都会被导出。如果你只想导出一个未训练的模型,就将此参数设置为False。在这种情况下,导出的模型将首先把所有parameters作为参arguments,顺序由model.state_dict().values()指定。
- verbose(bool,default False):如果指定,将会输出被导出的轨迹的调试描述。
- training(bool,default False):导出训练模型下的模型。目前,ONNX只面向推断模型的导出,所以一般不需要将该项设置为True。
- input_names(list of strings, default empty list):按顺序分配名称到图中的输入节点。
- output_names(list of strings, default empty list):按顺序分配名称到图中的输出节点。
|