onnx モデルは、conv の背後にあるノードを切り取り、最後のノード名と一致するように出力レイヤー名を設定し、出力ノードと一致するように出力レイヤーの形状を設定します。

一部のモデルの最後の畳み込み層の後の演算子は、推論エンジンでの実行に適していないため、conv の後の演算子を削除すると、CPU でのパフォーマンスが向上します。
内容: 
1. onnx 中間ノードの形状を取得する例
2. onnx モデルの出力を追加し、名前、タイプ、形状を設定します。例
3。onnx モデルの例を編集し


、例の緑色の部分を切り取ります。
import onnx 
import sys 
import json 
from onnx import shade_inference, TensorProto 

if len(sys.argv) < 2: 
    print('Usage: ' + sys.argv[0] + '<onnx_filename>' ) 
    exit(-1) 

onnx_file = sys.argv[1] 

# ONNX モデルをロード
model = onnx.load(onnx_file) 

chart = model.graph

出力 = model.graph.output 
if(len(outputs)! =3): 
    print ("これは ScoreBoxKpt モデルではありません!") 
    quit() 

Output_list=["output0","output1",

出力の出力の場合:
    スコア_ボックス_kpt の出力名: 
        print(f"出力名: {output.name}") 
    else: 
        print("これは適合モデルではありません!") 
        quit() 

def getConvList(endName): 
    stack=[] 
    stack.append(endName) 
    convList=[] 
    while(len(stack)): 
        graph.node のノードの
        name=stack.pop() : 
            node.output の名前の場合: 
                node.op_type==の場合Conv": 
                    node.name が convList にない場合: 
                        convList.append(node.name) 
                else:
                        入力がスタックにない場合: 
                    node.input の入力の場合:
                            stack.insert(0, input) 
    return convList 

Conv0=getConvList(output_list[0]) 
Conv1=getConvList(output_list[1]) 
Conv2=getConvList(output_list[2]) 

def save2json(save_dict, name): 
    if len(save_dict) == 0: 
        print("これは json を保存するものではありません") 
        return None 
    with open(name, 'w') as fp: 
        #{'a': 'Runoob', 'b': 7} 
        json.dump(save_dict 、 fp、sort_keys=False、indent=4、separators=(',', ': ')) #default=str 

save_dict = {output_list[0]:scoreConv、
             output_list[1]:boxConv、
             output_list[2]:kptConv 
            }

conv_list=Conv0+Conv1+Conv2 

#获取onnx中点の形状。
output_dim_dic={} 
inferred_onnx_model =shape_inference.infer_shapes(model) 
inferred_graph = inferred_onnx_model.graph 
inferred_value_info = inferred_graph.value_info
グラフ.node のノード: conv_list の
    node.name の場合: 
        inferred_value_info の value_info: 
            value_info.name==node.output の場合[0]: 
                Output_dim_dic[node.name]=value_info.type.tensor_type; 

#删除conv 後の onnx 点# 
conv_list 内の name の
ターゲット ノード インデックスを検索します: 
    target_node = None 
    for node ingraph.node: 
        if node.name == name: 
            target_node=node
    target_node.output の出力の場合は、
    output_names = []
            を中断します
        。output_names.append(output) 
    set1=set(output_names) 
    del_node = [] 
    have_new_del_node = False 
    while True: 
        graph.node のノードの場合、 
        have_new_del_node = False : 
            if node.name in del_node :
                続行
            set2=set(node.input) 
            if set1.intersection(set2): 
                output_names+=node.output          
                set1=set(output_names) 
                del_node.append(node.name) 
                have_new_del_node = True

 
        if have_new_del_node == False: 
    graph.node 内のノードを
            ブレークします。
 
        del_node のnode.name の場合に: 
            print(f"1remove node {node.name}") 
            model.graph.node.remove(node) 

have_new_del_node = False 
while True:
    グラフのノード 1 に対して
    have_new_del_node = False .node: 
        conv_list の node1.name の場合:
            続行
        set1=set(node1.output) 
        to_delete =True 
        (graph.node のノード 2 の場合): 
            set2=set(node2.input) 
            if set1.intersection(set2): 
                to_delete = False 
                Break 
        if to_delete == True: 
            print(f"2remove node {node1 .name}") 
            model.graph.node.remove(node1) 
            have_new_del_node=True 
    if have_new_del_node == False :
グラフのノードの
save_output_name=[]
        を中断します
    。 ノード: conv_list のノードの名前: 
     #增加输出层
        output_info = onnx.helper .ValueInfoProto()
        ノード.output[0]=node.name
        出力_情報.名 = ノード.出力[0] 
        、output_dim_dic[node.name].shape.dim の dim_value: 
            output_info.type.tensor_type.shape.dim.extend( [dim_value]) 
        Output_info.type.tensor_type.elem_type = TensorProto.FLOAT 
        print(output_info) 
        graph.output.extend([output_info])

        save_output_name.append(node.output[0])
 
Outputs = model.graph.output 
# 打印出力点名
出力の出力の
    場合: if Output.name in save_output_name : 
        continue 
    model.graph.output.remove(output) 
Outputs = model.graph.output #
出力の
打印出力节点名: 
    if Output.name in save_output_name : 
        continue 
    model.graph.output.remove(output) 
# 変更された ONNX モデルを保存
onnx.checker.check_model(model) 
onnx.save(model, "backbone.onnx") 
save2json(save_dict, 'conv_param .json」

おすすめ

転載: blog.csdn.net/soralaro/article/details/132621291