onnx 增刪改查，修改節點，刪除節點，修改input，output

细雨微光發表於2024-07-09

一、onnx 的資料型別，共有16種

elem_type: 1 --> float32
elem_type: 2 --> uint8
elem_type: 3 --> int8
elem_type: 4 --> uint16
elem_type: 5 --> int16
elem_type: 6 --> int32
elem_type: 7 --> int64
elem_type: 8 --> string
elem_type: 9 --> boolean
elem_type: 10 --> float16
elem_type: 11 --> float64
elem_type: 12 --> uint32
elem_type: 14 --> uint64
elem_type: 15 --> complex128
elem_type: 16 --> bfloat16

二、onnx 節點列印

onnx_print.py: python onnx_print.py model.onnx

import onnx
import numpy as np
import sys
ori_file=sys.argv[1]
onnx_model = onnx.load(ori_file)
#print(onnx_model)
#print(onnx_model.output)
graph = onnx_model.graph
print("model input: ")
print(graph.input)
for input_node in onnx_model.graph.input:
    print("input data name",input_node.name)
print("\nmodel output: ")
print(graph.output)
for output_node in onnx_model.graph.output:
    print("output data name",output_node.name)
print("\nall model graph")
print("all_node_number: " + str(len(graph.node))) #節點個數
num = 0
for each in graph.node:
    print(num,"node")
    print(each)
    num +=1

三、onnx 輸出資訊表示

向量模型舉例.....
模型輸入：
model input:
[name: "input_ids" #名稱
type {
  tensor_type {
    elem_type: 1 #輸入型別
    shape {
      dim {
        dim_value: 1 #輸入第一維度
      }
      dim {
        dim_value: 512 #輸入第二維度
      }
    }
  }
}
]

model output:
[name: "1329" #模型輸出的節點名稱，對應到onnx的node的output
type {
  tensor_type {
    elem_type: 1 #輸出型別
    shape {
      dim {
        dim_value: 1#輸的第一維度
      }
      dim {
        dim_value: 768 #輸出的第二維度
      }
    }
  }
}
]
節點：
input: "/encoder/encoder/layer.11/output/LayerNorm/Pow_output_0" #輸入節點名稱，對應上層output
output: "/encoder/encoder/layer.11/output/LayerNorm/ReduceMean_1_output_0" #輸出節點名，對應下層input，或者模型 output
name: "/encoder/encoder/layer.11/output/LayerNorm/ReduceMean_1" #當前節點名稱
op_type: "ReduceMean" #當前節點的op 操作
attribute { #節點屬性
　　name: "axes"
　　type: INTS
　　ints: -1
}

四、onnx 增，刪，改：最下面有示例

import onnx
import numpy as np
import sys
from onnx import helper
ori_file=sys.argv[1] #原始模型 model.onnx
onnxfile=sys.argv[2] #生成後的模型 model_modify.onnx
onnx_model = onnx.load(ori_file)
graph = onnx_model.graph
# 建立新節點
#new_output = helper.make_tensor_value_info(new_output_name,
#   model.graph.output[len(model.graph.output) - 1].type.tensor_type.elem_type,
#   model.graph.output[len(model.graph.output) - 1].type.tensor_type.shape)
modify_input_node=False
modify_output_node=False
delete_node=False
add_node=False

print("all_node_number: " + str(len(graph.node)))
all_node_len = len(graph.node)
node_index = all_node_len
#修改input節點名稱，以及型別;
if(modify_input_node):
    for index,eachNode in enumerate(graph.input):
        now_name = graph.input[index].name
        new_input = helper.make_tensor_value_info(now_name,
                       graph.input[index].type.tensor_type.elem_type,
                       [1,256]) #修改模型輸出維度
        graph.input.remove(graph.input[index]) #刪除舊節點，
        #graph.input.append(new_input)      #插入新節點,順序會亂
        graph.input.insert(index,new_input)      #插入新節點,可以保證之前的順序

if(modify_output_node):#修改 output  名稱或者維度，或者型別....
    for index,eachNode in enumerate(graph.output):
        now_name = eachNode.name
        if(now_name == "1329"):
            new_output = helper.make_tensor_value_info("last_output_1",
                           graph.output[index].type.tensor_type.elem_type,
                           [1,768]) # 維度是list，不能是*tensor_type.shape
            graph.output.remove(graph.output[index]) #刪除舊節點，
            #graph.input.append(new_input)      #插入新節點,順序會亂
            graph.output.insert(index,new_output)      #插入新節點,保證原始順序

if( delete_node == True ): #刪除節點
    print("need_modify index : " + str(node_index))
    num = 0
    for i in range(0,10): #刪除最後10層模型,也可指定index刪除,但是遍歷或許有問題
         graph.node.remove(graph.node[-1])
    last_node = graph.node[-1]
    print("\last_node: ",last_node)
    print("after_delete_node_number: " + str(len(graph.node)))


if(add_node == True):#最後一層增加softmax，但是也需要修改output節點名稱和維度
    last_index = len(graph.node)
    old_node = graph.node[-1]
    print(old_node.output)
    new_node = onnx.helper.make_node(
        'Softmax',
        name='Softmax_1234',
        inputs=old_node.output,
        outputs=["last_output_1"], #是列表，傳入字串報錯，but單個字串不報錯
    )
    #graph.node.add(new_node) #這樣操作會報錯，TypeError: No positional arguments allowed
    graph.node.insert(last_index,new_node) # 最後一個節點插入
    print("add : \n",graph.node[-1])
    print("elem_type: ",graph.output[-1].type.tensor_type.elem_type) #列印elem_type
'''
onnx.checker.check_model()
檢查模型的一致性，即模型在結構、格式和配置方面的正確性和完整性。
model：要檢查的模型。如果模型是一個路徑，函式會首先檢查模型路徑。如果模型的位元組大小超過2GB，應使用模型路徑來呼叫該函式。
full_check：如果為 True，函式還會執行形狀推斷檢查。
skip_opset_compatibility_check：如果為 True，函式將跳過運算元集相容性檢查。
check_custom_domain：如果為 True，函式將檢查所有域。否則，僅檢查內建域。
'''
#onnx.checker.check_model(onnx_model)
onnx.checker.check_model(onnx_model,full_check=True)
onnx.save(onnx_model, onnxfile)

效果示例如下：右側處理後結果

修改input維度：

修改output名稱：

刪除最後10個節點：僅僅展示效果，實際刪除後或許還需要修改 outout名稱為對應node output

最後一層新增個softmax：僅僅展示效果，實際刪除後或許還需要修改 outout名稱為對應node output

　　

相關文章