模型要求
维度要求:[NCW*H]
节点要求:不支持特殊节点
修改onnx输入输出维度
import onnx
model = onnx.load("end2end.onnx")
print(model.graph.input)
print(model.graph.output)
model.graph.input[0].type.tensor_type.shape.dim[0].dim_value=1
model.graph.output[0].type.tensor_type.shape.dim[0].dim_value=1
#若onnx输入输出为fp16,可以修改elem_type=1,变成fp32
onnx.checker.check_model(model)
onnx.save(model,"out.onnx")
删除onnx中特殊节点
import onnx
from onnx import numpy_helper
def remove_flatten_node(model_path):
model = onnx.load(model_path)
graph = model.graph
# Find all the flatten nodes in the model
flatten_nodes = []
for node in graph.node:
if node.op_type == 'Flatten':
flatten_nodes.append(node)
# Remove the flatten nodes and rewire the graph
for node in flatten_nodes:
input_tensor_name = node.input
output_tensor_name = node.output
input_tensor_shape = None
for value_info in graph.value_info:
if value_info.name == input_tensor_name:
input_tensor_shape = list(d.dim_value for d in value_info.type.tensor_type.shape.dim)
# Create a Reshape node to replace the Flatten node
reshape_node = onnx.helper.make_node('Reshape',
inputs=input_tensor_name,
outputs=output_tensor_name,
name=node.name + '_reshape')
# Calculate the output shape of the Reshape node
output_tensor_shape = [input_tensor_shape, -1] + input_tensor_shape[2:]
shape_tensor = numpy_helper.from_array(output_tensor_shape, name='shape')
graph.node.append(reshape_node)
graph.initializer.append(shape_tensor)
# Rewire the graph to bypass the Flatten node
for idx, n in enumerate(graph.node):
if n == node:
graph.node[idx] = reshape_node
break
for i, inp in enumerate(n.input):
if inp == output_tensor_name:
n.input[i] = input_tensor_name
# Remove unused initializers and nodes
initializers_to_remove = set()
for value_info in graph.value_info:
initializers_to_remove.add(value_info.name)
for node in graph.node:
for inp in node.input:
if inp in initializers_to_remove:
initializers_to_remove.remove(inp)
new_initializers = []
for initializer in graph.initializer:
if initializer.name not in initializers_to_remove:
new_initializers.append(initializer)
del graph.initializer[:]
graph.initializer.extend(new_initializers)
new_nodes = []
for node in graph.node:
if not (node.op_type == 'Flatten'):
new_nodes.append(node)
del graph.node[:]
graph.node.extend(new_nodes)
# Update the output shapes of the model
onnx.checker.check_model(model)
onnx.save(model, model_path)
onnxpath = 'C:\\Users\\tunicorn\\PycharmProjects\\ct\\out.onnx'
remove_flatten_node(onnxpath)