diff --git a/docs/onnx/tr_hash_v8_nms_free.metadata.json b/docs/onnx/tr_hash_v8_nms_free.metadata.json new file mode 100644 index 0000000..89c7a25 --- /dev/null +++ b/docs/onnx/tr_hash_v8_nms_free.metadata.json @@ -0,0 +1,23 @@ +{ + "architecture_version": 8, + "image_size": 640, + "num_classes": 80, + "num_cells": 34000, + "regression_width": 68, + "reg_max": 16, + "scale_factors": [ + 1, + 2, + 4 + ], + "grid_sizes": [ + 160, + 80, + 40, + 20 + ], + "p2_head": true, + "branch": "nms-free", + "requires_nms": false, + "output_semantics": "raw_ltrb_dfl_and_quality_class_logits" +} diff --git a/docs/onnx/tr_hash_v8_o2m.metadata.json b/docs/onnx/tr_hash_v8_o2m.metadata.json new file mode 100644 index 0000000..9e974b1 --- /dev/null +++ b/docs/onnx/tr_hash_v8_o2m.metadata.json @@ -0,0 +1,23 @@ +{ + "architecture_version": 8, + "image_size": 640, + "num_classes": 80, + "num_cells": 34000, + "regression_width": 68, + "reg_max": 16, + "scale_factors": [ + 1, + 2, + 4 + ], + "grid_sizes": [ + 160, + 80, + 40, + 20 + ], + "p2_head": true, + "branch": "o2m", + "requires_nms": true, + "output_semantics": "raw_ltrb_dfl_and_quality_class_logits" +} diff --git a/docs/onnx/tr_hash_v8_validation_report.md b/docs/onnx/tr_hash_v8_validation_report.md new file mode 100644 index 0000000..5750c66 --- /dev/null +++ b/docs/onnx/tr_hash_v8_validation_report.md @@ -0,0 +1,177 @@ +# TR-HASH Vision v8 ONNX Validation + +Validation was run from commit `0fcf05c146f84a857d392b7da7d2947a41eb6d62` +(`Make Vision v8 the only detector architecture`) with the native checkpoint +`AETHORIA-AI/TR-HASH-Vision-v8-2M-COCO-SFT`. + +The generated ONNX binaries are intentionally not committed to the source +repository. Upload them as GitHub Release assets and link the release asset URLs +from the release notes or from a follow-up update to this report. The release +assets should include both ONNX binaries and their export metadata sidecars: + +- `tr_hash_v8_o2m.onnx` +- `tr_hash_v8_o2m.json` +- `tr_hash_v8_nms_free.onnx` +- `tr_hash_v8_nms_free.json` + +Use the following hashes and sizes to verify the uploaded release artifacts: + +| Branch | Artifact | Size | SHA-256 | +|---|---:|---:|---| +| O2M | `tr_hash_v8_o2m.onnx` | 11,104,476 bytes / 10.590054 MiB | `24CADE1A285475DDD6D8AB25EF6775F34BB9B686A08355CD35A39342A493256A` | +| NMS-free | `tr_hash_v8_nms_free.onnx` | 11,108,683 bytes / 10.594066 MiB | `CA24D9577FF2224BC1A4949D09A2E1CC514DB04B5DFFECE6B42223DE3326FDD7` | + +Branch metadata is committed beside this report: + +| Branch | Metadata | +|---|---| +| O2M | [`tr_hash_v8_o2m.metadata.json`](tr_hash_v8_o2m.metadata.json) | +| NMS-free | [`tr_hash_v8_nms_free.metadata.json`](tr_hash_v8_nms_free.metadata.json) | + +## Reproducibility + +Download the checkpoint snapshot: + +```powershell +python -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='AETHORIA-AI/TR-HASH-Vision-v8-2M-COCO-SFT', local_dir=r'models/TR-HASH-Vision-v8-2M-COCO-SFT')" +``` + +Export both raw prediction branches: + +```powershell +$env:PYTHONPATH='.' +python scripts/export_onnx.py models\TR-HASH-Vision-v8-2M-COCO-SFT --output tr_hash_v8_o2m.onnx --branch o2m --opset 17 --check +python scripts/export_onnx.py models\TR-HASH-Vision-v8-2M-COCO-SFT\best_nms_free --output tr_hash_v8_nms_free.onnx --branch nms-free --opset 17 --check +``` + +Validate parity with the calibrated v8 defaults: + +```powershell +$env:PYTHONPATH='.' +python scripts/check_onnx_parity.py models\TR-HASH-Vision-v8-2M-COCO-SFT tr_hash_v8_o2m.onnx --branch auto --num-tests 5 --batch-size 1 +python scripts/check_onnx_parity.py models\TR-HASH-Vision-v8-2M-COCO-SFT\best_nms_free tr_hash_v8_nms_free.onnx --branch auto --num-tests 5 --batch-size 1 +``` + +To reproduce the exact historical parity commands, pass the thresholds +explicitly: + +```powershell +python scripts/check_onnx_parity.py models\TR-HASH-Vision-v8-2M-COCO-SFT tr_hash_v8_o2m.onnx --branch auto --tolerance 0.002 --num-tests 5 --batch-size 1 +python scripts/check_onnx_parity.py models\TR-HASH-Vision-v8-2M-COCO-SFT\best_nms_free tr_hash_v8_nms_free.onnx --branch auto --tolerance 0.0035 --num-tests 5 --batch-size 1 +``` + +## Contract + +Both ONNX models were exported with opset 17 and expose raw detector logits +only. Decode and post-processing remain outside the graph. + +| Branch | Input | Output | Post-processing | +|---|---|---|---| +| O2M | `pixel_values`, `float32`, `[1, 3, 640, 640]` | `predictions`, `float32`, `[1, 34000, 148]` | Decode plus NMS | +| NMS-free | `pixel_values`, `float32`, `[1, 3, 640, 640]` | `predictions`, `float32`, `[1, 34000, 148]` | Decode plus confidence filtering | + +Output channel layout is `68` LTRB/DFL regression logits followed by `80` +quality-class logits. The grid pyramid is `[160, 80, 40, 20]`, for `34,000` +prediction cells. + +## Runtime Versions + +| Component | Version | +|---|---| +| PyTorch | `2.6.0+cu118` | +| ONNX | `1.21.0` | +| ONNX Runtime | `1.23.2` | +| ONNX Runtime providers | `TensorrtExecutionProvider`, `CUDAExecutionProvider`, `CPUExecutionProvider` | +| GPU | NVIDIA GeForce GTX 1660 SUPER, driver `560.94` | + +The CPU benchmark used `CPUExecutionProvider` on +`AMD64 Family 23 Model 113 Stepping 0, AuthenticAMD`. + +The GPU benchmark used `CUDAExecutionProvider` on the GTX 1660 SUPER. On this +Windows environment, ONNX Runtime required `ort.preload_dlls()` plus explicit +`nvidia/*/bin` DLL directories from the installed CUDA 12 and cuDNN 9 wheels. + +## Export Results + +Both exports succeeded: + +```text +o2m: Forward OK: output shape torch.Size([1, 34000, 148]) +nms-free: Forward OK: output shape torch.Size([1, 34000, 148]) +``` + +## Parity Results + +The legacy strict raw-logit threshold `1e-4` failed for both branches on CPU +ONNX Runtime: + +```text +o2m: max_diff range 1.36e-03 to 1.74e-03, mean_diff about 1.0e-04 +nms-free: max_diff range 2.36e-03 to 3.00e-03, mean_diff about 1.0e-04 +``` + +Tensor localization showed the largest raw differences are concentrated in +regression logits on the finest `160x160` grid. Class-logit drift is smaller. +Decoded-output drift is substantially lower: + +| Branch | Decoded boxes max diff | Decoded class score max diff | +|---|---:|---:| +| O2M | `6.00814819e-05` | `3.02493572e-05` | +| NMS-free | `5.73396683e-05` | `1.35712326e-05` | + +The calibrated v8 raw-logit thresholds are: + +| Branch | Raw-logit tolerance | Justification | +|---|---:|---| +| O2M | `0.002` | Covers observed max drift up to `1.74e-03` with decoded box drift near `6e-05` | +| NMS-free | `0.0035` | Covers observed max drift up to `3.00e-03` with decoded box drift near `5.7e-05` | + +Calibrated parity output: + +```text +O2M, tolerance 0.002: + Test 1/5: max_diff=1.61e-03, mean_diff=1.02e-04 [PASS] + Test 2/5: max_diff=1.36e-03, mean_diff=1.03e-04 [PASS] + Test 3/5: max_diff=1.63e-03, mean_diff=1.17e-04 [PASS] + Test 4/5: max_diff=1.36e-03, mean_diff=1.08e-04 [PASS] + Test 5/5: max_diff=1.74e-03, mean_diff=1.06e-04 [PASS] + Parity PASSED: branch=o2m, tolerance=0.002 + +NMS-free, tolerance 0.0035: + Test 1/5: max_diff=2.65e-03, mean_diff=1.03e-04 [PASS] + Test 2/5: max_diff=2.36e-03, mean_diff=1.01e-04 [PASS] + Test 3/5: max_diff=2.83e-03, mean_diff=1.11e-04 [PASS] + Test 4/5: max_diff=2.51e-03, mean_diff=1.00e-04 [PASS] + Test 5/5: max_diff=3.00e-03, mean_diff=1.10e-04 [PASS] + Parity PASSED: branch=nms-free, tolerance=0.0035 +``` + +## Benchmarks + +Benchmarks used batch size 1, `10` warmup runs, and `50` measured runs. + +CPU ONNX Runtime: + +| Branch | Provider | Mean latency | P95 latency | +|---|---|---:|---:| +| O2M | `CPUExecutionProvider` | `230.712 ms` | `240.292 ms` | +| NMS-free | `CPUExecutionProvider` | `244.064 ms` | `258.964 ms` | + +NMS-free was `13.351 ms` slower than O2M on mean latency, a `5.79%` increase. + +GPU ONNX Runtime: + +| Branch | Provider | Mean latency | P95 latency | +|---|---|---:|---:| +| O2M | `CUDAExecutionProvider` | `32.997 ms` | `33.890 ms` | +| NMS-free | `CUDAExecutionProvider` | `33.855 ms` | `34.724 ms` | + +NMS-free was `0.858 ms` slower than O2M on mean latency, a `2.60%` increase. + +## Conclusion + +TR-HASH Vision v8 exports successfully to ONNX for both raw prediction +branches. Branch behavior matches the expected architecture: NMS-free computes +the extra one-to-one head path and is empirically slower than O2M. Deployment +validation should use decoded-output drift and the calibrated v8 raw-logit +thresholds above rather than the legacy `1e-4` raw-logit threshold. diff --git a/scripts/check_onnx_parity.py b/scripts/check_onnx_parity.py index e8f8339..908f02f 100644 --- a/scripts/check_onnx_parity.py +++ b/scripts/check_onnx_parity.py @@ -20,6 +20,12 @@ from complexity.generative.detection.exporting import ExportBranch, RawDetectorExport from complexity.generative.detection.hub import load_detector_checkpoint +DEFAULT_PARITY_TOLERANCE = 1e-4 +V8_PARITY_TOLERANCES = { + "o2m": 2e-3, + "nms-free": 3.5e-3, +} + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) @@ -34,8 +40,11 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--tolerance", type=float, - default=1e-4, - help="Max allowed absolute difference (default: %(default)s)", + default=None, + help=( + "Max allowed absolute difference. Defaults to calibrated v8 branch " + "thresholds when ONNX metadata is available, otherwise 1e-4." + ), ) parser.add_argument( "--num-tests", @@ -52,24 +61,38 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() -def branch_from_sidecar(onnx_path: Path, requested: ExportBranch) -> ExportBranch: - if requested != "auto": - return requested +def sidecar_metadata(onnx_path: Path) -> dict: metadata_path = onnx_path.with_suffix(".json") if not metadata_path.is_file(): + return {} + return json.loads(metadata_path.read_text()) + + +def branch_from_sidecar(metadata: dict, requested: ExportBranch) -> ExportBranch: + if requested != "auto": + return requested + if "branch" not in metadata: return "auto" - branch = json.loads(metadata_path.read_text()).get("branch", "auto") + branch = metadata.get("branch", "auto") if branch not in {"nms-free", "o2m"}: - raise ValueError(f"invalid export branch in {metadata_path}: {branch}") + raise ValueError(f"invalid export branch in ONNX metadata: {branch}") return branch +def calibrated_parity_tolerance(metadata: dict, branch: ExportBranch) -> float: + """Return the default raw-logit parity tolerance for an exported model.""" + + if metadata.get("architecture_version") == 8 and branch in V8_PARITY_TOLERANCES: + return V8_PARITY_TOLERANCES[branch] + return DEFAULT_PARITY_TOLERANCE + + def check_parity( checkpoint_path: Path, onnx_path: Path, *, branch: ExportBranch = "auto", - tolerance: float = 1e-4, + tolerance: float | None = None, num_tests: int = 5, batch_size: int = 1, ) -> bool: @@ -85,7 +108,14 @@ def check_parity( print(f"Loading PyTorch model: {checkpoint_path}") detector = load_detector_checkpoint(checkpoint_path, device="cpu") - export_model = RawDetectorExport(detector, branch_from_sidecar(onnx_path, branch)).eval() + metadata = sidecar_metadata(onnx_path) + resolved_branch = branch_from_sidecar(metadata, branch) + effective_tolerance = ( + calibrated_parity_tolerance(metadata, resolved_branch) + if tolerance is None + else tolerance + ) + export_model = RawDetectorExport(detector, resolved_branch).eval() image_size = detector.config.image_size print(f"Prediction branch: {export_model.branch}") @@ -105,7 +135,7 @@ def check_parity( absolute_difference = np.abs(pytorch_output - onnx_output) max_difference = float(absolute_difference.max()) mean_difference = float(absolute_difference.mean()) - passed = max_difference <= tolerance + passed = max_difference <= effective_tolerance all_passed &= passed status = "PASS" if passed else "FAIL" print( @@ -114,7 +144,10 @@ def check_parity( ) summary = "PASSED" if all_passed else "FAILED" - print(f"\nParity {summary}: branch={export_model.branch}, tolerance={tolerance}") + print( + f"\nParity {summary}: " + f"branch={export_model.branch}, tolerance={effective_tolerance}" + ) return all_passed diff --git a/tests/test_detector_export.py b/tests/test_detector_export.py index 305d5c1..060588d 100644 --- a/tests/test_detector_export.py +++ b/tests/test_detector_export.py @@ -13,7 +13,13 @@ load_detector_checkpoint, ) from complexity.generative.detection.exporting import RawDetectorExport -from scripts.check_onnx_parity import check_parity +from scripts.check_onnx_parity import ( + DEFAULT_PARITY_TOLERANCE, + V8_PARITY_TOLERANCES, + branch_from_sidecar, + calibrated_parity_tolerance, + check_parity, +) from scripts.export_onnx import export_onnx @@ -76,6 +82,35 @@ def test_auto_export_selects_the_production_branch() -> None: RawDetectorExport(classic, "nms-free") +@pytest.mark.parametrize( + ("branch", "expected"), + ( + ("o2m", 2e-3), + ("nms-free", 3.5e-3), + ), +) +def test_v8_exports_use_branch_calibrated_parity_tolerances( + branch: str, + expected: float, +) -> None: + metadata = {"architecture_version": 8, "branch": branch} + + assert V8_PARITY_TOLERANCES[branch] == expected + assert calibrated_parity_tolerance(metadata, branch) == expected + + +def test_legacy_or_unlabelled_exports_keep_strict_parity_tolerance() -> None: + assert calibrated_parity_tolerance({}, "o2m") == DEFAULT_PARITY_TOLERANCE + assert ( + calibrated_parity_tolerance({"architecture_version": 7}, "nms-free") + == DEFAULT_PARITY_TOLERANCE + ) + + +def test_sidecar_less_auto_branch_still_defers_to_model_resolution() -> None: + assert branch_from_sidecar({}, "auto") == "auto" + + @pytest.mark.parametrize("branch", ("nms-free", "o2m")) def test_dynamic_onnx_export_matches_the_selected_branch(tmp_path: Path, branch: str) -> None: pytest.importorskip("onnx")