From a27ff56cb1d01bf55d48417e3eb89b0fb2d5f950 Mon Sep 17 00:00:00 2001 From: Shiyi Zheng Date: Fri, 31 Jul 2026 07:10:22 +0800 Subject: [PATCH] fix(export): derive normalized image input ranges --- .../cpu/keypoint-detection_fp16_config.json | 67 ++++++++++++ .../cpu/keypoint-detection_fp32_config.json | 45 ++++++++ src/winml/modelkit/export/io.py | 103 +++++++++++++++--- src/winml/modelkit/export/policy.py | 2 +- tests/unit/export/test_io.py | 79 ++++++++++++++ tests/unit/export/test_io_specs.py | 77 +++++++++++++ 6 files changed, 356 insertions(+), 17 deletions(-) create mode 100644 examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp16_config.json create mode 100644 examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp32_config.json diff --git a/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp16_config.json b/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp16_config.json new file mode 100644 index 000000000..ae6e68ac8 --- /dev/null +++ b/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp16_config.json @@ -0,0 +1,67 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 256, + 192 + ], + "value_range": [ + -2.1179039478302, + 2.640000343322754 + ] + } + ], + "output_tensors": [ + { + "name": "heatmaps" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": {}, + "quant": { + "mode": "fp16", + "samples": 10, + "calibration_method": "minmax", + "weight_type": "uint8", + "activation_type": "uint8", + "per_channel": false, + "symmetric": false, + "weight_symmetric": null, + "activation_symmetric": null, + "save_calibration": false, + "distribution": "uniform", + "seed": null, + "calibration_load_path": null, + "calibration_save_path": null, + "op_types_to_quantize": null, + "nodes_to_exclude": null, + "task": "keypoint-detection", + "model_id": "usyd-community/vitpose-plus-huge", + "model_type": "vitpose", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "keypoint-detection", + "model_class": "VitPoseForPoseEstimation", + "model_type": "vitpose" + } +} \ No newline at end of file diff --git a/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp32_config.json b/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp32_config.json new file mode 100644 index 000000000..9cc610033 --- /dev/null +++ b/examples/recipes/usyd-community_vitpose-plus-huge/cpu/cpu/keypoint-detection_fp32_config.json @@ -0,0 +1,45 @@ +{ + "export": { + "opset_version": 17, + "batch_size": 1, + "export_params": true, + "do_constant_folding": true, + "verbose": false, + "dynamo": false, + "enable_hierarchy_tags": true, + "clean_onnx": false, + "hierarchy_tag_format": "full", + "input_tensors": [ + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 256, + 192 + ], + "value_range": [ + -2.1179039478302, + 2.640000343322754 + ] + } + ], + "output_tensors": [ + { + "name": "heatmaps" + } + ], + "compatibility": { + "transformers_attention": "eager" + } + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "keypoint-detection", + "model_class": "VitPoseForPoseEstimation", + "model_type": "vitpose" + } +} \ No newline at end of file diff --git a/src/winml/modelkit/export/io.py b/src/winml/modelkit/export/io.py index 3bedca226..9cacf49b6 100644 --- a/src/winml/modelkit/export/io.py +++ b/src/winml/modelkit/export/io.py @@ -34,6 +34,7 @@ import logging from typing import TYPE_CHECKING, Any, cast +import numpy as np from optimum.exporters.tasks import TasksManager from optimum.utils.input_generators import ( DEFAULT_DUMMY_SHAPES, @@ -211,6 +212,8 @@ def _populate_image_size_from_preprocessor( model_id: str | None, shape_kwargs: dict, hf_config: PretrainedConfig | None = None, + *, + preprocessor_config: dict | None = None, ) -> None: """Populate height/width in shape_kwargs from preprocessor metadata. @@ -233,7 +236,11 @@ def _populate_image_size_from_preprocessor( if "height" in shape_kwargs or "width" in shape_kwargs: return - config = _get_preprocessor_dict(model_id, hf_config) + config = ( + preprocessor_config + if preprocessor_config is not None + else _get_preprocessor_dict(model_id, hf_config) + ) size = config.get("size") if isinstance(size, int): @@ -263,17 +270,16 @@ def _get_preprocessor_dict( Resolution order: - 1. ``preprocessor_config.json`` fetched from the hub (standard HF vision), - used only when it carries a ``size`` key. - 2. Synthesized from a nested plain-dict attribute on ``hf_config`` - carrying ``input_size`` or ``image_size`` (e.g. - ``TimmWrapperConfig.pretrained_cfg``). Reached when the hub file is - unavailable *or* present but missing ``size`` (a partial config). + 1. ``preprocessor_config.json`` fetched from the hub (standard HF vision). + 2. Synthesized from a nested plain-dict attribute on ``hf_config`` carrying + ``input_size`` or ``image_size``. Its size is used when the hub file is + unavailable or present but missing ``size``. - Returns the dict in the standard preprocessor schema (``{"size": ...}``) - so downstream parsing logic does not need to know which source it came - from. Returns an empty dict when neither source yields a usable size. + Returns the full standard preprocessor schema so downstream logic can use + both shape and normalization metadata. Returns an empty dict when neither + source yields usable metadata. """ + config: dict = {} try: if model_id is None: raise OSError("No model_id provided") @@ -284,14 +290,58 @@ def _get_preprocessor_dict( config, _ = ImageProcessingMixin.get_image_processor_dict(model_id) if "size" in config: return config - # Partial preprocessor_config.json without a "size" key: fall through - # to synthesis so we don't silently use Optimum's 64x64 default. except (OSError, ValueError, KeyError) as e: logger.debug("Could not load preprocessor_config.json for %s: %s", model_id, e) if hf_config is not None: - return _synthesize_preprocessor_dict(hf_config) - return {} + synthesized = _synthesize_preprocessor_dict(hf_config) + if synthesized: + return {**config, **synthesized} + return config + + +def _normalized_image_value_range(config: dict) -> tuple[float, float] | None: + """Derive the normalized float32 image domain from processor metadata.""" + if config.get("do_rescale") is not True or config.get("do_normalize") is not True: + return None + + try: + rescale_factor = np.float32(config["rescale_factor"]) + image_mean, image_std = np.broadcast_arrays( + np.atleast_1d(np.asarray(config["image_mean"], dtype=np.float32)), + np.atleast_1d(np.asarray(config["image_std"], dtype=np.float32)), + ) + except (KeyError, TypeError, ValueError, OverflowError): + return None + + if ( + image_mean.ndim != 1 + or image_std.ndim != 1 + or image_mean.size == 0 + or not np.isfinite(rescale_factor) + or not np.all(np.isfinite(image_mean)) + or not np.all(np.isfinite(image_std)) + or np.any(image_std <= 0) + ): + return None + + raw_endpoints = np.asarray([0, 255], dtype=np.float32) + normalized_endpoints = ( + raw_endpoints[:, np.newaxis] * rescale_factor - image_mean + ) / image_std + low = np.min(normalized_endpoints) + high = np.nextafter(np.max(normalized_endpoints), np.float32(np.inf)) + return float(low), float(high) + + +def _is_image_float_input( + shape: tuple[int, ...], dtype: str, channels: int | None +) -> bool: + """Return whether a generated tensor is a floating image-shaped input.""" + if not dtype.startswith("float") or len(shape) < 4: + return False + expected_channels = {channels} if channels is not None else {1, 3, 4} + return bool(expected_channels.intersection((shape[1], shape[-1]))) def _synthesize_preprocessor_dict(hf_config: PretrainedConfig) -> dict: @@ -409,7 +459,13 @@ def generate_dummy_inputs( onnx_config.float_dtype = float_dtype shape_kwargs["batch_size"] = batch_size - _populate_image_size_from_preprocessor(model_id, shape_kwargs, hf_config) + preprocessor_config = _get_preprocessor_dict(model_id, hf_config) + _populate_image_size_from_preprocessor( + model_id, + shape_kwargs, + hf_config, + preprocessor_config=preprocessor_config, + ) _populate_sequence_length_from_config(hf_config, shape_kwargs) logger.debug( @@ -476,7 +532,13 @@ def resolve_io_specs( # Populate shapes from model config / preprocessor shape_kwargs["batch_size"] = batch_size - _populate_image_size_from_preprocessor(model_id, shape_kwargs, hf_config) + preprocessor_config = _get_preprocessor_dict(model_id, hf_config) + _populate_image_size_from_preprocessor( + model_id, + shape_kwargs, + hf_config, + preprocessor_config=preprocessor_config, + ) _populate_sequence_length_from_config(hf_config, shape_kwargs) # Generate dummy inputs for concrete shapes and dtypes, @@ -491,6 +553,15 @@ def resolve_io_specs( value_range_tuples = { name: (info["min"], info["max"]) for name, info in value_ranges.items() } + normalized_range = _normalized_image_value_range(preprocessor_config) + image_mean = preprocessor_config.get("image_mean") + channels = len(image_mean) if isinstance(image_mean, (list, tuple)) else None + if normalized_range is not None: + for name, shape, dtype in zip( + onnx_config.inputs, input_shapes, input_dtypes, strict=False + ): + if _is_image_float_input(shape, dtype, channels): + value_range_tuples[name] = normalized_range return { "inputs": onnx_config.inputs, diff --git a/src/winml/modelkit/export/policy.py b/src/winml/modelkit/export/policy.py index 2e54c4c84..943022516 100644 --- a/src/winml/modelkit/export/policy.py +++ b/src/winml/modelkit/export/policy.py @@ -161,7 +161,7 @@ def resolve_export_compatibility( def _catalog_targets() -> tuple[ExportPolicyTarget, ...]: - from ..session.ep_device import EP_DEVICE_SPECS + from ..session import EP_DEVICE_SPECS return tuple(ExportPolicyTarget(ep=spec.ep, device=spec.device) for spec in EP_DEVICE_SPECS) diff --git a/tests/unit/export/test_io.py b/tests/unit/export/test_io.py index 0ce30f475..0c5c23f77 100644 --- a/tests/unit/export/test_io.py +++ b/tests/unit/export/test_io.py @@ -35,6 +35,8 @@ from winml.modelkit.export import generate_dummy_inputs, resolve_io_specs from winml.modelkit.export.io import ( # Testing internal implementation _get_onnx_config, + _is_image_float_input, + _normalized_image_value_range, _populate_image_size_from_preprocessor, ) from winml.modelkit.models.winml.kv_cache import PastKeyValueInputGenerator @@ -792,6 +794,83 @@ def test_existing_height_blocks_nested_dict_too(self) -> None: assert shape_kwargs == {"height": 128} +class TestNormalizedImageValueRange: + """Tests processor-derived float image input ranges.""" + + def test_channel_extrema_and_float32_high_exclusive_bound(self) -> None: + config = { + "do_rescale": True, + "do_normalize": True, + "rescale_factor": 1 / 255, + "image_mean": [0.485, 0.456, 0.406], + "image_std": [0.229, 0.224, 0.225], + } + + assert _normalized_image_value_range(config) == ( + -2.1179039478302, + 2.640000343322754, + ) + + def test_scalar_metadata(self) -> None: + config = { + "do_rescale": True, + "do_normalize": True, + "rescale_factor": 1 / 255, + "image_mean": 0.5, + "image_std": 0.5, + } + + assert _normalized_image_value_range(config) == (-1.0, 1.0000001192092896) + + def test_scalar_std_broadcasts_across_channel_means(self) -> None: + config = { + "do_rescale": True, + "do_normalize": True, + "rescale_factor": 1 / 255, + "image_mean": [0.25, 0.5, 0.75], + "image_std": 0.5, + } + + assert _normalized_image_value_range(config) == (-1.5, 1.5000001192092896) + + @pytest.mark.parametrize( + "config", + [ + {}, + {"do_rescale": False, "do_normalize": True}, + {"do_rescale": True, "do_normalize": False}, + { + "do_rescale": True, + "do_normalize": True, + "rescale_factor": 1 / 255, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5], + }, + ], + ) + def test_missing_disabled_or_incompatible_metadata_falls_back(self, config: dict) -> None: + assert _normalized_image_value_range(config) is None + + @pytest.mark.parametrize( + "shape,dtype,channels,expected", + [ + ((1, 3, 256, 192), "float32", 3, True), + ((1, 3, 256, 192), "float32", None, True), + ((1, 256, 192, 3), "float16", 3, True), + ((1, 80, 3000), "float32", 3, False), + ((1, 3, 256, 192), "int64", 3, False), + ], + ) + def test_only_image_shaped_float_inputs_are_selected( + self, + shape: tuple[int, ...], + dtype: str, + channels: int | None, + expected: bool, + ) -> None: + assert _is_image_float_input(shape, dtype, channels) is expected + + # ============================================================================= # PastKeyValueInputGenerator — shared KV cache dummy input generation # ============================================================================= diff --git a/tests/unit/export/test_io_specs.py b/tests/unit/export/test_io_specs.py index 73b3d99b8..64e078fa4 100644 --- a/tests/unit/export/test_io_specs.py +++ b/tests/unit/export/test_io_specs.py @@ -17,6 +17,7 @@ from __future__ import annotations from typing import ClassVar +from unittest.mock import patch import pytest @@ -329,6 +330,82 @@ def test_resnet_vision_ranges(self, resnet_config) -> None: assert ranges["pixel_values"] == (0, 1) + def test_normalized_processor_metadata_overrides_image_range(self, resnet_config) -> None: + """Normalized image metadata replaces Optimum's raw float range.""" + processor_config = { + "size": 32, + "do_rescale": True, + "do_normalize": True, + "rescale_factor": 1 / 255, + "image_mean": [0.485, 0.456, 0.406], + "image_std": [0.229, 0.224, 0.225], + } + + with patch( + "transformers.image_processing_utils.ImageProcessingMixin.get_image_processor_dict", + return_value=(processor_config, {}), + ): + specs = resolve_io_specs( + "resnet", + "image-classification", + resnet_config, + model_id="local/vision-model", + ) + + assert specs["value_ranges"]["pixel_values"] == ( + -2.1179039478302, + 2.640000343322754, + ) + + def test_disabled_normalization_preserves_image_fallback(self, resnet_config) -> None: + """Disabled normalization retains Optimum's intercepted float range.""" + processor_config = { + "size": 32, + "do_rescale": True, + "do_normalize": False, + "rescale_factor": 1 / 255, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + } + + with patch( + "transformers.image_processing_utils.ImageProcessingMixin.get_image_processor_dict", + return_value=(processor_config, {}), + ): + specs = resolve_io_specs( + "resnet", + "image-classification", + resnet_config, + model_id="local/vision-model", + ) + + assert specs["value_ranges"]["pixel_values"] == (0, 1) + + def test_image_processor_metadata_does_not_change_text_inputs(self, bert_config) -> None: + """Processor metadata cannot replace non-image tensor ranges.""" + processor_config = { + "size": 32, + "do_rescale": True, + "do_normalize": True, + "rescale_factor": 1 / 255, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], + } + + with patch( + "transformers.image_processing_utils.ImageProcessingMixin.get_image_processor_dict", + return_value=(processor_config, {}), + ): + specs = resolve_io_specs( + "bert", + "fill-mask", + bert_config, + model_id="local/text-model", + ) + + assert specs["value_ranges"]["input_ids"] == (0, bert_config.vocab_size) + assert specs["value_ranges"]["attention_mask"] == (0, 2) + def test_gpt2_text_ranges(self, gpt2_config) -> None: """GPT-2 text inputs: input_ids uses vocab_size, position_ids uses n_positions.""" specs = resolve_io_specs("gpt2", "text-generation", gpt2_config)