From ce8f459735d7151f91cd9d9ccbd6e37c3ebda36c Mon Sep 17 00:00:00 2001 From: Bowen Bao Date: Fri, 5 Apr 2024 13:52:01 -0700 Subject: [PATCH] Skip full model shape inference if model > 2GB | feat(optimizer) (#1340) Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom): * #1334 * __->__ #1340 --- onnxscript/optimizer/__init__.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/onnxscript/optimizer/__init__.py b/onnxscript/optimizer/__init__.py index 98c2038d9..8a4326979 100644 --- a/onnxscript/optimizer/__init__.py +++ b/onnxscript/optimizer/__init__.py @@ -56,9 +56,15 @@ def optimize( ) for _ in range(num_iterations): if onnx_shape_inference: - model = onnx.shape_inference.infer_shapes( - model, check_type=True, strict_mode=True, data_prop=True - ) + if model.ByteSize() < 1024 * 1024 * 1024 * 2: + model = onnx.shape_inference.infer_shapes( + model, check_type=True, strict_mode=True, data_prop=True + ) + else: + logger.warning( + "The model size is too large for full model shape inference. " + "Skipping this step." + ) inline_simple_functions(model) modified = fold_constants(