Support models of 2GB or more in make_dynamic_shape_fixed - #31663
Open
adityasingh2400 wants to merge 1 commit into
Open
Support models of 2GB or more in make_dynamic_shape_fixed#31663adityasingh2400 wants to merge 1 commit into
adityasingh2400 wants to merge 1 commit into
Conversation
The tool fails on any model at or above the protobuf message limit. fix_output_shapes calls onnx.shape_inference.infer_shapes, which serializes the model into a single protobuf, so the run dies with "Message onnx.ModelProto exceeds maximum protobuf size of 2GB" before anything is written. Saving the result has the same problem. Run shape inference through onnx.shape_inference.infer_shapes_path for models of that size. It keeps the initializer data in a separate file, so the 2GB proto limit applies to the graph alone. Only the inferred output shapes are read back, so the initializer data stays on disk rather than being loaded a second time. Models below the limit keep taking the existing in memory path. Also add --use_external_data_format so the fixed model can be written with its initializers in a separate file. It is turned on automatically for models of 2GB or more, since those cannot be saved any other way. Fixes microsoft#16773 Signed-off-by: Aditya Singh <adisin650@gmail.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #16773
onnxruntime.tools.make_dynamic_shape_fixed fails on any model at or above the protobuf message limit. fix_output_shapes calls onnx.shape_inference.infer_shapes, which serializes the model into a single protobuf, so the run dies with "Message onnx.ModelProto exceeds maximum protobuf size of 2GB" before anything is written. Saving the result has the same problem.
Shape inference now goes through onnx.shape_inference.infer_shapes_path for models of that size. It keeps the initializer data in a separate file, so the 2GB limit applies to the graph alone. Only the inferred output shapes are read back, so the initializer data stays on disk rather than being loaded a second time. Models below the limit keep taking the existing in-memory path, so nothing changes for them.
This also adds --use_external_data_format so the fixed model can be written with its initializers in a separate file. It turns on automatically for models of 2GB or more, since those cannot be saved any other way.
Verification: added two tests to tools/python/util/test/test_onnx_model_utils.py. One drops PROTOBUF_SIZE_LIMIT to exercise the path based shape inference branch without building a 2GB model, the other drives the command line tool end to end and checks the .data file is written and the shapes are fixed. All 8 tests in that file pass, and both new tests fail against the unmodified code. ruff check and ruff format are clean on the changed files, the two remaining ruff findings in onnx_model_utils.py predate this change.