From 6fb54fc607d94de50a18997430d7c87b0d921e87 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 16 Dec 2022 22:45:57 -0800 Subject: [PATCH] Add ms domain during saving onnx model in onnx_model.py (#13978) Add domain "com.microsoft" during saving model if needed. --- onnxruntime/python/tools/transformers/onnx_model.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index fa4294962e..2dcd2db901 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -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