神经网络模型模型转ONNX

近期由于业务需要,需要将训练好的模型转为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/ONNXhttp://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框架下可执行的模型。

 

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

posted @ 2019-05-16 02:43  Neal_Pu  阅读(6524)  评论(0编辑  收藏  举报