saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make symbolic shape inference script support external weight (#11909)

Browse source 
* add support for external data

* fix format

* fix format

* fix typo

* fix typo
This commit is contained in:
Chi Lo 2022-06-20 13:07:45 -07:00 committed by GitHub
parent c1577d08ca
commit 457ce6cb89
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 35 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
onnxruntime/python/tools/symbolic_shape_infer.py
View file

@ -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!")