From 85529d8173f0430ba776044f55dc4e55f94f6251 Mon Sep 17 00:00:00 2001 From: Shiyi Zheng Date: Thu, 16 Jul 2026 10:58:19 +0800 Subject: [PATCH] feat(depth-pro): support FP16 conversion and evaluation --- .../cpu/cpu/depth-estimation_fp16_config.json | 67 +++++++++++++++++++ .../cpu/cpu/depth-estimation_fp32_config.json | 45 +++++++++++++ .../modelkit/models/winml/depth_estimation.py | 26 +++++-- src/winml/modelkit/quant/fp16.py | 12 ++++ .../models/winml/test_depth_estimation.py | 23 +++++++ tests/unit/test_quant_passes.py | 35 ++++++++++ 6 files changed, 201 insertions(+), 7 deletions(-) create mode 100644 examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_fp16_config.json create mode 100644 examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_fp32_config.json diff --git a/examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_fp16_config.json b/examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_fp16_config.json new file mode 100644 index 000000000..b809da782 --- /dev/null +++ b/examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_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, + 1536, + 1536 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "predicted_depth" + }, + { + "name": "field_of_view" + } + ] + }, + "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": "depth-estimation", + "model_id": "apple/DepthPro-hf", + "model_type": "depth_pro", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "depth-estimation", + "model_class": "AutoModelForDepthEstimation", + "model_type": "depth_pro" + } +} diff --git a/examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_fp32_config.json b/examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_fp32_config.json new file mode 100644 index 000000000..3828b2688 --- /dev/null +++ b/examples/recipes/apple_DepthPro-hf/cpu/cpu/depth-estimation_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, + 1536, + 1536 + ], + "value_range": [ + 0, + 1 + ] + } + ], + "output_tensors": [ + { + "name": "predicted_depth" + }, + { + "name": "field_of_view" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "depth-estimation", + "model_class": "AutoModelForDepthEstimation", + "model_type": "depth_pro" + } +} diff --git a/src/winml/modelkit/models/winml/depth_estimation.py b/src/winml/modelkit/models/winml/depth_estimation.py index 4fd3d17ab..42c16e6fd 100644 --- a/src/winml/modelkit/models/winml/depth_estimation.py +++ b/src/winml/modelkit/models/winml/depth_estimation.py @@ -12,19 +12,25 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Any, cast +from dataclasses import dataclass +from typing import Any, cast +import torch from transformers.modeling_outputs import DepthEstimatorOutput from .base import WinMLPreTrainedModel -if TYPE_CHECKING: - import torch - logger = logging.getLogger(__name__) +@dataclass +class WinMLDepthEstimatorOutput(DepthEstimatorOutput): + """Depth output with optional architecture-specific camera metadata.""" + + field_of_view: torch.FloatTensor | None = None + + class WinMLModelForDepthEstimation(WinMLPreTrainedModel): """WinML model for monocular depth estimation. @@ -33,7 +39,7 @@ class WinMLModelForDepthEstimation(WinMLPreTrainedModel): ``image_processor.post_process_depth_estimation()``. """ - def forward(self, **kwargs: Any) -> DepthEstimatorOutput: + def forward(self, **kwargs: Any) -> WinMLDepthEstimatorOutput: """Run depth estimation inference. Accepts all processor outputs via ``**kwargs`` and passes them @@ -41,7 +47,7 @@ def forward(self, **kwargs: Any) -> DepthEstimatorOutput: architecture-agnostic. Returns: - DepthEstimatorOutput with the ``predicted_depth`` tensor populated. + Depth output with ``predicted_depth`` and available camera metadata. """ formatted = self._format_inputs(**kwargs) outputs = self._run_inference(formatted) @@ -54,4 +60,10 @@ def forward(self, **kwargs: Any) -> DepthEstimatorOutput: # transformers' Output fields are annotated FloatTensor (legacy, over-narrow); # the ONNX session returns a real float Tensor. depth: torch.FloatTensor = cast("torch.FloatTensor", predicted_depth) - return DepthEstimatorOutput(predicted_depth=depth) + field_of_view: torch.FloatTensor | None = cast( + "torch.FloatTensor | None", outputs.get("field_of_view") + ) + return WinMLDepthEstimatorOutput( + predicted_depth=depth, + field_of_view=field_of_view, + ) diff --git a/src/winml/modelkit/quant/fp16.py b/src/winml/modelkit/quant/fp16.py index b3a6fe3a3..7b46b0444 100644 --- a/src/winml/modelkit/quant/fp16.py +++ b/src/winml/modelkit/quant/fp16.py @@ -47,6 +47,8 @@ def convert_to_fp16( from onnx import TensorProto from onnxruntime.transformers.float16 import convert_float_to_float16 + from ..onnx import EXTERNAL_DATA_THRESHOLD, get_model_size + # Skip if model is already FP16 (check floating-point initializer dtypes) fp32_types = {TensorProto.FLOAT, TensorProto.DOUBLE, TensorProto.BFLOAT16} initializers = model.graph.initializer @@ -64,9 +66,19 @@ def convert_to_fp16( if op_block_list: logger.info(" Keeping ops in FP32: %s", op_block_list) + # ORT's default shape-inference path serializes the complete ModelProto in + # memory. Protobuf cannot serialize very large models (for example, + # DepthPro is ~3.6 GiB), even when their tensors came from external data. + # Export and optimization have already populated shape information, so + # bypass that redundant, size-unsafe step for external-data-scale models. + disable_shape_infer = get_model_size(model) >= EXTERNAL_DATA_THRESHOLD + if disable_shape_infer: + logger.info(" Skipping in-memory shape inference for large model") + converted: ModelProto = convert_float_to_float16( model, keep_io_types=keep_io_types, + disable_shape_infer=disable_shape_infer, op_block_list=op_block_list, ) diff --git a/tests/unit/models/winml/test_depth_estimation.py b/tests/unit/models/winml/test_depth_estimation.py index c64677376..995a07fda 100644 --- a/tests/unit/models/winml/test_depth_estimation.py +++ b/tests/unit/models/winml/test_depth_estimation.py @@ -79,6 +79,29 @@ def test_attribute_and_dict_access(self): assert out.predicted_depth is depth assert out["predicted_depth"] is depth + def test_field_of_view_passthrough_when_available(self): + """DepthPro camera metadata remains available to its post-processor.""" + depth = torch.zeros((1, 16, 16)) + field_of_view = torch.full((1,), 55.0) + model = _make_model( + {"predicted_depth": depth, "field_of_view": field_of_view} + ) + + out = model.forward(pixel_values=torch.zeros((1, 3, 16, 16))) + + assert isinstance(out, DepthEstimatorOutput) + assert out.field_of_view is field_of_view + assert out["field_of_view"] is field_of_view + + def test_field_of_view_is_optional(self): + """Single-output depth architectures retain their existing behavior.""" + model = _make_model({"predicted_depth": torch.zeros((1, 16, 16))}) + + out = model.forward(pixel_values=torch.zeros((1, 3, 16, 16))) + + assert out.field_of_view is None + assert "field_of_view" not in out + def test_falls_back_to_first_output_when_name_differs(self): """Non-standard output names use the first tensor (architecture-agnostic).""" depth = torch.full((1, 8, 8), 2.5) diff --git a/tests/unit/test_quant_passes.py b/tests/unit/test_quant_passes.py index d78cbe410..1e3d48665 100644 --- a/tests/unit/test_quant_passes.py +++ b/tests/unit/test_quant_passes.py @@ -337,6 +337,41 @@ def fake_convert(model, *, keep_io_types, op_block_list): assert calls == [{"keep_io_types": False, "op_block_list": ["Gather"]}] +class TestConvertToFP16: + @pytest.mark.parametrize( + ("model_size", "expected_disable_shape_infer"), + [(1, False), (100 * 1024 * 1024, True)], + ) + def test_disables_shape_inference_for_large_models( + self, + monkeypatch: pytest.MonkeyPatch, + model_size: int, + expected_disable_shape_infer: bool, + ) -> None: + """Large external-data models must avoid protobuf serialization in ORT.""" + import onnx + import onnxruntime.transformers.float16 as ort_fp16 + + import winml.modelkit.onnx as onnx_utils + from winml.modelkit.quant.fp16 import convert_to_fp16 + + model = onnx.helper.make_model( + onnx.helper.make_graph([], "g", [], []), + opset_imports=[onnx.helper.make_opsetid("", 17)], + ) + calls: list[dict[str, Any]] = [] + + def fake_convert(candidate, **kwargs): # type: ignore[no-untyped-def] + calls.append(kwargs) + return candidate + + monkeypatch.setattr(onnx_utils, "get_model_size", lambda _model: model_size) + monkeypatch.setattr(ort_fp16, "convert_float_to_float16", fake_convert) + + assert convert_to_fp16(model) is model + assert calls[0]["disable_shape_infer"] is expected_disable_shape_infer + + # --------------------------------------------------------------------------- # RTNPass — config field wiring # ---------------------------------------------------------------------------