一部のモデルの最後の畳み込み層の後の演算子は、推論エンジンでの実行に適していないため、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」
onnx モデルは、conv の背後にあるノードを切り取り、最後のノード名と一致するように出力レイヤー名を設定し、出力ノードと一致するように出力レイヤーの形状を設定します。
おすすめ
転載: blog.csdn.net/soralaro/article/details/132621291
ランキング