pytoch模型转成onnx

pytoch模型转成onnx,第1张

pytoch模型转成onnx 什么是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):按顺序分配名称到图中的输出节点。

欢迎分享,转载请注明来源:内存溢出

原文地址: http://outofmemory.cn/zaji/5651859.html

(0)
打赏 微信扫一扫 微信扫一扫 支付宝扫一扫 支付宝扫一扫
上一篇 2022-12-16
下一篇 2022-12-16

发表评论

登录后才能评论

评论列表(0条)

保存