From 457ce6cb8910bdfa7e0117030e6c83838b946ea6 Mon Sep 17 00:00:00 2001
From: Chi Lo <54722500+chilo-ms@users.noreply.github.com>
Date: Mon, 20 Jun 2022 13:07:45 -0700
Subject: [PATCH] Make symbolic shape inference script support external weight
 (#11909)

* add support for external data

* fix format

* fix format

* fix typo

* fix typo
---
 .../python/tools/symbolic_shape_infer.py      | 36 ++++++++++++++++++-
 1 file changed, 35 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/python/tools/symbolic_shape_infer.py b/onnxruntime/python/tools/symbolic_shape_infer.py
index 617b010ad6..da9e3b626d 100755
--- a/onnxruntime/python/tools/symbolic_shape_infer.py
+++ b/onnxruntime/python/tools/symbolic_shape_infer.py
@@ -2389,6 +2389,29 @@ def parse_arguments():
         type=int,
         default=0,
     )
+    parser.add_argument(
+        "--save_as_external_data",
+        help="Saving an ONNX model to external data",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--all_tensors_to_one_file",
+        help="Saving all the external data to one file",
+        action="store_true",
+        default=False,
+    )
+    parser.add_argument(
+        "--external_data_location",
+        help="The file location to save the external file",
+        default="./",
+    )
+    parser.add_argument(
+        "--external_data_size_threshold",
+        help="The size threshold for external data",
+        type=int,
+        default=1024,
+    )
     return parser.parse_args()
 
 
@@ -2406,5 +2429,16 @@ if __name__ == "__main__":
         args.verbose,
     )
     if args.output and out_mp:
-        onnx.save(out_mp, args.output)
+        if args.save_as_external_data:
+            onnx.save_model(
+                out_mp,
+                args.output,
+                save_as_external_data=True,
+                all_tensors_to_one_file=args.all_tensors_to_one_file,
+                location=args.external_data_location,
+                size_threshold=args.external_data_size_threshold,
+                convert_attribute=False,
+            )
+        else:
+            onnx.save(out_mp, args.output)
         logger.info("Done!")