Skip to content

Instantly share code, notes, and snippets.

@napsternxg
Last active January 24, 2024 15:58
Show Gist options
  • Save napsternxg/60a726eb7735d7e94162c5f69be248f3 to your computer and use it in GitHub Desktop.
Save napsternxg/60a726eb7735d7e94162c5f69be248f3 to your computer and use it in GitHub Desktop.
Edit Onnx Model Ops
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment