diff --git a/examples/arm/QAT_example/qat_loop.py b/examples/arm/QAT_example/qat_loop.py index faca0a20f1a..cb91ff9748c 100644 --- a/examples/arm/QAT_example/qat_loop.py +++ b/examples/arm/QAT_example/qat_loop.py @@ -30,7 +30,7 @@ from collections import Counter from dataclasses import dataclass from pathlib import Path -from typing import Any, cast, Iterable +from typing import Any, cast, Iterable, Sequence import numpy as np import torch @@ -57,6 +57,10 @@ ) from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner # noqa: E402 from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower # noqa: E402 +from executorch.exir.passes.quantize_io_pass import ( # noqa: E402 + quantize_input, + quantize_output, +) from executorch.extension.export_util.utils import save_pte_program # noqa: E402 DEFAULT_WIDTH = 768 @@ -80,6 +84,15 @@ class QuantizedModel: coverage: dict[str, Any] +@dataclass(frozen=True) +class ExportedArtifact: + name: str + artifact_dir: Path + pte_path: Path | None + io_qparams_path: Path | None + io_qparams_payload: dict[str, Any] | None + + class RIFEWrapper(torch.nn.Module): def __init__(self, flownet: torch.nn.Module) -> None: super().__init__() @@ -672,6 +685,17 @@ def make_vgf_compile_spec(output_dir: Path) -> VgfCompileSpec: return compile_spec +def quant_args_to_dict(quant_args: Sequence[Any]) -> dict[str, Any]: + scale, zero_point, qmin, qmax, dtype = quant_args + return { + "scale": float(scale), + "zero_point": int(zero_point), + "qmin": int(qmin), + "qmax": int(qmax), + "dtype": str(dtype), + } + + def build_ptq_model( model: torch.nn.Module, calibration_inputs: list[tuple[torch.Tensor, torch.Tensor]], @@ -770,8 +794,9 @@ def export_vgf_artifact( *, output_dir: Path, name: str, + io_quantization: str, save_pte: bool, -) -> Path: +) -> ExportedArtifact: artifact_dir = output_dir / f"vgf_{name}" artifact_dir.mkdir(parents=True, exist_ok=True) compile_spec = make_vgf_compile_spec(artifact_dir) @@ -781,6 +806,12 @@ def export_vgf_artifact( inputs[0].contiguous(memory_format=torch.channels_last), inputs[1].contiguous(memory_format=torch.channels_last), ) + pte_path = artifact_dir / f"rife_{name}_vgf.pte" if save_pte else None + io_qparams_path = ( + artifact_dir / f"rife_{name}_vgf_io_qparams.json" + if io_quantization != "none" + else None + ) print(f"[vgf] exporting {name} model to {artifact_dir} ...") aten_dialect = export(model, args=inputs, strict=True) @@ -789,11 +820,43 @@ def export_vgf_artifact( partitioner=[partitioner], compile_config=EdgeCompileConfig(_check_ir_validity=False), ) + io_qparams_payload: dict[str, Any] | None = None + if io_quantization != "none": + edge_program = edge_program_manager.exported_program() + input_qparams = [ + quantize_input(edge_program, input_index) + for input_index in range(len(inputs)) + ] + output_qparams = [quantize_output(edge_program, 0)] + edge_program.graph_module.graph.eliminate_dead_code() + edge_program.graph_module.recompile() + io_qparams_payload = { + "input_qparams": [ + quant_args_to_dict(quant_args) for quant_args in input_qparams + ], + "output_qparams": [ + quant_args_to_dict(quant_args) for quant_args in output_qparams + ], + } + if io_qparams_path is None: + raise RuntimeError("Expected an IO qparams output path") + write_json(io_qparams_path, io_qparams_payload) if save_pte: - pte_path = artifact_dir / f"rife_{name}_vgf.pte" + if pte_path is None: + raise RuntimeError("Expected a PTE output path") save_pte_program(edge_program_manager.to_executorch(), str(pte_path)) print(f"[vgf] wrote {pte_path}") - return artifact_dir + if io_qparams_payload is not None: + if io_qparams_path is None: + raise RuntimeError("Expected an IO qparams output path") + print(f"[vgf] wrote {io_qparams_path}") + return ExportedArtifact( + name=name, + artifact_dir=artifact_dir, + pte_path=pte_path, + io_qparams_path=io_qparams_path, + io_qparams_payload=io_qparams_payload, + ) def write_json(path: Path, payload: dict[str, Any]) -> None: @@ -1175,9 +1238,10 @@ def export_requested_vgf_artifacts( triples: list[FrameTriple], ptq_model: torch.nn.Module | None, qat_model: torch.nn.Module | None, -) -> None: +) -> dict[str, ExportedArtifact]: + artifacts: dict[str, ExportedArtifact] = {} if args.export_vgf == "none": - return + return artifacts inputs, _target = load_triple_tensors( triples[0], @@ -1189,23 +1253,26 @@ def export_requested_vgf_artifacts( if args.export_vgf in ("ptq", "all"): if ptq_model is None: raise RuntimeError("PTQ VGF export requested, but PTQ was not built.") - export_vgf_artifact( + artifacts["ptq"] = export_vgf_artifact( ptq_model, inputs, output_dir=args.output_dir, name="ptq", + io_quantization=args.io_quantization, save_pte=args.save_pte, ) if args.export_vgf in ("qat", "all"): if qat_model is None: raise RuntimeError("QAT VGF export requested, but QAT was not built.") - export_vgf_artifact( + artifacts["qat"] = export_vgf_artifact( qat_model, inputs, output_dir=args.output_dir, name="qat", + io_quantization=args.io_quantization, save_pte=args.save_pte, ) + return artifacts def update_row_with_prediction_metrics( @@ -1230,6 +1297,247 @@ def update_row_with_delta_metrics( row[f"{prefix}_{key}"] = value +def qparam_dtype(dtype_text: str) -> torch.dtype: + if dtype_text == "torch.int8": + return torch.int8 + if dtype_text == "torch.uint8": + return torch.uint8 + raise ValueError(f"Unsupported runtime qparam dtype: {dtype_text}") + + +def quantize_tensor_for_runtime( + tensor: torch.Tensor, + quant_args: dict[str, Any], +) -> torch.Tensor: + scale = float(quant_args["scale"]) + zero_point = int(quant_args["zero_point"]) + qmin = int(quant_args["qmin"]) + qmax = int(quant_args["qmax"]) + dtype = qparam_dtype(str(quant_args["dtype"])) + quantized = torch.round(tensor / scale + zero_point) + quantized = quantized.clamp(qmin, qmax).to(dtype) + return quantized.contiguous(memory_format=torch.channels_last) + + +def dequantize_tensor_from_runtime( + tensor: torch.Tensor, + quant_args: dict[str, Any], +) -> torch.Tensor: + scale = float(quant_args["scale"]) + zero_point = int(quant_args["zero_point"]) + dequantized = (tensor.to(torch.float32) - zero_point) * scale + return dequantized.contiguous(memory_format=torch.channels_last) + + +def try_get_runtime() -> tuple[Any | None, list[str]]: + warnings: list[str] = [] + try: + from executorch.runtime import Runtime + except Exception as error: + warnings.append( + "ExecuTorch runtime Python bindings are unavailable; " + f"skipping runtime numerics ({type(error).__name__}: {error})." + ) + return None, warnings + + try: + runtime = Runtime.get() + except Exception as error: + warnings.append( + "ExecuTorch runtime initialization failed; " + f"skipping runtime numerics ({type(error).__name__}: {error})." + ) + return None, warnings + + if not runtime.backend_registry.is_available("VgfBackend"): + warnings.append( + "VgfBackend is not available in this ExecuTorch runtime build; " + "skipping runtime numerics." + ) + return None, warnings + return runtime, warnings + + +def execute_runtime_artifact( + method: Any, + input_qparams: list[dict[str, Any]] | None, + output_qparams: list[dict[str, Any]] | None, + inputs: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + runtime_inputs: tuple[torch.Tensor, torch.Tensor] + if input_qparams is None: + runtime_inputs = inputs + else: + runtime_inputs = ( + quantize_tensor_for_runtime(inputs[0], input_qparams[0]), + quantize_tensor_for_runtime(inputs[1], input_qparams[1]), + ) + outputs = method.execute(runtime_inputs) + if not outputs: + raise RuntimeError("Runtime returned no outputs") + output = outputs[0] + if not isinstance(output, torch.Tensor): + output = torch.as_tensor(output) + if output_qparams is None: + return output + return dequantize_tensor_from_runtime(output, output_qparams[0]) + + +def dump_runtime_comparison_outputs( + output_dir: Path, + mode_name: str, + row_name: str, + reference_output: torch.Tensor, + runtime_output: torch.Tensor, +) -> None: + safe_row_name = row_name.replace("..", "__").replace("/", "_").replace("\\", "_") + dump_dir = output_dir / "runtime_debug" / mode_name / safe_row_name + dump_dir.mkdir(parents=True, exist_ok=True) + reference_cpu = reference_output.detach().cpu() + runtime_cpu = runtime_output.detach().cpu() + diff_cpu = runtime_cpu - reference_cpu + np.save(dump_dir / f"{mode_name}.npy", reference_cpu.numpy()) + np.save(dump_dir / f"{mode_name}_runtime.npy", runtime_cpu.numpy()) + np.save(dump_dir / f"{mode_name}_runtime_minus_{mode_name}.npy", diff_cpu.numpy()) + + +def load_runtime_method_for_artifact( + runtime: Any, + mode_name: str, + artifact: ExportedArtifact, +) -> tuple[Any, list[dict[str, Any]] | None, list[dict[str, Any]] | None] | str: + if artifact.pte_path is None: + return f"No .pte was written for {mode_name}; skipping runtime numerics." + try: + program = runtime.load_program(str(artifact.pte_path)) + method = program.load_method("forward") + if method is None: + raise RuntimeError(f"No forward method found in {artifact.pte_path}") + except Exception as error: + return ( + f"{mode_name} runtime initialization failed; skipping runtime numerics " + f"({type(error).__name__}: {error})." + ) + return ( + method, + ( + None + if artifact.io_qparams_payload is None + else artifact.io_qparams_payload["input_qparams"] + ), + ( + None + if artifact.io_qparams_payload is None + else artifact.io_qparams_payload["output_qparams"] + ), + ) + + +def evaluate_runtime_artifact_rows( + args: argparse.Namespace, + triples: list[FrameTriple], + rows: list[dict[str, Any]], + mode_name: str, + prefix: str, + method: Any, + input_qparams: list[dict[str, Any]] | None, + output_qparams: list[dict[str, Any]] | None, + reference_model: torch.nn.Module | None, +) -> str | None: + for index, triple in enumerate(triples): + inputs, target = load_triple_tensors( + triple, + height=args.height, + width=args.width, + preprocess=args.preprocess, + random_seed=args.random_seed, + ) + try: + runtime_output = execute_runtime_artifact( + method, + input_qparams, + output_qparams, + inputs, + ) + except Exception as error: + for processed_row in rows[:index]: + for key in list(processed_row): + if key.startswith(f"{prefix}_"): + del processed_row[key] + mode_name = prefix.removesuffix("_runtime") + return ( + f"{mode_name} runtime execution failed; skipping remaining runtime numerics " + f"({type(error).__name__}: {error})." + ) + update_row_with_prediction_metrics( + rows[index], + prefix=prefix, + prediction=runtime_output, + target=target, + ) + if reference_model is not None: + reference_output = run_model(reference_model, inputs) + update_row_with_delta_metrics( + rows[index], + prefix=f"{prefix}_vs_{mode_name}", + prediction=runtime_output, + reference=reference_output, + ) + dump_runtime_comparison_outputs( + args.output_dir, + mode_name, + rows[index]["name"], + reference_output, + runtime_output, + ) + return None + + +def evaluate_runtime_artifacts( + args: argparse.Namespace, + triples: list[FrameTriple], + artifacts: dict[str, ExportedArtifact], + rows: list[dict[str, Any]], + reference_models: dict[str, torch.nn.Module], +) -> list[str]: + if not artifacts: + return [] + runtime, warnings = try_get_runtime() + if runtime is None: + for warning in warnings: + print(f"[runtime] warning: {warning}") + return warnings + + for mode_name, artifact in sorted(artifacts.items()): + prefix = f"{mode_name}_runtime" + runtime_artifact = load_runtime_method_for_artifact( + runtime, mode_name, artifact + ) + if isinstance(runtime_artifact, str): + warnings.append(runtime_artifact) + continue + method, input_qparams, output_qparams = runtime_artifact + execution_warning = evaluate_runtime_artifact_rows( + args, + triples, + rows, + mode_name, + prefix, + method, + input_qparams, + output_qparams, + reference_models.get(mode_name), + ) + if execution_warning is None: + continue + warnings.append(execution_warning) + continue + + for warning in warnings: + print(f"[runtime] warning: {warning}") + return warnings + + def print_progress(index: int, total: int, row: dict[str, Any]) -> None: progress = f"{index + 1}/{total} {row['name']}" if "eager_psnr" in row: @@ -1240,6 +1548,10 @@ def print_progress(index: int, total: int, row: dict[str, Any]) -> None: progress += f" ptq_psnr={row['ptq_psnr']:.3f}" if "qat_psnr" in row: progress += f" qat_psnr={row['qat_psnr']:.3f}" + if "ptq_runtime_psnr" in row: + progress += f" ptq_runtime_psnr={row['ptq_runtime_psnr']:.3f}" + if "qat_runtime_psnr" in row: + progress += f" qat_runtime_psnr={row['qat_runtime_psnr']:.3f}" print(progress) @@ -1338,6 +1650,7 @@ def evaluate_triples( def build_aggregate( args: argparse.Namespace, rows: list[dict[str, Any]], + runtime_warnings: list[str], ) -> dict[str, Any]: aggregate: dict[str, Any] = { "triple_count": len(rows), @@ -1348,6 +1661,7 @@ def build_aggregate( "io_quantization": args.io_quantization, "data_source": "random" if rows and rows[0].get("synthetic") else "frames", "random_seed": args.random_seed, + "runtime_warnings": runtime_warnings, } for prefix in ( "eager", @@ -1357,6 +1671,10 @@ def build_aggregate( "ptq_vs_eager", "qat", "qat_vs_eager", + "ptq_runtime", + "ptq_runtime_vs_ptq", + "qat_runtime", + "qat_runtime_vs_qat", ): aggregate.update(aggregate_metrics(rows, prefix)) return aggregate @@ -1401,7 +1719,9 @@ def print_summary(aggregate: dict[str, Any]) -> None: ("eager", "eager"), ("eager_channels_last", "eager_channels_last"), ("ptq", "ptq"), + ("ptq_runtime", "ptq_runtime"), ("qat", "qat"), + ("qat_runtime", "qat_runtime"), ): if f"{prefix}_psnr_mean" not in aggregate: continue @@ -1410,6 +1730,26 @@ def print_summary(aggregate: dict[str, Any]) -> None: f"psnr={format_metric(aggregate.get(f'{prefix}_psnr_mean'))} " f"ssim={format_metric(aggregate.get(f'{prefix}_ssim_mean'))}" ) + if prefix == "ptq_runtime": + mae_key = "ptq_runtime_vs_ptq_mae_mean" + sqnr_key = "ptq_runtime_vs_ptq_sqnr_mean" + if mae_key in aggregate: + summary += ( + f" vs_ptq_mae={format_metric(aggregate.get(mae_key))}" + f" vs_ptq_sqnr={format_metric(aggregate.get(sqnr_key))}" + ) + print(summary) + continue + if prefix == "qat_runtime": + mae_key = "qat_runtime_vs_qat_mae_mean" + sqnr_key = "qat_runtime_vs_qat_sqnr_mean" + if mae_key in aggregate: + summary += ( + f" vs_qat_mae={format_metric(aggregate.get(mae_key))}" + f" vs_qat_sqnr={format_metric(aggregate.get(sqnr_key))}" + ) + print(summary) + continue mae_key = f"{prefix}_vs_eager_mae_mean" sqnr_key = f"{prefix}_vs_eager_sqnr_mean" if mae_key in aggregate: @@ -1418,6 +1758,11 @@ def print_summary(aggregate: dict[str, Any]) -> None: f" vs_eager_sqnr={format_metric(aggregate.get(sqnr_key))}" ) print(summary) + runtime_warnings = cast(list[str], aggregate.get("runtime_warnings", [])) + if runtime_warnings: + print("runtime_warnings:") + for warning in runtime_warnings: + print(f" - {warning}") def main() -> int: @@ -1450,7 +1795,20 @@ def main() -> int: ptq_model, qat_model, ) - aggregate = build_aggregate(args, rows) + artifacts = export_requested_vgf_artifacts(args, triples, ptq_model, qat_model) + reference_models = { + name: reference_model + for name, reference_model in (("ptq", ptq_model), ("qat", qat_model)) + if reference_model is not None + } + runtime_warnings = evaluate_runtime_artifacts( + args, + triples, + artifacts, + rows, + reference_models, + ) + aggregate = build_aggregate(args, rows, runtime_warnings) write_outputs( args, aggregate,