神经网络模型模型转ONNX

2024年4月22日 128次阅读 来源: pytorch

近期由于业务需要，需要将训练好的模型转为ONNX格式，为此颇费了一番功夫，在此总结一下吧。。

1、ONNX是一种神经网络模型保存的中间格式，支持多种格式的模型转为ONNX，也支持使用ONNX导入多种格式的模型，具体见https://github.com/onnx/tutorials;目前其实ONNX对于模型的支持还不是太好，主要表现在一些op还不能够支持；

2、在PyTorch下要将模型保存成ONNX格式需要使用torch.onnx.export()函数，使用该函数的时候需要传入下面参数：

　　–model:待保存的model，也就是你在程序中已经训练好或者初始化好的模型

　　–input_shape:指定输入数据的大小，也就是输入数据的形状，是一个包含输入形状元组的列表；

　　–name:模型的名称，即模型的保存路径；

　　–verbrose:True或者False,用来指定输出模型时是否将模型的结构打印出来；

　　–input_names:输入数据节点的名称，数据类型为包含字符串的列表；一般将这个名称设为[‘data’]；

　　–output_names:输出数据节点的名称，类型与输入数据的节点名称相同；

　　在成功导出模型后，可以使用ONNX再对模型进行检查：

import onnx # Load the ONNX model
model = onnx.load("alexnet.onnx") # Check that the IR is well formed
onnx.checker.check_model(model) # Print a human readable representation of the graph
onnx.helper.printable_graph(model.graph)

目前PyTorch还不支持导入ONNX格式的模型。

3、使用MXNET导出模型为ONNX时，参考地址：https://cwiki.apache.org/confluence/display/MXNET/ONNX，http://mxnet.incubator.apache.org/versions/master/tutorials/onnx/export_mxnet_to_onnx.html。MXNet模型的保存格式为.json文件+.params文件，.json文件里保存的是模型的结构，.params文件中保存的是模型的参数。使用onnx_mxnet.export_export_model()方法就可以实现将模型从mxnet转为ONNX格式，该方法需要传入的参数为：

　　–sym：.json文件，也就是保存了网络结构的文件

　　–params:参数文件

　　–input_shape:输入数据的形状，是一个包含形状元组的列表

　　–input_type:输入数据的类型；

　　–模型的保存路径

4、从MXNet导入ONNX格式模型：需要使用mxnet.contrib.onnx.onnx2mx.import_model.import_model(model_file)，这里返回的是sym, arg_arams,aux_params，也就是网络结构symbol对象，保存参数的字典, 再将其转为MXNet的module对象(使用mxnet.module.Module())，即可将模型恢复到mxnet框架下可执行的模型。

 

最后，好久没有记录日常学习积累的东西了，趁着失眠开个好头吧，晚安。。。

    原文作者：pytorch
    原文地址: https://www.cnblogs.com/puheng/p/10873289.html
    本文转自网络文章，转载此文章仅为分享知识，如有侵权，请联系博主进行删除。
点赞