onnxruntime 使用,删除、修改

输出中间层数据(python)

import onnx
from onnx import helper

model = onnx.load("C:/Users/Desktop/fastscnn_0907_fix.onnx")
#名字命名为想要输出层对应的节点名,参数类型、维度,需要自己给出
info =  helper.make_tensor_value_info('552',onnx.TensorProto.INT64, [1,1280, 1280])
#将构建的输出插入到模型中
model.graph.output.insert(0, info)
onnx.save(model, 'C:/Users/Desktop/out.onnx')

删除节点

import onnx
from onnx import helper

model = onnx.load("C:/Users/Desktop/fastscnn_0907_fix.onnx")
graph = model.graph
nodes = graph.node
for i in range(len(nodes)):
    print(i,nodes[i])
#通过打印查看想要删除的节点,并记录删除节点位置i
#删除节点,255为所有节点nodes的索引
graph.node.remove(nodes[255])
#检查模型
onnx.checker.check_model(model)
#保存新模型
onnx.save(model, 'C:/Users/Desktop/out.onnx')

修改节点,可通过先删除,再插入的方式

import onnx
from onnx import helper

model = onnx.load("C:/Users/Desktop/fastscnn.onnx")
graph = model.graph
nodes = graph.node
for i in range(len(nodes)):
    print(i,nodes[i])
#通过打印查看想要删除的节点,并记录删除节点位置i
#删除节点,255为所有节点nodes的索引
graph.node.remove(nodes[255])
#插入新节点
new_node = onnx.helper.make_node(
    "Unsqueeze",
    inputs=[''],
    outputs=['output'],
    axes = [0],
  )  
graph.node.insert(255, new_node)
#检查模型
onnx.checker.check_model(model)
#保存新模型
onnx.save(model, 'C:/Users/Desktop/out.onnx')

你可能感兴趣的:(python)