模型转换之【MXNet 转 ONNX】—— 笔记

MXNet 安装

• MXNet 主页安装教程

nvcc --version 查看自己所使用的 Cuda 版本

nvcc -V

nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2020 NVIDIA Corporation
Built on Mon_Nov_30_19:08:53_PST_2020
Cuda compilation tools, release 11.2, V11.2.67
Build cuda_11.2.r11.2/compiler.29373293_0

因此我的安装命令如下

 pip install mxnet-cu110
 # 或者
 pip install mxnet
 
pip install onnx

pip install onnxruntime

MXNet 转 ONNX

官方教程链接如下：

• MXNet Exporting to ONNX
  

本博文做简单整理

import mxnet as mx
import numpy as np
from mxnet.contrib import onnx as onnx_mxnet


import logging
logging.basicConfig(level=logging.INFO)

# Download pre-trained resnet model - json and params by running following code.
# 这里也可以自己合并链接到浏览器手动下载 | 默认会下载到当前目录
path='http://data.mxnet.io/models/imagenet/'
[mx.test_utils.download(path+'resnet/18-layers/resnet-18-0000.params'),
 mx.test_utils.download(path+'resnet/18-layers/resnet-18-symbol.json'),
 mx.test_utils.download(path+'synset.txt')]


 # Downloaded input symbol and params files
sym = './resnet-18-symbol.json'
params = './resnet-18-0000.params'

# Standard Imagenet input - 3 channels, 224*224
input_shape = (1,3,224,224)

# Path of the output file
onnx_file = './mxnet_exported_resnet50.onnx'

# 模型转换已经封装很好了，一行命令即可
converted_model_path = onnx_mxnet.export_model(sym, params, [input_shape], np.float32, onnx_file)


from onnx import checker
import onnx


# Check validity of ONNX model 检查导出 onnx 的可用性
# Load onnx model
model_proto = onnx.load_model(converted_model_path)

# Check if converted ONNX protobuf is valid
checker.check_graph(model_proto.graph)

简单分析

MXNet 框架目前热度没那么高，不过官方的文档非常的清晰易懂，建议有兴趣自行查阅【好像有不少404】