diff --git a/examples/recipes/google_owlv2-base-patch16-finetuned/cpu/cpu/zero-shot-object-detection_fp16_config.json b/examples/recipes/google_owlv2-base-patch16-finetuned/cpu/cpu/zero-shot-object-detection_fp16_config.json new file mode 100644 index 000000000..3856c9fef --- /dev/null +++ b/examples/recipes/google_owlv2-base-patch16-finetuned/cpu/cpu/zero-shot-object-detection_fp16_config.json @@ -0,0 +1,97 @@ +{ + "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": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 49408 + ] + }, + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 960, + 960 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + }, + { + "name": "pred_boxes" + }, + { + "name": "text_embeds" + }, + { + "name": "image_embeds" + } + ] + }, + "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": "zero-shot-object-detection", + "model_id": "google/owlv2-base-patch16-finetuned", + "model_type": "owlv2", + "fp16_keep_io_types": true, + "fp16_op_block_list": null + }, + "compile": null, + "loader": { + "task": "zero-shot-object-detection", + "model_class": "AutoModelForZeroShotObjectDetection", + "model_type": "owlv2" + } +} diff --git a/examples/recipes/google_owlv2-base-patch16-finetuned/cpu/cpu/zero-shot-object-detection_fp32_config.json b/examples/recipes/google_owlv2-base-patch16-finetuned/cpu/cpu/zero-shot-object-detection_fp32_config.json new file mode 100644 index 000000000..93ed0b917 --- /dev/null +++ b/examples/recipes/google_owlv2-base-patch16-finetuned/cpu/cpu/zero-shot-object-detection_fp32_config.json @@ -0,0 +1,75 @@ +{ + "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": "input_ids", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 49408 + ] + }, + { + "name": "pixel_values", + "dtype": "float32", + "shape": [ + 1, + 3, + 960, + 960 + ], + "value_range": [ + 0, + 1 + ] + }, + { + "name": "attention_mask", + "dtype": "int32", + "shape": [ + 1, + 16 + ], + "value_range": [ + 0, + 2 + ] + } + ], + "output_tensors": [ + { + "name": "logits" + }, + { + "name": "pred_boxes" + }, + { + "name": "text_embeds" + }, + { + "name": "image_embeds" + } + ] + }, + "optim": {}, + "quant": null, + "compile": null, + "loader": { + "task": "zero-shot-object-detection", + "model_class": "AutoModelForZeroShotObjectDetection", + "model_type": "owlv2" + } +} diff --git a/src/winml/modelkit/export/htp/exporter.py b/src/winml/modelkit/export/htp/exporter.py index c61a162d3..bc052e391 100644 --- a/src/winml/modelkit/export/htp/exporter.py +++ b/src/winml/modelkit/export/htp/exporter.py @@ -20,6 +20,7 @@ from __future__ import annotations import contextlib +import inspect import logging import sys import time @@ -445,8 +446,12 @@ def _convert_model_to_onnx( output_path = str(Path(output_path).resolve()) Path(output_path).parent.mkdir(parents=True, exist_ok=True) - # Input names from config, fallback to inputs dict keys + # Input names from config, fallback to inputs dict keys. Keyword inputs + # invoke forward() by name, but the TorchScript exporter assigns these + # ONNX names positionally to the traced graph inputs. Align names with + # forward() order so a config's tensor order cannot swap graph bindings. input_names = export_config.get_input_names() or list(inputs.keys()) + input_names = self._resolve_keyword_input_names(model, inputs, input_names) # Output names: infer from traced hierarchy, validate against config traced_outputs = self._hierarchy_builder.get_outputs() if self._hierarchy_builder else None @@ -495,6 +500,54 @@ def _convert_model_to_onnx( else: torch.onnx.export(model, (), output_path, kwargs=inputs, **onnx_kwargs) + @staticmethod + def _resolve_keyword_input_names( + model: nn.Module, + inputs: dict[str, torch.Tensor], + configured_input_names: list[str], + ) -> list[str]: + """Return ONNX names in the order produced by keyword tracing. + + ``torch.onnx.export`` invokes the model safely with keyword arguments, + then applies ``input_names`` positionally to graph inputs emitted in + ``forward`` signature order. Reorder only when every configured, + generated input has an unambiguous signature match. Models implementing + ``get_export_args`` own a separate positional protocol, so their config + order remains authoritative. + """ + if hasattr(model, "get_export_args"): + return configured_input_names + + try: + parameters = inspect.signature(model.forward).parameters + except (TypeError, ValueError): + logger.debug( + "Could not inspect %s.forward; preserving configured input order.", + type(model).__name__, + ) + return configured_input_names + + configured_set = set(configured_input_names) + generated_set = set(inputs) + forward_input_names = [ + name for name in parameters if name in configured_set and name in generated_set + ] + if ( + len(forward_input_names) == len(configured_input_names) + and set(forward_input_names) == configured_set + ): + return forward_input_names + + logger.debug( + "Could not align every configured input for %s.forward; " + "preserving configured input order. configured=%s generated=%s forward=%s", + type(model).__name__, + configured_input_names, + list(inputs), + list(parameters), + ) + return configured_input_names + @staticmethod def _get_optimum_patcher(model: nn.Module, task: str | None) -> Any: """Get Optimum's model patcher for TorchScript tracing compatibility. diff --git a/tests/unit/export/test_pytorch_export.py b/tests/unit/export/test_pytorch_export.py index 5b6aeae26..39c8e9a90 100644 --- a/tests/unit/export/test_pytorch_export.py +++ b/tests/unit/export/test_pytorch_export.py @@ -87,12 +87,22 @@ def __init__(self): self.vision_conv = nn.Conv2d(3, 8, kernel_size=3, padding=1) self.pool = nn.AdaptiveAvgPool2d(1) - def forward(self, text_input, pixel_values): - text_out = self.text_fc(text_input) + def forward(self, input_ids, pixel_values, attention_mask): + text_out = self.text_fc(input_ids.float() * attention_mask.float()) vis_out = self.pool(self.vision_conv(pixel_values)).squeeze(-1).squeeze(-1) return text_out + vis_out +class PositionalExportOrderModel(nn.Module): + """Model whose explicit positional export protocol owns input ordering.""" + + def get_export_args(self, inputs): + return inputs["narrow"], inputs["wide"] + + def forward(self, wide, narrow): + return torch.cat((wide, narrow), dim=1) + + # ============================================================================= # TestInputTensorSpecToTensor # ============================================================================= @@ -350,20 +360,21 @@ def test_normalize_false_skips_normalization(self, tmp_path) -> None: assert not _all_value_info_have_shape(onnx_model) def test_mismatched_input_order_exports_successfully(self, tmp_path) -> None: - """Export succeeds when InputTensorSpec order differs from forward() param order. + """ONNX names bind correctly when specs differ from forward() order. Regression test for CLIP bug: OnnxConfig listed pixel_values before input_ids, - but CLIPModel.forward() expected input_ids first. With positional args this - caused 'not enough values to unpack (expected 4, got 2)'. The fix uses - kwargs= in torch.onnx.export so name-based binding makes order irrelevant. + but a multimodal forward expects text inputs first. Keyword arguments make + invocation order-independent; ONNX input names must separately follow the + forward signature because torch.onnx.export assigns them positionally. """ model = MismatchedInputOrderModel() - # Intentionally list pixel_values FIRST — opposite of forward(text_input, pixel_values) + # Intentionally list pixel_values FIRST — opposite of forward(). config = WinMLExportConfig( input_tensors=[ InputTensorSpec(name="pixel_values", dtype="float32", shape=(1, 3, 8, 8)), - InputTensorSpec(name="text_input", dtype="float32", shape=(1, 16)), + InputTensorSpec(name="input_ids", dtype="int32", shape=(1, 16)), + InputTensorSpec(name="attention_mask", dtype="int32", shape=(1, 16)), ], ) @@ -374,9 +385,37 @@ def test_mismatched_input_order_exports_successfully(self, tmp_path) -> None: onnx_model = onnx.load(str(tmp_path / "model.onnx")) onnx.checker.check_model(onnx_model) - input_names = {i.name for i in onnx_model.graph.input} - assert "pixel_values" in input_names - assert "text_input" in input_names + input_bindings = { + tensor.name: ( + onnx.TensorProto.DataType.Name(tensor.type.tensor_type.elem_type), + tuple(dim.dim_value for dim in tensor.type.tensor_type.shape.dim), + ) + for tensor in onnx_model.graph.input + } + assert input_bindings == { + "input_ids": ("INT32", (1, 16)), + "pixel_values": ("FLOAT", (1, 3, 8, 8)), + "attention_mask": ("INT32", (1, 16)), + } + + def test_get_export_args_preserves_configured_positional_names(self, tmp_path) -> None: + """An explicit get_export_args protocol keeps config order authoritative.""" + model = PositionalExportOrderModel() + config = WinMLExportConfig( + input_tensors=[ + InputTensorSpec(name="narrow", dtype="float32", shape=(1, 3)), + InputTensorSpec(name="wide", dtype="float32", shape=(1, 5)), + ], + ) + + export_pytorch(model, tmp_path / "model.onnx", config) + + onnx_model = onnx.load(str(tmp_path / "model.onnx")) + input_shapes = { + tensor.name: tuple(dim.dim_value for dim in tensor.type.tensor_type.shape.dim) + for tensor in onnx_model.graph.input + } + assert input_shapes == {"narrow": (1, 3), "wide": (1, 5)} class TestStaleExternalDataCleanup: