|
|
|
|
|
import onnx
|
|
|
from onnx import helper
|
|
|
from onnx import AttributeProto, TensorProto, GraphProto
|
|
|
import time
|
|
|
import sys
|
|
|
|
|
|
def check_update_model(model):
|
|
|
if (model.graph.name == ''):
|
|
|
model.graph.name = 'tackle_checker_error'
|
|
|
|
|
|
try:
|
|
|
onnx.checker.check_model(model)
|
|
|
except onnx.checker.ValidationError as e:
|
|
|
print("model is invalid: %s"%(e))
|
|
|
# 目前只考虑这个错误“ Nodes in a graph must be topologically sorted, however input ...”，
|
|
|
# 就是模型的节点顺序有问题，onnxruntime 是可以推断的，为了其他引擎比如 trt可以推断，调整一下节点顺序。
|
|
|
else:
|
|
|
time.sleep(1)
|
|
|
print("模型正常，返回原模型")
|
|
|
return model
|
|
|
|
|
|
print('update model ...')
|
|
|
new_model = onnx.ModelProto(ir_version=model.ir_version,
|
|
|
producer_name=model.producer_name,
|
|
|
producer_version=model.producer_version,
|
|
|
opset_import=model.opset_import)
|
|
|
|
|
|
new_model.graph.name = model.graph.name
|
|
|
# 只考虑模型是单输入单输出的情况。
|
|
|
# 这里不需要 copy.deepcopy(model.graph.input[0]), python 是写时copy。\
|
|
|
list_total_input_name = []
|
|
|
for inp in model.graph.input:
|
|
|
new_model.graph.input.append(inp)
|
|
|
list_total_input_name.append(inp.name)
|
|
|
for oup in model.graph.output:
|
|
|
new_model.graph.output.append(oup)
|
|
|
|
|
|
|
|
|
for weight in model.graph.initializer:
|
|
|
new_model.graph.initializer.append(weight)
|
|
|
list_total_input_name.append(weight.name)
|
|
|
|
|
|
list_nodes_invalid = []
|
|
|
for node in model.graph.node:
|
|
|
node_invalid = False
|
|
|
for node_input_name in node.input:
|
|
|
if (list_total_input_name.count(node_input_name) == 0):
|
|
|
list_nodes_invalid.append(node)
|
|
|
node_invalid = True
|
|
|
break
|
|
|
if (node_invalid):
|
|
|
continue
|
|
|
for node_output_name in node.output:
|
|
|
list_total_input_name.append(node_output_name)
|
|
|
new_model.graph.node.append(node)
|
|
|
|
|
|
for node in list_nodes_invalid:
|
|
|
for node_input_name in node.input:
|
|
|
assert list_total_input_name.count(node_input_name) > 0
|
|
|
for node_output_name in node.output:
|
|
|
list_total_input_name.append(node_output_name)
|
|
|
new_model.graph.node.append(node)
|
|
|
|
|
|
return new_model
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
onnx_in_name = sys.argv[1]
|
|
|
onnx_out_name = sys.argv[2]
|
|
|
|
|
|
onnx_model = onnx.load(onnx_in_name)
|
|
|
new_model = check_update_model(onnx_model)
|
|
|
onnx.checker.check_model(new_model)
|
|
|
|
|
|
onnx.save(new_model, onnx_out_name) |