编辑ONNX模型
概述
onnx模型编辑方法有2种:
- onnx-graphsurgeon : 手工编辑onnx模型
- onnx-modifier : 图形化编辑onnx模型
虽然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]