ONNX模型节点Node修改及Head输出简化方法

​
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)

​

你可能感兴趣的:(java,前端,服务器)