yolov5的onnx模型去除transpose层

import onnx
from onnx import helper, checker
from onnx import TensorProto
import re
import argparse
model = ""
onnx_model = onnx.load(model)
graph = onnx_model.graph
# print(graph)
node = graph.node

# node[213].output[0] = node[212].output[0]
# node[213].output[0] = node[213].input[0]
# for idx in graph.node:
#     print(idx)
# graph.node[]
graph.node.remove(graph.node[215])
#
graph.node.remove(graph.node[214])
graph.node.remove(graph.node[213])
# node[212].name = 'output'
for i in range(len(node)):
    # print(node[i])
    if node[i].name == 'Conv_212':
        node[i].output[0]="output"
    if node[i].name == 'Conv_193':
        node[i].output[0]="390"
    if node[i].name == 'Conv_174':
        node[i].output[0]="391"
onnx.checker.check_model(onnx_model)
onnx.save(onnx_model,"best_del_transpose.onnx")

你可能感兴趣的:(深度学习记录,人工智能,分类)