saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-05-16 21:00:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add ms domain during saving onnx model in onnx_model.py (#13978)

Browse source 
Add domain "com.microsoft" during saving model if needed.
This commit is contained in:
Tianlei Wu 2022-12-16 22:45:57 -08:00 committed by GitHub
parent cc0a6213e4
commit 6fb54fc607
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
onnxruntime/python/tools/transformers/onnx_model.py
View file

@ -917,6 +917,16 @@ class OnnxModel:
):
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
# Add ms domain if needed
ms_opset = [opset for opset in model.opset_import if opset.domain == "com.microsoft"]
# Check whether there is custom op in top level graph (our fusion is on top level right now).
# May need to extend to subgraph if our fusion are extended to subgraphs.
ms_node = [node for node in model.graph.node if node.domain == "com.microsoft"]
if ms_node and not ms_opset:
opset = model.opset_import.add()
opset.version = 1
opset.domain = "com.microsoft"
if save_as_external_data:
# Save model to external data, which is needed for model size > 2GB
output_dir = Path(output_path).parent