Skip to content

Support models of 2GB or more in make_dynamic_shape_fixed - #31663

Open
adityasingh2400 wants to merge 1 commit into
microsoft:mainfrom
adityasingh2400:fix-make-dynamic-shape-fixed-large-models
Open

Support models of 2GB or more in make_dynamic_shape_fixed#31663
adityasingh2400 wants to merge 1 commit into
microsoft:mainfrom
adityasingh2400:fix-make-dynamic-shape-fixed-large-models

Conversation

@adityasingh2400

Copy link
Copy Markdown

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.

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

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Add --use_external_data_format to onnxruntime.tools.make_dynamic_shape_fixed

1 participant