欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

pytorch导出onnx模型支持动态图

程序员文章站 2024-01-19 12:37:16
...

 

torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, aten=False, export_raw_ir=False, operator_export_type=None, opset_version=None, _retain_param_name=True, do_constant_folding=False, example_outputs=None, strip_doc_string=True, dynamic_axes=None, keep_initializers_as_inputs=None)

功能:
将模型以ONNX格式导出并保存.

参数:

  • model (torch.nn.Module) – 要导出的模型.
  • args (tuple of arguments) – 模型的输入, 任何非Tensor参数都将硬编码到导出的模型中;任何Tensor参数都将成为导出的模型的输入,并按照他们在args中出现的顺序输入。因为export运行模型,所以我们需要提供一个输入张量x。只要是正确的类型和大小,其中的值就可以是随机的。请注意,除非指定为动态轴,否则输入尺寸将在导出的ONNX图形中固定为所有输入尺寸。在此示例中,我们使用输入batch_size 1导出模型,但随后dynamic_axes 在torch.onnx.export()。因此,导出的模型将接受大小为[batch_size,3、100、100]的输入,其中batch_size可以是可变的。
  • export_params (bool, default True) – 如果指定为True或默认, 参数也会被导出. 如果你要导出一个没训练过的就设为 False.
  • 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) –按顺序分配名称到图中的输出节点
  • dynamic_axes – {‘input’ : {0 : ‘batch_size’}, ‘output’ : {0 : ‘batch_size’}}) # variable lenght axes
shape(input_1) = ('b', 3, 'w', 'h')
and shape(input_2) = ('b', 4)
and shape(output)  = ('b', 'd', 5)
torch.onnx.export(model, args, f, export_params=True, verbose=False, training=False, input_names=None, output_names=None, dynamic_axes = {‘input_1’:[0, 2, 3], ‘input_2’:[0], ‘output’:[0, 1]})

你好,请问pytorch模型转换onnx时的input_names和output_names都是固定的吗?还是根据网络定义的?如果根据网络定义,如何查看呢?是按print(model)查看还是torchsummary?

答:保持数量对应,名字都是随意命名的

 

参考:[pytorch][转载]torch.onnx.export():将pytorch模型转换为.onnx模型

 

 

 

 

 

 

 

 

相关标签: pytorch