Implement tensor similarity evaluator

pyproject.toml

@@ -63,6 +63,7 @@ dependencies = [
   "torchinfo>=1.8",
   "torchmetrics[detection]>=1.0",
   "torchvision>=0.24",
+  "tqdm>=4.67",
   "transformers>=4.57",
   "uvicorn[standard]>=0.42.0",
   "windowsml==2.0.300",
@@ -82,7 +83,6 @@ optional-dependencies.dev = [
   "ruff>=0.14",
   "seaborn>=0.13",
   "timm>=1.0.20",
-  "tqdm>=4.67",
   "types-colorama>=0.4.15.20250801",
 ]
 optional-dependencies.audio = [ "soundfile>=0.13" ]

scripts/e2e_eval/run_pytorch_baseline.py

@@ -347,7 +347,7 @@ def main() -> None:
 
         from winml.modelkit.eval.evaluate import get_evaluator_class
 
-        evaluator_cls = get_evaluator_class(task)
+        evaluator_cls = get_evaluator_class(eval_config)
         task_evaluator = evaluator_cls(eval_config, pytorch_model)
 
         metrics = task_evaluator.compute()

src/winml/modelkit/commands/eval.py

@@ -16,7 +16,7 @@
 from rich.console import Console
 
 from ..utils import cli as cli_utils
-from ..utils.eval_utils import TASK_SCHEMAS, TaskSchema
+from ..utils.eval_utils import EVAL_MODES, TASK_SCHEMAS, EvalMode, TaskSchema
 from ..utils.logging import configure_logging
 
 
@@ -149,6 +149,18 @@
     default=False,
     help="Print expected dataset schema for the given --task and exit.",
 )
+@click.option(
+    "--mode",
+    type=click.Choice(EVAL_MODES, case_sensitive=False),
+    default="onnx",
+    show_default=True,
+    help=(
+        "Evaluation mode. "
+        "'onnx' (default): evaluate the ONNX candidate on the dataset. "
+        "'compare': compare ONNX vs HF reference output tensors on identical "
+        "random inputs and report tensor-similarity metrics per output tensor."
+    ),
+)
 @cli_utils.build_config_option()
 @cli_utils.verbosity_options()
 @click.pass_context
@@ -175,6 +187,7 @@ def eval(
     dataset_script: str | None,
     trust_remote_code: bool,
     show_schema: bool,
+    mode: EvalMode,
     config_file: Path | None,
 ) -> None:
     r"""Evaluate a model for a task.
@@ -510,7 +523,8 @@ def display_eval_report(result: EvalResult, console: Console) -> None:
     console.print()
     console.print(f"[dim]Task:[/dim]       {cfg.task}")
     console.print(f"[dim]Device:[/dim]     {cfg.device}")
-    console.print(f"[dim]Dataset:[/dim]    {ds.path}")
+    if ds.path:
+        console.print(f"[dim]Dataset:[/dim]    {ds.path}")
     console.print(f"[dim]Samples:[/dim]    {ds.samples}")
     if cfg.model_path:
         console.print(f"[dim]ONNX:[/dim]       {cfg.model_path}")

src/winml/modelkit/eval/__init__.py

@@ -35,6 +35,7 @@
     from .metrics.top_k_accuracy import TopKAccuracyMetric
     from .object_detection_evaluator import WinMLObjectDetectionEvaluator
     from .question_answering_evaluator import WinMLQuestionAnsweringEvaluator
+    from .tensor_similarity_evaluator import TensorSimilarityEvaluator
     from .text_classification_evaluator import WinMLTextClassificationEvaluator
     from .token_classification_evaluator import WinMLTokenClassificationEvaluator
     from .zero_shot_classification_evaluator import WinMLZeroShotClassificationEvaluator
@@ -67,6 +68,8 @@
         ".zero_shot_classification_evaluator:WinMLZeroShotClassificationEvaluator",
     "WinMLZeroShotImageClassificationEvaluator":
         ".zero_shot_image_classification_evaluator:WinMLZeroShotImageClassificationEvaluator",
+    "TensorSimilarityEvaluator":
+        ".tensor_similarity_evaluator:TensorSimilarityEvaluator",
     # Metrics (defer numpy / scipy / torch / torchmetrics until first use)
     "ClassificationMetric":
         ".metrics.classification:ClassificationMetric",
@@ -117,6 +120,7 @@ def __dir__() -> list[str]:
     "MeanIoUMetric",
     "PseudoPerplexityMetric",
     "SpearmanCorrelationMetric",
+    "TensorSimilarityEvaluator",
     "TopKAccuracyMetric",
     "WinMLDepthEstimationEvaluator",
     "WinMLEvaluationConfig",

src/winml/modelkit/eval/config.py

@@ -12,6 +12,7 @@
 from typing import Any
 
 from ..utils.constants import EPNameOrAlias
+from ..utils.eval_utils import EvalMode
 
 
 @dataclass
@@ -93,6 +94,13 @@ class WinMLEvaluationConfig:
             device-to-provider mapping when provided.
         dataset: Dataset configuration.
         output_path: Path to write JSON results.
+        mode: Evaluation mode (see :data:`EvalMode`).
+
+            - ``"onnx"`` (default): evaluate the ONNX candidate on the
+              labeled dataset.
+            - ``"compare"``: compare ONNX vs HF reference output tensors
+              on identical random inputs and report tensor-similarity
+              metrics per output tensor.
 
     Usage:
         config = WinMLEvaluationConfig(
@@ -109,6 +117,7 @@ class WinMLEvaluationConfig:
     ep: EPNameOrAlias | None = None
     dataset: DatasetConfig = field(default_factory=DatasetConfig)
     output_path: Path | None = field(default=None, metadata={"cli_name": "output"})
+    mode: EvalMode = "onnx"
 
     def to_dict(self) -> dict:
         """Convert to dictionary for serialization."""
@@ -127,6 +136,8 @@ def to_dict(self) -> dict:
         result["dataset"] = self.dataset.to_dict()
         if self.output_path is not None:
             result["output_path"] = str(self.output_path)
+        if self.mode != "onnx":
+            result["mode"] = self.mode
         return result
 
     @classmethod
@@ -155,4 +166,5 @@ def from_dict(cls, data: dict) -> WinMLEvaluationConfig:
             ep=data.get("ep"),
             dataset=dataset,
             output_path=(Path(data["output_path"]) if data.get("output_path") else None),
+            mode=data.get("mode", "onnx"),
         )

src/winml/modelkit/eval/evaluate.py

@@ -59,16 +59,19 @@
         "winml.modelkit.eval.zero_shot_image_classification_evaluator:WinMLZeroShotImageClassificationEvaluator",
     "depth-estimation":
         "winml.modelkit.eval.depth_estimation_evaluator:WinMLDepthEstimationEvaluator",
+    "compare-tensor":
+        "winml.modelkit.eval.tensor_similarity_evaluator:TensorSimilarityEvaluator",
 }
 
 
-def get_evaluator_class(task: str) -> type[WinMLEvaluator]:
+def get_evaluator_class(config: WinMLEvaluationConfig) -> type[WinMLEvaluator]:
     """Return the evaluator class for *task*, or raise ValueError if unsupported."""
-    spec = _EVALUATOR_REGISTRY.get(task)
+    key = "compare-tensor" if config.mode == "compare" else config.task
+    spec = _EVALUATOR_REGISTRY.get(key)
     if spec is None:
         supported = ", ".join(sorted(_EVALUATOR_REGISTRY))
         raise ValueError(
-            f"Task '{task}' is not supported by `winml eval`. "
+            f"Task '{key}' is not supported by `winml eval`. "
             f"Supported tasks: {supported}."
         )
     module_path, class_name = spec.rsplit(":", 1)
@@ -286,8 +289,17 @@ def evaluate(config: WinMLEvaluationConfig) -> EvalResult:
     copies via ``dataclasses.replace`` and ``deepcopy`` so the original
     config and any module-level defaults remain untouched.
     """
-    config = replace(config, task=_resolve_task(config), dataset=deepcopy(config.dataset))
-    if config.dataset.path is None:
+    from ..utils.eval_utils import EVAL_MODES
+
+    mode = config.mode if config.mode is not None else "onnx"
+    if mode not in EVAL_MODES:
+        raise ValueError(
+            f"Invalid mode {mode!r}; expected one of {EVAL_MODES} or None."
+        )
+    config = replace(
+        config, mode=mode, task=_resolve_task(config), dataset=deepcopy(config.dataset)
+    )
+    if config.mode != "compare" and config.dataset.path is None:
         default = _DEFAULT_DATASETS.get(config.task)
         if default is None:
             raise ValueError(
@@ -318,7 +330,7 @@ def evaluate(config: WinMLEvaluationConfig) -> EvalResult:
 
     from ..utils.eval_utils import DatasetValidationError
 
-    cls = get_evaluator_class(config.task)
+    cls = get_evaluator_class(config)
     try:
         console.print("[bold]Loading dataset and evaluating...[/bold]")
         task_evaluator = cls(config, model)
@@ -351,15 +363,16 @@ def print_config(config: WinMLEvaluationConfig) -> None:
     if config.ep is not None:
         output_console.print(f"[bold blue]EP:[/bold blue] {config.ep}")
     output_console.print(f"[bold blue]Precision:[/bold blue] {config.precision}")
-    output_console.print(f"[bold blue]Dataset:[/bold blue] {ds.path}")
-    if ds.name:
-        output_console.print(f"[bold blue]Dataset name:[/bold blue] {ds.name}")
-    output_console.print(f"[bold blue]Split:[/bold blue] {ds.split}")
-    output_console.print(f"[bold blue]Samples:[/bold blue] {ds.samples}")
-    output_console.print(f"[bold blue]Shuffle:[/bold blue] {ds.shuffle} (seed={ds.seed})")
-    output_console.print(f"[bold blue]Streaming:[/bold blue] {ds.streaming}")
-    if ds.columns_mapping:
-        cols = ", ".join(f"{k}={v}" for k, v in ds.columns_mapping.items())
-        output_console.print(f"[bold blue]Columns:[/bold blue] {cols}")
+    if config.mode != "compare":
+        output_console.print(f"[bold blue]Dataset:[/bold blue] {ds.path}")
+        if ds.name:
+            output_console.print(f"[bold blue]Dataset name:[/bold blue] {ds.name}")
+        output_console.print(f"[bold blue]Split:[/bold blue] {ds.split}")
+        output_console.print(f"[bold blue]Samples:[/bold blue] {ds.samples}")
+        output_console.print(f"[bold blue]Shuffle:[/bold blue] {ds.shuffle} (seed={ds.seed})")
+        output_console.print(f"[bold blue]Streaming:[/bold blue] {ds.streaming}")
+        if ds.columns_mapping:
+            cols = ", ".join(f"{k}={v}" for k, v in ds.columns_mapping.items())
+            output_console.print(f"[bold blue]Columns:[/bold blue] {cols}")
     if config.output_path is not None:
         output_console.print(f"[bold blue]Output:[/bold blue] {config.output_path}")

src/winml/modelkit/eval/metrics/tensor_similarity.py

@@ -0,0 +1,137 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+"""Tensor-similarity metrics for compare-mode (ONNX vs HF reference parity).
+
+A :class:`TensorSimilarityMetric` instance accumulates per-sample scalar
+metrics across ``(prediction, reference)`` tensor pairs via ``update``
+and reports ``{f"{metric}_{stat}": float}`` (4 stats per metric: mean,
+std, min, max) via ``compute``. The per-sample math (SQNR, PSNR, cosine,
+MSE, max abs diff) mirrors the team-wide ``eval_tensors`` library so
+numbers match bit-for-bit on the same ``.npy`` pair.
+"""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+
+
+_SCALAR_METRICS = (
+    "sqnr_db",
+    "psnr_db",
+    "cosine_similarity",
+    "mse",
+    "max_abs_diff",
+)
+
+
+def _sqnr_db(ref: np.ndarray, test: np.ndarray) -> float:
+    """``10 * log10(sum(ref^2) / sum((ref-test)^2))``. ``+inf`` if identical."""
+    signal = float(np.sum(ref * ref))
+    noise = float(np.sum((ref - test) ** 2))
+    if noise == 0.0:
+        return math.inf
+    if signal == 0.0:
+        return -math.inf
+    return 10.0 * math.log10(signal / noise)
+
+
+def _mse(ref: np.ndarray, test: np.ndarray) -> float:
+    return float(np.mean((ref - test) ** 2))
+
+
+def _max_abs_diff(ref: np.ndarray, test: np.ndarray) -> float:
+    return float(np.max(np.abs(ref - test)))
+
+
+def _psnr_db(ref: np.ndarray, mse_val: float) -> float:
+    """``10 * log10(peak^2 / mse)``, ``peak = max(|ref|)``."""
+    if mse_val == 0.0:
+        return math.inf
+    peak = float(np.max(np.abs(ref)))
+    if peak == 0.0:
+        return -math.inf
+    return 10.0 * math.log10((peak * peak) / mse_val)
+
+
+def _cosine_similarity(ref: np.ndarray, test: np.ndarray) -> float:
+    """``dot(ref, test) / (||ref|| * ||test||)``, asymmetric zero handling.
+
+    Both inputs all-zero -> ``1.0`` (identical zero vectors).
+    Exactly one input all-zero -> ``0.0`` (a dead vector against a live
+    one is NOT a perfect match, even though the angle is undefined).
+    """
+    norm_ref = float(np.linalg.norm(ref))
+    norm_test = float(np.linalg.norm(test))
+    if norm_ref == 0.0 and norm_test == 0.0:
+        return 1.0
+    if norm_ref == 0.0 or norm_test == 0.0:
+        return 0.0
+    return float(np.dot(ref, test) / (norm_ref * norm_test))
+
+
+class TensorSimilarityMetric:
+    """Streaming per-sample tensor-parity metrics.
+
+    Each ``update(prediction, reference)`` computes the 5 scalar metrics
+    on the pair and appends them to internal per-metric lists. ``compute``
+    aggregates each list to ``mean`` / ``std`` / ``min`` / ``max`` and
+    returns a flat ``{f"{metric}_{stat}": float}`` dict ready for direct
+    consumption by the generic eval report renderer. ``mean`` and ``std``
+    are computed over only the finite values so a single bit-identical
+    sample (``sqnr_db = +inf``, ``psnr_db = +inf``) does not poison
+    the aggregate.
+    """
+
+    def __init__(self) -> None:
+        self._per_sample: dict[str, list[float]] = {m: [] for m in _SCALAR_METRICS}
+
+    def update(self, prediction: np.ndarray, reference: np.ndarray) -> None:
+        """Compute all scalar metrics on one pair and append to per-metric lists."""
+        if prediction.shape != reference.shape:
+            raise ValueError(
+                f"shape mismatch: prediction {prediction.shape} vs "
+                f"reference {reference.shape}",
+            )
+        ref = reference.astype(np.float64).ravel()
+        test = prediction.astype(np.float64).ravel()
+
+        mse_val = _mse(ref, test)
+        self._per_sample["sqnr_db"].append(_sqnr_db(ref, test))
+        self._per_sample["psnr_db"].append(_psnr_db(ref, mse_val))
+        self._per_sample["cosine_similarity"].append(_cosine_similarity(ref, test))
+        self._per_sample["mse"].append(mse_val)
+        self._per_sample["max_abs_diff"].append(_max_abs_diff(ref, test))
+
+    def compute(self) -> dict[str, float]:
+        """Return ``{f"{metric}_{stat}": float}`` for stats mean/std/min/max."""
+        result: dict[str, float] = {}
+        for metric, values in self._per_sample.items():
+            if not values:
+                continue
+            finite = [v for v in values if math.isfinite(v)]
+            if finite:
+                arr = np.asarray(finite, dtype=np.float64)
+                mean_val = float(arr.mean())
+                std_val = float(arr.std())
+            elif all(v == math.inf for v in values):
+                mean_val, std_val = math.inf, 0.0
+            elif all(v == -math.inf for v in values):
+                mean_val, std_val = -math.inf, 0.0
+            else:
+                # Any NaN, or a mix of +inf and -inf: un-summarizable.
+                mean_val, std_val = math.nan, math.nan
+            result[f"{metric}_mean"] = mean_val
+            result[f"{metric}_std"] = std_val
+            result[f"{metric}_min"] = float(min(values))
+            result[f"{metric}_max"] = float(max(values))
+        return result
+
+    def reset(self) -> None:
+        """Clear all accumulated per-sample values."""
+        for k in self._per_sample:
+            self._per_sample[k] = []