Edit Onnx Model Ops

onnx_edit.py

import onnx

model_path = "./model.onnx"
fixed_model_path = model_path.replace(".onnx", ".fixed.onnx")

# # Load the ONNX model which should have last layer as Sigmoid. 
# LGBM Models may sometime not add the Sigmoid op during export when using regression loss
onnx_model = onnx.load(model_path)
print(onnx_model)
onnx.checker.check_model(onnx_model)
# Edit node
for node in onnx_model.graph.node:
  print(node.name)
  if node.name == "Identity":
    print(node)
    node.op_type = "Sigmoid"
    print(node)

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, fixed_model_path)

# Check model has new node
onnx_model = onnx.load(fixed_model_path)
for node in onnx_model.graph.node:
  print(node.name)
  if node.name == "Identity":
    print(node)
# Should show sigmoid op