pytoch模型转成onnx

onnx定义了一种可扩展的计算图模型一系列内置的运算(op)和标准数据类型.每一个计算流图都定义为由节点组成的列表,并构建有向无环图.其中,每一个节点都有一个或者多个输入与输出,每一个节点都称为一个一个OP.这相当于一种通用的计算图,不同深度学习框架构建的计算图都能转换成它.

Caffe2, PyTorch, Microsoft Cognitive Toolkit, Apache MXNet 和其他工具都在对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):按顺序分配名称到图中的输出节点。