Skip to content

Latest commit

 

History

History
104 lines (79 loc) · 3.26 KB

README.md

File metadata and controls

104 lines (79 loc) · 3.26 KB
Raw

ONNX2X

ONNX作为微软的网络模型中间表示被各个框架广泛应用，包括Pytroch，TensorFlow，OneFlow，Keras，Paddle等多种深度学习训练框架。因此，一直在思考一个问题，一个TensorFlow导出来的ONNX模型是否可以借助ONNX被Pytorch框架使用呢？ONNX的理想是作为所有框架的模型的中间交换，那么我们只需要再实现ONNX到各个框架的逆转就可以完成这件事情了。本工程的目的即是尝试支持ONNX转换到各种训练框架，主要为了锻炼算子对齐和更深入的了解ONNX。

代码结构

- onnx2pytorch onnx转pytorch代码实现
- onnx2pytorch.py onnx转pytorch测试代码
- convert_models.md 转换ONNX Model Zoo里面的模型对应的命令和结果记录
- README.md

运行环境

  • pytorch >= 1.1.0
  • onnx>=1.8.1

  • =1.6.0

  • onnxoptimizer>=0.2.3

使用方法

使用下面的命令将各个训练框架导出的ONNX模型转换成Pytorch模型

python .\onnx2pytorch.py ...

参数列表如下:

  • --onnx_path 字符串，必选参数，代表onnx模型的路径
  • --pytorch_path 字符串，必选参数，代表转换出的Pytorch模型保存路径
  • --simplify_path 字符串，可选参数，代表ONNX模型简化（例如删除Dropout和常量OP）后保存的ONNX模型路径
  • --input_shape 字符串，必选参数，代表ONNX模型的输入数据层的名字和维度信息

使用示例

python .\onnx2pytorch.py --onnx_path .\models\mobilenetv2-7.onnx --simplify_path .\models\mobilenetv2-7-simplify.onnx --pytorch_path .\models\mobilenetv2-7.pth --input_shape input:1,3,224,224

模型转换失败处理方法

  • 将onnx2pytorch.py里面的model = convert.ConvertModel(onnx_model, debug=False)这行代码里面的debug设置False重新运行模型即可定位到转换失败的OP，然后你可以在工程提出issue或者自己解决然后给本工程PR。

ONNX2Pytorch

已支持的ONNX OP

  • Conv
  • BatchNormalization
  • GlobalAvgragePool
  • AvgPool
  • MaxPool
  • BatchNorm
  • Flatten
  • Reshape
  • Relu
  • Add
  • Gemm
  • Sigmoid
  • Mul
  • Concat
  • Resize (还有一些问题需要解决，当前版本支持固定倍数方法)
  • Transpose
  • LRN
  • Clip
  • Pad2d
  • Split
  • ReduceMean
  • LeakyRelu

已验证支持的模型

基于ONNXRuntime和Pytorch推理之后特征值mse小于1e-7，视为转换成功

分类模型

  • zfnet512-9.onnx
  • resnet50-v2-7.onnx
  • mobilenetv2-7.onnx
  • mobilenetv2-1.0.onnx
  • bvlcalexnet-9.onnx
  • googlenet-9.onnx
  • squeezenet1.1-7.onnx
  • shufflenet-v2-10.onnx
  • inception-v1-9.onnx
  • inception-v2-9.onnx
  • vgg19-caffe2-9.onnx
  • rcnn-ilsvrc13-9.onnx

检测模型

  • yolov5s-simple.onnx

分割模型

TODO

  • 支持更多模型
  • 重构工程，并解决某些模型转为Pytorch模型之后Netron可视化看不到某些OP的问题
  • 一些部署工作，比如Keras导出的ONNX转为Pytorch模型后，二次导出ONNX递交给NCNN推理

相关链接