【ONNX】各深度学习框架的模型转ONNX

注：大多框架的模型(pytorch、caffe2、mxnet)在加载的时候(如果有大佬知道)，都需要知道输入的shape，caffe2甚至需要输入的name(caffe2只是在转onnx时需要知道input_name和input_size，如果哪位大佬知道如何在caffe2模型中获取input name或Input size可以告诉我一下)，tensorflow需要知道输出的name。。。。在这方面还是cntk和caffe好

pytorch

pytorch安装

Linux和Windows现都已支持Stable(1.0)版本
官网安装
linux cpu安装:

conda install pytorch-cpu torchvision-cpu -c pytorch

版本查询：

import torch
print(torch.__version__)

pytorch转onnx

pytorch在导入模型时，需要有定义模型类的文件(.py格式)

python pytorch2onnx.py ptmodel_path ptmodel_class_py_path insize_n insize_c insize_w insize_h saved_onnx_path_name

ptmodel_path:模型路径，pth、pkl、pt格式不限，pytorch保存模型的时候选择的是保存整个模型而不是只保存网络训练参数。
ptmodel_class_py_path:模型类定义文件的路径(可以不放在当前路径下)，py文件，为pytorch定义模型的类的文件
insize_n、insize_c、insize_w、insize_h:输入的维度
saved_onnx_path_name:保存为onnx模型的路径+文件名(没有路径会保存在当前路径下)

关于pytorch模型的题外话

感觉pytorch保存模型实在是太麻烦了、就算把网络结构和训练参数都保存了，依然需要在导入的时候在某处存在这个模型类的定义。对pytorch的探索可能还是太浅，现在暂且只能做到把这个类定义文件拷贝到当前路径下，用完再删除。至于只保存训练参数的情况就没写了，可以写，应该需要通过这个py文件导入这个模型类，实例化一个模型，再加载训练参数。
还有一点就是必须手动给出网络输入的维度。
我觉得pytorch以后保存模型会直接保存为onnx(官方提供了直接保存onnx的接口)，而不是在外部将pytorch模型转为onnx模型。有了onnx模型，pytorch在部署上就会更加轻松。因为onnx模型就像是tensorflow的pb模型一样，是个固化模型结构。pytorch本身的模型和tensorflow的ckpt模型类似，在重新加载模型和训练上灵活性很强。

cntk

cntk安装

官网安装
linux python3.6 cpu 安装

pip install https://cntk.ai/PythonWheel/CPU-Only/cntk-2.6-cp36-cp36m-linux_x86_64.whl

版本查询

import cntk
print(cntk.__version__)

cntk转onnx

cntk比较简单

python cntk2onnx.py cntk_model_path saved_onnx_path_name

cntk_model_path:cntk模型路径
saved_onnx_path_name:保存为onnx模型的路径+文件名(没有路径会保存在当前路径下)

mxnet

mxnet安装

官网安装
linux python安装

pip install mxnet

版本查询

import mxnet
print(mxnet.__version__)

mxnet转onnx

python mxnet2onnx.py symbol_json_path params_path insize_n insize_c insize_w insize_h saved_onnx_path_name

symbol_json_path:描述网络结构的文件路径，json后缀
params_path:存储网络参数的文件路径，param后缀
insize_n、insize_c、insize_w、insize_h:输入的维度
saved_onnx_path_name:保存为onnx模型的路径+文件名(没有路径会保存在当前路径下)

caffe2

caffe2安装

由于pytorch1.0后，caffe2被集成在pytorch里了，因此只用安装pytorch即可

caffe2转onnx

python caffe22onnx.py predict_path init_path insize_n insize_c insize_w insize_h input_name saved_onnx_path_name

predict_path:描述网络结构的文件路径，pb后缀
init_path:存储网络参数的文件路径，pb后缀
insize_n、insize_c、insize_w、insize_h:输入的维度
input_name:输入的名字
saved_onnx_path_name:保存为onnx模型的路径+文件名(没有路径会保存在当前路径下)

tensorflow

tensorflow安装

cpu pip安装

pip install tensorflow

版本查询

import tensorflow
print(tensorflow.__version__)

onnx-tensorflow安装

pip install onnx-tf

tensorflow转onnx

tensorflow在转onnx的时候可能会出现维度问题，而onnx转tensorflow这个工具现在还在实验中。

python tensorflow2onnx.py tensorflow_model_path output_layer_name saved_onnx_path_name

tensorflow_model_path:tensorflow模型路径，pb后缀，不是ckpt模型
output_layer_name:输出的名字
saved_onnx_path_name:保存为onnx模型的路径+文件名(没有路径会保存在当前路径下)

caffe

caffe->caffe2->onnx或caffe->onnx

现在caffe转caffe2支持的层有(通过看caffe_translator.py可知)：

Input
VideoData
Data
Convolution3D
Convolution
Deconvolution
Crop
ReLU
Pooling
Pooling3D
LRN
InnerProduct
Dropout
Softmax
SoftmaxWithLoss
Accuracy
Concat
TanH
InstanceNorm
BatchNorm
Eltwise
Scale
Reshape
Sigmoid
ROIPooling
PReLU
Reduction

caffe->onnx(自己编写的)支持的算子

Convolution
Concat
Dropout
InnerProduct(Reshape+MatMul+Add)
LRN
Pooling
ReLU
Softmax
Eltwise
Upsample
BatchNorm
Scale
PRelu

caffe->caffe2

caffe_translator.py

python -m caffe2.python.caffe_translator deploy.prototxt pretrained.caffemodel

caffe->caffe2->onnx比较麻烦的是，在caffe2->onnx的时候需要知道输入的shape和name，所以最好提前在caffe的prototxt文件中先把这些信息获取了。