From 430db19aeee07b37182ea3f1e2aee04f2b298a2d Mon Sep 17 00:00:00 2001 From: Hualiang Xie Date: Tue, 28 Apr 2026 16:23:10 +0800 Subject: [PATCH] add has_unloaded_external_data --- src/winml/modelkit/onnx/__init__.py | 3 +- src/winml/modelkit/onnx/utils.py | 25 ++++++ src/winml/modelkit/pattern/base.py | 17 +++- .../analyze/pattern/test_pattern_matching.py | 78 ++++++++++++++++++ tests/unit/onnx/test_onnx_utils.py | 82 +++++++++++++++++++ 5 files changed, 203 insertions(+), 2 deletions(-) create mode 100644 tests/unit/onnx/test_onnx_utils.py diff --git a/src/winml/modelkit/onnx/__init__.py b/src/winml/modelkit/onnx/__init__.py index 521f9c0f8..bf19e1b20 100644 --- a/src/winml/modelkit/onnx/__init__.py +++ b/src/winml/modelkit/onnx/__init__.py @@ -20,7 +20,7 @@ from .metadata import capture_metadata, restore_metadata from .persistence import cleanup_onnx, load_onnx, save_onnx from .shape import infer_onnx_shapes, infer_shapes -from .utils import EXTERNAL_DATA_THRESHOLD, check_onnx_model, get_model_size +from .utils import EXTERNAL_DATA_THRESHOLD, check_onnx_model, get_model_size, has_unloaded_external_data __all__ = [ @@ -36,6 +36,7 @@ "generate_inputs_from_onnx", "get_io_config", "get_model_size", + "has_unloaded_external_data", "infer_onnx_shapes", "infer_shapes", "is_compiled_onnx", diff --git a/src/winml/modelkit/onnx/utils.py b/src/winml/modelkit/onnx/utils.py index 87426a545..8aaeff267 100644 --- a/src/winml/modelkit/onnx/utils.py +++ b/src/winml/modelkit/onnx/utils.py @@ -13,6 +13,20 @@ EXTERNAL_DATA_THRESHOLD = 100 * 1024 * 1024 # 100 MiB +def has_unloaded_external_data(model: onnx.ModelProto) -> bool: + """Return True if the model contains external-data tensors whose bytes are not in memory. + + When a model is loaded with ``load_external_data=False``, tensors keep their + ``data_location == EXTERNAL`` annotation but ``raw_data`` stays empty. + The ONNX checker needs the sidecar ``.data`` file on disk to validate those + tensors, which may not be available in the current working directory. + """ + return any( + t.data_location == onnx.TensorProto.EXTERNAL and not t.raw_data + for t in _get_all_tensors(model) + ) + + def get_model_size(model: onnx.ModelProto) -> int: """Calculate the total size of an ONNX model in bytes. @@ -35,11 +49,22 @@ def check_onnx_model( full_check: bool = False, skip_opset_compatibility_check: bool = False, check_custom_domain: bool = False, + skip_if_unloaded_external_data: bool = False, ) -> None: """Same as ``onnx.checker.check_model``, but handles >2GiB models. Uses a temp file on disk for large models. + + Args: + skip_if_unloaded_external_data: When True, skip validation if the model + has tensors with ``data_location == EXTERNAL`` but no ``raw_data`` + (i.e. loaded with ``load_external_data=False``). The ONNX checker + needs the sidecar ``.data`` file on disk to validate those tensors, + which is not available in that case. """ + if skip_if_unloaded_external_data and has_unloaded_external_data(model): + return + tmp_dir = None if get_model_size(model) >= EXTERNAL_DATA_THRESHOLD: diff --git a/src/winml/modelkit/pattern/base.py b/src/winml/modelkit/pattern/base.py index ecbb3f07b..be502d3fa 100644 --- a/src/winml/modelkit/pattern/base.py +++ b/src/winml/modelkit/pattern/base.py @@ -1238,6 +1238,12 @@ def __init__(self, onnx_model: ModelProto, raise_on_invalid_model: bool = True) # Maps tensor name -> numpy array value (for constants/initializers) self.tensor_values: dict[str, np.ndarray] = {} + # Names of initializers whose external data was not loaded into memory. + # Constant-value constraints are skipped for these tensors so that + # topology-based pattern matching still works on models loaded with + # load_external_data=False. + self._external_unloaded_names: set[str] = set() + # Maps tensor name -> shape tuple self.tensor_shapes: dict[str, tuple] = {} @@ -1255,7 +1261,7 @@ def __init__(self, onnx_model: ModelProto, raise_on_invalid_model: bool = True) if raise_on_invalid_model: try: - check_onnx_model(self.model) + check_onnx_model(self.model, skip_if_unloaded_external_data=True) except onnx.checker.ValidationError as e: raise InvalidPatternMatcherModelError( f"Model failed ONNX validation: {e}", @@ -1320,6 +1326,11 @@ def _build_lookups(self) -> None: for initializer in self.graph.initializer: self.producer_lookup.setdefault(initializer.name, (initializer.name, 0, "Initializer")) if initializer.name: + if initializer.data_location == onnx.TensorProto.EXTERNAL and not initializer.raw_data: + # External data not loaded; record the name so constant-value + # checks can be skipped rather than returning a false failure. + self._external_unloaded_names.add(initializer.name) + continue self.tensor_values[initializer.name] = numpy_helper.to_array(initializer) for node_idx, node in enumerate(self.graph.node): @@ -1519,6 +1530,10 @@ def _check_constant_constraints( # Get tensor value if input_tensor_name not in self.tensor_values: + # If the initializer's external data was not loaded, skip the value + # check: topology matched, but we can't verify the scalar constant. + if input_tensor_name in self._external_unloaded_names: + continue return False actual_value = self.tensor_values[input_tensor_name] diff --git a/tests/unit/analyze/pattern/test_pattern_matching.py b/tests/unit/analyze/pattern/test_pattern_matching.py index f92f1f661..68cbdcf5f 100644 --- a/tests/unit/analyze/pattern/test_pattern_matching.py +++ b/tests/unit/analyze/pattern/test_pattern_matching.py @@ -7,9 +7,12 @@ from pathlib import Path +import numpy as np import onnx import onnx.helper as oh +import onnx.numpy_helper as onph import pytest +from onnx import TensorProto from winml.modelkit.pattern import ( Gelu2Pattern, @@ -328,3 +331,78 @@ def test_raise_on_invalid_model_false_still_works(self): matcher.register_pattern(Gelu2Pattern()) results = matcher.match() assert len(results) == 1 + + +def _make_gelu_model_with_external_initializers() -> onnx.ModelProto: + """Gelu2 model where weight initializers simulate load_external_data=False. + + The graph topology is identical to a normal Gelu2 model; only the weight + tensors have data_location=EXTERNAL with empty raw_data, mimicking what + onnx.load(..., load_external_data=False) produces for large models. + """ + X = oh.make_tensor_value_info("X", TensorProto.FLOAT, [1, 8]) # noqa: N806 + Y = oh.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 8]) # noqa: N806 + + def _external_tensor(name: str, values: list, dtype=TensorProto.FLOAT, dims=None): + t = onph.from_array( + np.array(values, dtype=np.float32 if dtype == TensorProto.FLOAT else np.int64), + name=name, + ) + t.data_location = TensorProto.EXTERNAL + t.ClearField("raw_data") + entry = t.external_data.add() + entry.key = "location" + entry.value = "model.onnx.data" + return t + + sqrt2 = _external_tensor("sqrt2", [1.4142135]) + half = _external_tensor("half", [0.5]) + one = _external_tensor("one_val", [1.0]) + + div = oh.make_node("Div", ["X", "sqrt2"], ["div_out"], name="div_node") + erf = oh.make_node("Erf", ["div_out"], ["erf_out"], name="erf_node") + add = oh.make_node("Add", ["erf_out", "one_val"], ["add_out"], name="add_node") + mul1 = oh.make_node("Mul", ["X", "add_out"], ["mul1_out"], name="mul1_node") + mul2 = oh.make_node("Mul", ["mul1_out", "half"], ["Y"], name="mul2_node") + + graph = oh.make_graph( + [div, erf, add, mul1, mul2], + "gelu_external", + [X], + [Y], + initializer=[sqrt2, half, one], + ) + return oh.make_model(graph, opset_imports=[oh.make_opsetid("", 13)]) + + +class TestPatternMatchingWithUnloadedExternalData: + """PatternMatcher must not fail when model initializers have unloaded external data. + + This covers the regression where models loaded with load_external_data=False + triggered an ONNX ValidationError in check_onnx_model because the sidecar + .data file could not be found, causing pattern matching to be silently skipped. + """ + + def test_pattern_matcher_does_not_raise(self): + model = _make_gelu_model_with_external_initializers() + matcher = PatternMatcher(model) + assert matcher is not None + + def test_gelu_pattern_still_matched_on_topology(self): + """Topology-based pattern matching succeeds even without tensor values.""" + model = _make_gelu_model_with_external_initializers() + matcher = PatternMatcher(model) + matcher.register_pattern(Gelu2Pattern()) + results = matcher.match() + assert len(results) == 1, f"Expected 1 Gelu2Pattern match, got {len(results)}" + + def test_unloaded_initializers_absent_from_tensor_values(self): + """External-data initializers must not pollute tensor_values with zero arrays.""" + model = _make_gelu_model_with_external_initializers() + matcher = PatternMatcher(model) + # None of the external initializer names should have values loaded. + external_names = {init.name for init in model.graph.initializer} + for name in external_names: + assert name not in matcher.tensor_values, ( + f"Initializer '{name}' should not be in tensor_values when external data is unloaded" + ) diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py new file mode 100644 index 000000000..9108ffa68 --- /dev/null +++ b/tests/unit/onnx/test_onnx_utils.py @@ -0,0 +1,82 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- +"""Tests for modelkit.onnx.utils — check_onnx_model and has_unloaded_external_data.""" + +from __future__ import annotations + +import numpy as np +import onnx +import pytest +from onnx import TensorProto, helper, numpy_helper + +from winml.modelkit.onnx import check_onnx_model, has_unloaded_external_data + + +def _make_simple_model() -> onnx.ModelProto: + """Minimal valid ONNX model with an inline initializer.""" + x_info = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) + y_info = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]) + weight = numpy_helper.from_array(np.ones((4,), dtype=np.float32), name="W") + node = helper.make_node("Add", ["X", "W"], ["Y"], name="add") + graph = helper.make_graph([node], "g", [x_info], [y_info], initializer=[weight]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +def _make_model_with_unloaded_external_data() -> onnx.ModelProto: + """Model whose initializer mimics load_external_data=False: data_location=EXTERNAL, raw_data empty.""" + x_info = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) + y_info = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]) + + weight = numpy_helper.from_array(np.ones((4,), dtype=np.float32), name="W") + # Simulate what onnx.load(..., load_external_data=False) produces for large tensors. + weight.data_location = TensorProto.EXTERNAL + weight.ClearField("raw_data") + entry = weight.external_data.add() + entry.key = "location" + entry.value = "model.onnx.data" + + node = helper.make_node("Add", ["X", "W"], ["Y"], name="add") + graph = helper.make_graph([node], "g", [x_info], [y_info], initializer=[weight]) + return helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) + + +class TestHasUnloadedExternalData: + def test_inline_model_returns_false(self): + assert has_unloaded_external_data(_make_simple_model()) is False + + def test_unloaded_external_tensor_returns_true(self): + assert has_unloaded_external_data(_make_model_with_unloaded_external_data()) is True + + def test_loaded_external_data_returns_false(self): + """After raw_data is populated the tensor is considered loaded.""" + model = _make_model_with_unloaded_external_data() + init = model.graph.initializer[0] + init.raw_data = np.ones((4,), dtype=np.float32).tobytes() + assert has_unloaded_external_data(model) is False + + +class TestCheckOnnxModel: + def test_valid_inline_model_does_not_raise(self): + check_onnx_model(_make_simple_model()) + + def test_unloaded_external_data_does_not_raise_when_skip_enabled(self): + """check_onnx_model must not raise when skip_if_unloaded_external_data=True.""" + check_onnx_model( + _make_model_with_unloaded_external_data(), + skip_if_unloaded_external_data=True, + ) + + def test_unloaded_external_data_raises_by_default(self): + """check_onnx_model raises by default when external data is missing on disk.""" + with pytest.raises(onnx.checker.ValidationError): + check_onnx_model(_make_model_with_unloaded_external_data()) + + def test_invalid_model_raises(self): + """check_onnx_model should still raise for structurally invalid models.""" + model = _make_simple_model() + # Corrupt the opset version to something invalid. + model.opset_import[0].version = 0 + with pytest.raises(onnx.checker.ValidationError): + check_onnx_model(model)