diff --git a/onnxscript/rewriter/ort_fusions/_test_utils.py b/onnxscript/rewriter/ort_fusions/_test_utils.py index 24e9bcce61..35092ddf45 100644 --- a/onnxscript/rewriter/ort_fusions/_test_utils.py +++ b/onnxscript/rewriter/ort_fusions/_test_utils.py @@ -2,6 +2,9 @@ # Licensed under the MIT License. from __future__ import annotations +import os +import tempfile + import numpy as np import onnx_ir as ir import onnxruntime @@ -12,13 +15,17 @@ def ort_run(model_name: str, model, inputs): providers = ["CPUExecutionProvider"] - model_proto = ir.serde.serialize_model(model) options = onnxruntime.SessionOptions() options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL - session = onnxruntime.InferenceSession( - model_proto.SerializeToString(), options, providers=providers - ) - return session.run(None, inputs) + # Save the model to a temporary file and load it from disk. This uses external data + # to store tensors so that large models are not limited by the 2GB protobuf + # serialization limit, which can otherwise raise + # ``google.protobuf.message.EncodeError: Failed to serialize proto``. + with tempfile.TemporaryDirectory() as temp_dir: + model_path = os.path.join(temp_dir, f"{model_name}.onnx") + ir.save(model, model_path, external_data=f"{model_name}.onnx.data") + session = onnxruntime.InferenceSession(model_path, options, providers=providers) + return session.run(None, inputs) def assert_allclose(outputs, expected_outputs, rtol=1e-3, atol=1e-3):