编辑ONNX模型

1 分钟阅读

概述

onnx模型编辑方法有2种:

虽然onnx-modifier有图形界面编辑onnx,但是实际用起来问题特别多。onnx-graphsurgeon用起来更加强大,简洁。

安装方法如下:

pip install onnx_graphsurgeon --index-url https://pypi.ngc.nvidia.com

也可以源码编译安装:

make build
pip install onnx_graphsurgeon/dist/onnx_graphsurgeon-*-py2.py3-none-any.whl

基本操作

使用ipython命令进行交互式操作

import onnx_graphsurgeon as gs
import onnx
import numpy as np
graph = gs.import_onnx(onnx.load("test.onnx"))
graph #查看所有信息
graph.inputs #查看所有输入信息
graph.nodes #查看所有输出信息
graph.nodes[0] #查看node0的信息
onnx.save(gs.export_onnx(graph), "modified.onnx") # 保存onnx模型

新增节点

# 新增节点,对input做transpose
trans_out = gs.Variable("trans_out",dtype=np.float32, shape=(4,6,5,7))
trans_op=gs.Node(op="Transpose", inputs=[graph.inputs[0]], outputs=[trans_out], attrs={"perm":[0,2,1,3]})
# node0的输入改为trans_out
graph.nodes[0].inputs[0]=trans_out
# 插入到nodes中
graph.nodes.insert(0,trans_op)
# 修改input shape
graph.inputs[0].shape=[4,5,6,7]