From f8c017214fd4b56361d06cb587ac3d102bccfe32 Mon Sep 17 00:00:00 2001 From: Praveen Perera Date: Wed, 16 Sep 2026 22:00:13 +0000 Subject: [PATCH 1/3] Verify simplified ONNX outputs --- scripts/export_models.py | 113 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 111 insertions(+), 2 deletions(-) diff --git a/scripts/export_models.py b/scripts/export_models.py index e6d01ac..ead82b7 100644 --- a/scripts/export_models.py +++ b/scripts/export_models.py @@ -5,7 +5,9 @@ # "torch>=2.6", # "numpy", # "onnx", +# "onnxruntime", # "onnxscript", +# "onnxsim", # ] # /// """Download and export ONNX models + PLDA params for speakrs. @@ -19,7 +21,9 @@ """ import os +import stat import sys +import tempfile from typing import Any, cast import numpy as np @@ -31,6 +35,98 @@ os.environ.setdefault("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1") +def simplify_onnx_graph(path: str, inputs: dict[str, np.ndarray]) -> float: + """Simplify an ONNX graph after verifying equivalent CPU outputs.""" + import onnx + import onnxruntime as ort + from onnxsim import simplify + + original = onnx.load(path) + simplified, valid = simplify(original) + if not valid: + raise RuntimeError( + f"onnxsim could not validate the simplified graph for {path}" + ) + + onnx.checker.check_model(simplified) + original_session = ort.InferenceSession( + original.SerializeToString(), providers=["CPUExecutionProvider"] + ) + simplified_session = ort.InferenceSession( + simplified.SerializeToString(), providers=["CPUExecutionProvider"] + ) + + original_input_names = [value.name for value in original_session.get_inputs()] + simplified_input_names = [value.name for value in simplified_session.get_inputs()] + if original_input_names != simplified_input_names: + raise RuntimeError( + f"simplified input names changed for {path}: " + f"{original_input_names} != {simplified_input_names}" + ) + if set(inputs) != set(original_input_names): + raise RuntimeError( + f"validation inputs do not match {path}: " + f"{sorted(inputs)} != {sorted(original_input_names)}" + ) + + original_output_names = [value.name for value in original_session.get_outputs()] + simplified_output_names = [value.name for value in simplified_session.get_outputs()] + if original_output_names != simplified_output_names: + raise RuntimeError( + f"simplified output names changed for {path}: " + f"{original_output_names} != {simplified_output_names}" + ) + + original_outputs = original_session.run(None, inputs) + simplified_outputs = simplified_session.run(None, inputs) + max_abs_diff = 0.0 + for name, original_output, simplified_output in zip( + original_output_names, original_outputs, simplified_outputs, strict=True + ): + if original_output.shape != simplified_output.shape: + raise RuntimeError( + f"simplified output shape changed for {path}:{name}: " + f"{original_output.shape} != {simplified_output.shape}" + ) + if original_output.dtype != simplified_output.dtype: + raise RuntimeError( + f"simplified output dtype changed for {path}:{name}: " + f"{original_output.dtype} != {simplified_output.dtype}" + ) + if not np.isfinite(original_output).all(): + raise RuntimeError(f"original output is not finite for {path}:{name}") + if not np.isfinite(simplified_output).all(): + raise RuntimeError(f"simplified output is not finite for {path}:{name}") + + output_diff = ( + float(np.max(np.abs(original_output - simplified_output))) + if original_output.size + else 0.0 + ) + max_abs_diff = max(max_abs_diff, output_diff) + if not np.allclose(original_output, simplified_output, rtol=1e-5, atol=1e-6): + raise RuntimeError( + f"simplified output differs for {path}:{name}: " + f"max absolute difference = {output_diff:.3e}" + ) + + directory = os.path.dirname(path) or "." + original_mode = stat.S_IMODE(os.stat(path).st_mode) + file_descriptor, temporary_path = tempfile.mkstemp( + prefix=f".{os.path.basename(path)}.", suffix=".tmp", dir=directory + ) + os.close(file_descriptor) + try: + onnx.save(simplified, temporary_path) + os.chmod(temporary_path, original_mode) + os.replace(temporary_path, path) + finally: + if os.path.exists(temporary_path): + os.unlink(temporary_path) + + return max_abs_diff + + def main() -> None: models_dir = sys.argv[1] token = os.environ.get("HF_TOKEN") @@ -72,6 +168,8 @@ def export_segmentation(pipeline: Any, models_dir: str) -> None: seg_model.eval() dummy = torch.randn(1, 1, 160000) + dummy_b32 = torch.randn(32, 1, 160000) + dummy_b64 = torch.randn(64, 1, 160000) with torch.no_grad(): torch.onnx.export( seg_model, @@ -85,7 +183,7 @@ def export_segmentation(pipeline: Any, models_dir: str) -> None: ) torch.onnx.export( seg_model, - (torch.randn(32, 1, 160000),), + (dummy_b32,), os.path.join(models_dir, "segmentation-3.0-b32.onnx"), input_names=["input"], output_names=["output"], @@ -94,7 +192,7 @@ def export_segmentation(pipeline: Any, models_dir: str) -> None: ) torch.onnx.export( seg_model, - (torch.randn(64, 1, 160000),), + (dummy_b64,), os.path.join(models_dir, "segmentation-3.0-b64.onnx"), input_names=["input"], output_names=["output"], @@ -102,6 +200,17 @@ def export_segmentation(pipeline: Any, models_dir: str) -> None: dynamo=False, ) + validation_inputs = ( + ("segmentation-3.0.onnx", dummy), + ("segmentation-3.0-b32.onnx", dummy_b32), + ("segmentation-3.0-b64.onnx", dummy_b64), + ) + for filename, input_tensor in validation_inputs: + path = os.path.join(models_dir, filename) + inputs = {"input": input_tensor.detach().cpu().numpy()} + max_abs_diff = simplify_onnx_graph(path, inputs) + print(f" {filename} simplification verified (max diff = {max_abs_diff:.2e})") + sz = os.path.getsize(os.path.join(models_dir, "segmentation-3.0.onnx")) / 1e6 print(f" segmentation-3.0.onnx ({sz:.1f} MB)") bsz = os.path.getsize(os.path.join(models_dir, "segmentation-3.0-b32.onnx")) / 1e6 From 6cbddb0527c075852747450bf4f0fdf919f35061 Mon Sep 17 00:00:00 2001 From: Praveen Perera Date: Wed, 16 Sep 2026 22:02:19 +0000 Subject: [PATCH 2/3] Snapshot original ONNX graph --- scripts/export_models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/scripts/export_models.py b/scripts/export_models.py index ead82b7..dd5e0c8 100644 --- a/scripts/export_models.py +++ b/scripts/export_models.py @@ -42,6 +42,7 @@ def simplify_onnx_graph(path: str, inputs: dict[str, np.ndarray]) -> float: from onnxsim import simplify original = onnx.load(path) + original_bytes = original.SerializeToString() simplified, valid = simplify(original) if not valid: raise RuntimeError( @@ -50,7 +51,7 @@ def simplify_onnx_graph(path: str, inputs: dict[str, np.ndarray]) -> float: onnx.checker.check_model(simplified) original_session = ort.InferenceSession( - original.SerializeToString(), providers=["CPUExecutionProvider"] + original_bytes, providers=["CPUExecutionProvider"] ) simplified_session = ort.InferenceSession( simplified.SerializeToString(), providers=["CPUExecutionProvider"] From 7e85fb75a26e7f0f4d7ad215215f0a3792abc4c9 Mon Sep 17 00:00:00 2001 From: Praveen Perera Date: Wed, 16 Sep 2026 22:10:45 +0000 Subject: [PATCH 3/3] Verify ONNX interface metadata --- scripts/export_models.py | 36 ++++++++++++++++++++++++++---------- 1 file changed, 26 insertions(+), 10 deletions(-) diff --git a/scripts/export_models.py b/scripts/export_models.py index dd5e0c8..324e274 100644 --- a/scripts/export_models.py +++ b/scripts/export_models.py @@ -57,26 +57,42 @@ def simplify_onnx_graph(path: str, inputs: dict[str, np.ndarray]) -> float: simplified.SerializeToString(), providers=["CPUExecutionProvider"] ) - original_input_names = [value.name for value in original_session.get_inputs()] - simplified_input_names = [value.name for value in simplified_session.get_inputs()] - if original_input_names != simplified_input_names: + original_inputs = original_session.get_inputs() + simplified_inputs = simplified_session.get_inputs() + original_input_metadata = [ + (value.name, value.type, tuple(value.shape)) for value in original_inputs + ] + simplified_input_metadata = [ + (value.name, value.type, tuple(value.shape)) for value in simplified_inputs + ] + if original_input_metadata != simplified_input_metadata: raise RuntimeError( - f"simplified input names changed for {path}: " - f"{original_input_names} != {simplified_input_names}" + f"simplified input metadata changed for {path}: " + f"{original_input_metadata} != {simplified_input_metadata}" ) + original_input_names = [value.name for value in original_inputs] if set(inputs) != set(original_input_names): raise RuntimeError( f"validation inputs do not match {path}: " f"{sorted(inputs)} != {sorted(original_input_names)}" ) - original_output_names = [value.name for value in original_session.get_outputs()] - simplified_output_names = [value.name for value in simplified_session.get_outputs()] - if original_output_names != simplified_output_names: + original_outputs_metadata = original_session.get_outputs() + simplified_outputs_metadata = simplified_session.get_outputs() + original_output_metadata = [ + (value.name, value.type, tuple(value.shape)) + for value in original_outputs_metadata + ] + simplified_output_metadata = [ + (value.name, value.type, tuple(value.shape)) + for value in simplified_outputs_metadata + ] + if original_output_metadata != simplified_output_metadata: raise RuntimeError( - f"simplified output names changed for {path}: " - f"{original_output_names} != {simplified_output_names}" + f"simplified output metadata changed for {path}: " + f"{original_output_metadata} != {simplified_output_metadata}" ) + original_output_names = [value.name for value in original_outputs_metadata] original_outputs = original_session.run(None, inputs) simplified_outputs = simplified_session.run(None, inputs)