神经网络模型模型转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/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框架下可执行的模型。
最后,好久没有记录日常学习积累的东西了,趁着失眠开个好头吧,晚安。。。