import onnx
input_path = "deploy.onnx"
onnx_model = onnx.load(input_path)
graph = onnx_model.graph
node = graph.node
#要删除的onnx节点
x = {"onnx_node1", "onnx_node2"}
de = []
num = 0
for i in range(len(graph.node)):
if node[i].name in x:
de.append(i)
num = num+1
de.sort()
de.reverse()
for i in range(num):
graph.node.remove(graph.node[de[i]])
onnx.save(onnx_model, "newtest.onnx")
input_path = "newtest.onnx"
output_path = "resize.onnx"
input_names = ["data"]
#要保留的head output
output_names = ["output1", "output2"]
onnx.utils.extract_model(input_path, output_path, input_names, output_names)