From b2b6c9f8a0da09f6a52ba4c2d3c96bec43f54c49 Mon Sep 17 00:00:00 2001 From: Baris Demir Date: Fri, 26 Jun 2026 09:41:58 +0100 Subject: [PATCH] Arm backend: Add RIFE PTQ/QAT accuracy flow Add a Practical-RIFE PTQ/QAT comparison script under the Arm QAT example area. The flow loads original RIFE weights, uses real frame triples, preserves uint8 image IO by default, and writes JSON, CSV, Markdown, and quantization coverage reports. If a requested triples list is missing, create deterministic random data for smoke testing. The reports clearly label this as smoke-test data, not real accuracy data. Signed-off-by: Baris Demir Change-Id: Ia60c2a43a8f600405fda1b782965f560c69f945b --- examples/arm/QAT_example/qat_loop.py | 1334 ++++++++++++++++++++++++++ 1 file changed, 1334 insertions(+) create mode 100644 examples/arm/QAT_example/qat_loop.py diff --git a/examples/arm/QAT_example/qat_loop.py b/examples/arm/QAT_example/qat_loop.py new file mode 100644 index 00000000000..83ad484a4fe --- /dev/null +++ b/examples/arm/QAT_example/qat_loop.py @@ -0,0 +1,1334 @@ +#!/usr/bin/env python3 +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Measure Practical-RIFE PTQ and QAT accuracy on real frame triples. + +The script is intended as a small handover flow for comparing eager FP32, +PTQ, and QAT PT2E models before generating customer VGF artifacts. The +expected triple order is: + +* first frame: model input 0 +* third frame: model input 1 +* middle frame: quality target + +By default the fixed input shape is 768x384, matching the OPPO request. +""" + +from __future__ import annotations + +import argparse +import csv +import importlib +import json +import math +import re +import sys +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Any, cast, Iterable + +import numpy as np +import torch +import torch.nn.functional as F +from PIL import Image +from torch.export import export +from torchao.quantization.pt2e import ( + move_exported_model_to_eval, + move_exported_model_to_train, +) +from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e, + prepare_pt2e, + prepare_qat_pt2e, +) + +EXECUTORCH_ROOT = Path(__file__).resolve().parents[3] +sys.path.insert(0, str(EXECUTORCH_ROOT)) + +from executorch.backends.arm.quantizer import ( # noqa: E402 + get_symmetric_quantization_config, + get_uint8_io_quantization_config, + VgfQuantizer, +) +from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner # noqa: E402 +from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower # noqa: E402 +from executorch.extension.export_util.utils import save_pte_program # noqa: E402 + +DEFAULT_WIDTH = 768 +DEFAULT_HEIGHT = 384 +RIFE_SHAPE_ALIGNMENT = 64 +IMAGE_SUFFIXES = {".bmp", ".jpeg", ".jpg", ".png", ".webp"} + + +@dataclass(frozen=True) +class FrameTriple: + name: str + input0: Path | None + input1: Path | None + target: Path | None + synthetic: bool = False + + +@dataclass(frozen=True) +class QuantizedModel: + model: torch.nn.Module + coverage: dict[str, Any] + + +class RIFEWrapper(torch.nn.Module): + def __init__(self, flownet: torch.nn.Module) -> None: + super().__init__() + self.flownet = flownet + + def forward(self, img0: torch.Tensor, img1: torch.Tensor) -> torch.Tensor: + imgs = torch.cat((img0, img1), dim=1) + _, _, merged = self.flownet(imgs, 0.5, [16, 8, 4, 2, 1]) + return merged[-1] + + +def validate_input_shape(height: int, width: int) -> None: + if height <= 0 or width <= 0: + raise ValueError("--height and --width must be positive.") + if height % RIFE_SHAPE_ALIGNMENT != 0 or width % RIFE_SHAPE_ALIGNMENT != 0: + raise ValueError( + "RIFE inputs must be multiples of " + f"{RIFE_SHAPE_ALIGNMENT}. Got height={height}, width={width}." + ) + + +def load_rife_model(model_root: Path, checkpoint: Path | None) -> torch.nn.Module: + if not model_root.exists(): + raise FileNotFoundError(f"Practical-RIFE repo not found: {model_root}") + checkpoint = checkpoint or model_root / "train_log" / "flownet.pkl" + sys.path.insert(0, str(model_root)) + ifnet_module = importlib.import_module("train_log.IFNet_HDv3") + ifnet = cast(Any, ifnet_module) + + if not checkpoint.exists(): + raise FileNotFoundError( + "RIFE checkpoint was not found. Pass --checkpoint or place " + f"flownet.pkl at {checkpoint}." + ) + + checkpoint_data = torch.load( + checkpoint, + map_location=torch.device("cpu"), + weights_only=True, + ) + if not isinstance(checkpoint_data, dict): + raise TypeError(f"Unexpected RIFE checkpoint type: {type(checkpoint_data)}") + + state_dict = { + key.removeprefix("module."): value for key, value in checkpoint_data.items() + } + if hasattr(ifnet, "backwarp_tenGrid"): + ifnet.backwarp_tenGrid.clear() + flownet = ifnet.IFNet() + incompatible_keys = flownet.load_state_dict(state_dict, strict=False) + if incompatible_keys.missing_keys: + raise RuntimeError( + f"Missing RIFE checkpoint keys: {incompatible_keys.missing_keys}" + ) + return RIFEWrapper(flownet).eval() + + +def resolve_path(path_text: str, base_dir: Path) -> Path: + path = Path(path_text) + if not path.is_absolute(): + path = base_dir / path + return path + + +def split_triple_line(line: str) -> list[str]: + return [part for part in re.split(r"[\s,]+", line.strip()) if part] + + +def triple_list_has_header(first_line: str) -> bool: + fields = {part.lower() for part in split_triple_line(first_line)} + return {"input0", "input1", "target"}.issubset(fields) + + +def load_triples_list(path: Path) -> list[FrameTriple]: + triples: list[FrameTriple] = [] + base_dir = path.parent + with path.open(encoding="utf-8", newline="") as file: + first_line = file.readline() + file.seek(0) + if triple_list_has_header(first_line): + reader = csv.DictReader(file) + for index, row in enumerate(reader): + if row is None: + continue + triples.append( + FrameTriple( + name=row.get("name") or f"triple_{index:05d}", + input0=resolve_path(row["input0"], base_dir), + input1=resolve_path(row["input1"], base_dir), + target=resolve_path(row["target"], base_dir), + ) + ) + return triples + + for line_number, line in enumerate(file, start=1): + if not line.strip() or line.lstrip().startswith("#"): + continue + parts = split_triple_line(line) + if len(parts) not in (3, 4): + raise ValueError( + "Triple lines must contain 'input0 input1 target' or " + "'name input0 input1 target'. " + f"Bad line {line_number}: {line.rstrip()}" + ) + if len(parts) == 3: + name = f"triple_{len(triples):05d}" + input0, input1, target = parts + else: + name, input0, input1, target = parts + triples.append( + FrameTriple( + name=name, + input0=resolve_path(input0, base_dir), + input1=resolve_path(input1, base_dir), + target=resolve_path(target, base_dir), + ) + ) + return triples + + +def image_files(directory: Path) -> list[Path]: + return sorted( + path + for path in directory.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ) + + +def discover_triples(root: Path, stride: int) -> list[FrameTriple]: + if stride <= 0: + raise ValueError("--sequence-stride must be positive.") + + triples: list[FrameTriple] = [] + directories = [root] + sorted(path for path in root.rglob("*") if path.is_dir()) + for directory in directories: + images = image_files(directory) + if len(images) < 3: + continue + relative = directory.relative_to(root) if directory != root else Path("root") + prefix = "__".join(relative.parts) + for index in range(0, len(images) - 2, stride): + triples.append( + FrameTriple( + name=f"{prefix}_{index:05d}", + input0=images[index], + input1=images[index + 2], + target=images[index + 1], + ) + ) + return triples + + +def make_random_triples(count: int) -> list[FrameTriple]: + return [ + FrameTriple( + name=f"random_{index:05d}", + input0=None, + input1=None, + target=None, + synthetic=True, + ) + for index in range(count) + ] + + +def validate_triples(triples: Iterable[FrameTriple]) -> list[FrameTriple]: + validated = list(triples) + if not validated: + raise ValueError("No frame triples were found.") + missing = [ + str(path) + for triple in validated + for path in (triple.input0, triple.input1, triple.target) + if path is not None and not path.exists() + ] + if missing: + raise FileNotFoundError("Missing frame files:\n" + "\n".join(missing[:20])) + return validated + + +def load_image(path: Path) -> torch.Tensor: + image = Image.open(path).convert("RGB") + array = np.asarray(image, dtype=np.float32) / 255.0 + return torch.from_numpy(array).permute(2, 0, 1).contiguous() + + +def preprocess_image( + image: torch.Tensor, + *, + height: int, + width: int, + mode: str, +) -> torch.Tensor: + current_height, current_width = image.shape[-2:] + if current_height == height and current_width == width: + return image + + if mode == "center-crop": + if current_height < height or current_width < width: + raise ValueError( + "Image is smaller than the requested crop. Use " + "--preprocess resize or provide larger frames." + ) + top = (current_height - height) // 2 + left = (current_width - width) // 2 + return image[:, top : top + height, left : left + width].contiguous() + + if mode == "resize": + resized = F.interpolate( + image.unsqueeze(0), + size=(height, width), + mode="bilinear", + align_corners=False, + ) + return resized.squeeze(0).contiguous() + + raise ValueError(f"Unsupported preprocess mode: {mode}") + + +def load_triple_tensors( + triple: FrameTriple, + *, + height: int, + width: int, + preprocess: str, + random_seed: int, +) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]: + if triple.synthetic: + index_text = triple.name.rsplit("_", maxsplit=1)[-1] + generator = torch.Generator().manual_seed(random_seed + int(index_text)) + input0 = torch.rand((1, 3, height, width), generator=generator) + input1 = torch.rand((1, 3, height, width), generator=generator) + target = torch.rand((1, 3, height, width), generator=generator) + return (input0, input1), target + + if triple.input0 is None or triple.input1 is None or triple.target is None: + raise ValueError(f"Frame triple is missing paths: {triple.name}") + + frames = [ + preprocess_image( + load_image(path), + height=height, + width=width, + mode=preprocess, + ) + for path in (triple.input0, triple.input1, triple.target) + ] + input0, input1, target = [frame.unsqueeze(0) for frame in frames] + return (input0, input1), target + + +def gaussian_window( + channels: int, + device: torch.device, + dtype: torch.dtype, + height: int, + width: int, +) -> torch.Tensor: + window_size = min(11, height, width) + if window_size % 2 == 0: + window_size -= 1 + window_size = max(window_size, 1) + sigma = max(window_size / 6.0, 1e-3) + + coords = torch.arange(window_size, device=device, dtype=dtype) - window_size // 2 + kernel_1d = torch.exp(-(coords**2) / (2 * sigma**2)) + kernel_1d = kernel_1d / kernel_1d.sum() + kernel_2d = torch.outer(kernel_1d, kernel_1d) + return kernel_2d.expand(channels, 1, window_size, window_size).contiguous() + + +def batch_psnr(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + mse = F.mse_loss(prediction, target, reduction="none").mean(dim=(1, 2, 3)) + return 10.0 * torch.log10(1.0 / mse.clamp_min(1e-12)) + + +def batch_ssim(prediction: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + channels = prediction.shape[1] + kernel = gaussian_window( + channels, + prediction.device, + prediction.dtype, + prediction.shape[-2], + prediction.shape[-1], + ) + padding = kernel.shape[-1] // 2 + c1 = 0.01**2 + c2 = 0.03**2 + + mu_x = F.conv2d(prediction, kernel, padding=padding, groups=channels) + mu_y = F.conv2d(target, kernel, padding=padding, groups=channels) + mu_x_sq = mu_x.pow(2) + mu_y_sq = mu_y.pow(2) + mu_xy = mu_x * mu_y + sigma_x_sq = ( + F.conv2d(prediction * prediction, kernel, padding=padding, groups=channels) + - mu_x_sq + ) + sigma_y_sq = ( + F.conv2d(target * target, kernel, padding=padding, groups=channels) - mu_y_sq + ) + sigma_xy = ( + F.conv2d(prediction * target, kernel, padding=padding, groups=channels) - mu_xy + ) + + ssim_map = ((2 * mu_xy + c1) * (2 * sigma_xy + c2)) / ( + (mu_x_sq + mu_y_sq + c1) * (sigma_x_sq + sigma_y_sq + c2) + ) + return ssim_map.mean(dim=(1, 2, 3)) + + +def tensor_metrics( + prediction: torch.Tensor, reference: torch.Tensor +) -> dict[str, float]: + prediction = prediction.detach().cpu().to(torch.float64) + reference = reference.detach().cpu().to(torch.float64) + error = prediction - reference + abs_error = error.abs() + mse = error.pow(2).mean().item() + reference_norm = reference.norm().item() + error_norm = error.norm().item() + prediction_norm = prediction.norm().item() + + if reference_norm > 0.0 and error_norm > 0.0: + sqnr = 20.0 * math.log10(reference_norm / error_norm) + elif error_norm == 0.0: + sqnr = float("inf") + else: + sqnr = float("-inf") + + if reference_norm > 0.0 and prediction_norm > 0.0: + cosine = torch.dot(reference.flatten(), prediction.flatten()).item() / ( + reference_norm * prediction_norm + ) + else: + cosine = float("nan") + + return { + "mae": float(abs_error.mean().item()), + "max_abs": float(abs_error.max().item()), + "mse": float(mse), + "rmse": float(math.sqrt(mse)), + "sqnr": float(sqnr), + "cosine": float(cosine), + } + + +def quality_metrics( + prediction: torch.Tensor, + target: torch.Tensor, +) -> dict[str, float]: + clamped = prediction.detach().cpu().clamp(0.0, 1.0) + target = target.detach().cpu() + metrics = tensor_metrics(prediction, target) + metrics.update( + { + "psnr": float(batch_psnr(clamped, target).mean().item()), + "ssim": float(batch_ssim(clamped, target).mean().item()), + "output_min": float(prediction.detach().cpu().min().item()), + "output_max": float(prediction.detach().cpu().max().item()), + } + ) + return metrics + + +def mean(values: list[float]) -> float: + finite = [value for value in values if math.isfinite(value)] + if not finite: + return float("nan") + return float(sum(finite) / len(finite)) + + +def aggregate_metrics(rows: list[dict[str, Any]], prefix: str) -> dict[str, float]: + aggregate = {} + for metric in ( + "mae", + "max_abs", + "mse", + "rmse", + "sqnr", + "cosine", + "psnr", + "ssim", + "output_min", + "output_max", + ): + values = [ + float(row[f"{prefix}_{metric}"]) + for row in rows + if f"{prefix}_{metric}" in row and row[f"{prefix}_{metric}"] is not None + ] + if values: + aggregate[f"{prefix}_{metric}_mean"] = mean(values) + aggregate[f"{prefix}_{metric}_min"] = float(min(values)) + aggregate[f"{prefix}_{metric}_max"] = float(max(values)) + return aggregate + + +def format_metric(value: Any) -> str: + if value is None: + return "n/a" + if isinstance(value, float): + if math.isnan(value): + return "nan" + if math.isinf(value): + return "inf" if value > 0 else "-inf" + return f"{value:.6g}" + return str(value) + + +def run_model( + model: torch.nn.Module, + inputs: tuple[torch.Tensor, torch.Tensor], +) -> torch.Tensor: + with torch.no_grad(): + output = model(*inputs) + if not isinstance(output, torch.Tensor): + raise TypeError(f"Expected tensor output, got {type(output)}") + return output + + +def target_name(target: Any) -> str: + if isinstance(target, str): + return target + name = getattr(target, "__name__", None) + if name: + return str(name) + return str(target) + + +def node_arg_names(node: torch.fx.Node) -> list[str]: + names: list[str] = [] + + def collect(value: Any) -> Any: + if isinstance(value, torch.fx.Node): + names.append(value.name) + return value + + torch.fx.node.map_arg((node.args, node.kwargs), collect) + return names + + +def is_quantize_node(name: str) -> bool: + return name.startswith(("quantize_per_tensor", "quantize_per_channel")) + + +def is_dequantize_node(name: str) -> bool: + return name.startswith(("dequantize_per_tensor", "dequantize_per_channel")) + + +def is_quantized_op_name(name: str) -> bool: + return ( + "quantized_decomposed" in name + or name.startswith("quantized.") + or ".quantized_" in name + ) + + +def quantization_coverage(graph_module: torch.nn.Module) -> dict[str, Any]: + graph = getattr(graph_module, "graph", None) + if graph is None: + return { + "status": "unavailable", + "reason": f"Model has no FX graph: {type(graph_module)}", + } + + quantized_value_names: set[str] = set() + rows: list[dict[str, Any]] = [] + status_counter: Counter[str] = Counter() + op_counter: Counter[tuple[str, str]] = Counter() + + for node in graph.nodes: + name = target_name(node.target) + input_names = node_arg_names(node) + quantized_inputs = [ + input_name + for input_name in input_names + if input_name in quantized_value_names + ] + status = "not_compute" + output_quantized = False + + if node.op == "placeholder": + status = "graph_input_fp32" + elif node.op == "get_attr": + status = "parameter_or_buffer" + elif node.op == "output": + status = "graph_output" + elif is_quantize_node(name): + status = "quantize_boundary" + output_quantized = True + elif is_dequantize_node(name): + status = "dequantize_boundary" + elif is_quantized_op_name(name): + status = "quantized_op" + output_quantized = True + elif quantized_inputs: + status = "quantized_context_op" + output_quantized = True + elif node.op == "call_function": + status = "fp32_or_fallback_op" + + if output_quantized: + quantized_value_names.add(node.name) + + rows.append( + { + "index": len(rows), + "name": node.name, + "op": node.op, + "target": name, + "status": status, + "input_nodes": input_names, + "quantized_input_nodes": quantized_inputs, + "output_quantized": output_quantized, + } + ) + status_counter[status] += 1 + op_counter[(name, status)] += 1 + + op_rows = [ + {"target": target, "status": status, "count": count} + for (target, status), count in sorted( + op_counter.items(), key=lambda item: (-item[1], item[0]) + ) + ] + return { + "status": "ok", + "summary": { + "node_count": len(rows), + "status_counts": dict(status_counter), + "unique_op_status_count": len(op_rows), + }, + "op_status_counts": op_rows, + "nodes": rows, + } + + +def make_quantizer( + *, + is_qat: bool, + io_quantization: str, +) -> VgfQuantizer: + compile_spec = cast( + VgfCompileSpec, VgfCompileSpec()._set_preserve_io_quantization(True) + ) + quantizer = VgfQuantizer(compile_spec, use_composable_quantizer=True) + global_config = get_symmetric_quantization_config(is_qat=is_qat) + quantizer.set_global(global_config) + + if io_quantization == "int8": + quantizer.set_io(global_config) + elif io_quantization == "uint8": + quantizer.set_io(get_uint8_io_quantization_config(is_qat=is_qat)) + elif io_quantization == "none": + quantizer.set_io(None) + else: + raise ValueError(f"Unsupported IO quantization: {io_quantization}") + return quantizer + + +def make_vgf_compile_spec(output_dir: Path) -> VgfCompileSpec: + compile_spec = cast( + VgfCompileSpec, VgfCompileSpec()._set_preserve_io_quantization(True) + ) + compile_spec.dump_intermediate_artifacts_to(str(output_dir)) + return compile_spec + + +def build_ptq_model( + model: torch.nn.Module, + calibration_inputs: list[tuple[torch.Tensor, torch.Tensor]], + *, + io_quantization: str, +) -> QuantizedModel: + exported_model = export(model, calibration_inputs[0], strict=True).module( + check_guards=False + ) + quantizer = make_quantizer(is_qat=False, io_quantization=io_quantization) + prepared_model = prepare_pt2e(exported_model, quantizer) + prepared_model = move_exported_model_to_eval(prepared_model) + + with torch.no_grad(): + for inputs in calibration_inputs: + prepared_model(*inputs) + + converted_model = convert_pt2e(prepared_model) + return QuantizedModel( + model=converted_model, + coverage=quantization_coverage(converted_model), + ) + + +def build_qat_model( + model: torch.nn.Module, + training_samples: list[tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]], + *, + steps: int, + lr: float, + io_quantization: str, +) -> QuantizedModel: + exported_model = export(model, training_samples[0][0], strict=True).module( + check_guards=False + ) + quantizer = make_quantizer(is_qat=True, io_quantization=io_quantization) + prepared_model = prepare_qat_pt2e(exported_model, quantizer) + prepared_model = move_exported_model_to_train(prepared_model) + optimizer = torch.optim.SGD(prepared_model.parameters(), lr=lr, momentum=0.9) + + for step in range(steps): + total_loss = 0.0 + for inputs, target in training_samples: + prediction = prepared_model(*inputs) + loss = F.mse_loss(prediction, target) + optimizer.zero_grad() + loss.backward() + optimizer.step() + total_loss += float(loss.detach().item()) + print( + f"qat_step={step + 1}/{steps} " + f"loss={total_loss / len(training_samples):.6f}" + ) + + prepared_model = move_exported_model_to_eval(prepared_model) + converted_model = convert_pt2e(prepared_model) + return QuantizedModel( + model=converted_model, + coverage=quantization_coverage(converted_model), + ) + + +def export_vgf_artifact( + model: torch.nn.Module, + inputs: tuple[torch.Tensor, torch.Tensor], + *, + output_dir: Path, + name: str, + save_pte: bool, +) -> Path: + artifact_dir = output_dir / f"vgf_{name}" + artifact_dir.mkdir(parents=True, exist_ok=True) + compile_spec = make_vgf_compile_spec(artifact_dir) + partitioner = VgfPartitioner(compile_spec) + + print(f"[vgf] exporting {name} model to {artifact_dir} ...") + aten_dialect = export(model, args=inputs, strict=True) + edge_program_manager = to_edge_transform_and_lower( + aten_dialect, + partitioner=[partitioner], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + if save_pte: + pte_path = artifact_dir / f"rife_{name}_vgf.pte" + save_pte_program(edge_program_manager.to_executorch(), str(pte_path)) + print(f"[vgf] wrote {pte_path}") + return artifact_dir + + +def write_json(path: Path, payload: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") + + +def write_csv(path: Path, rows: list[dict[str, Any]]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fieldnames = sorted({key for row in rows for key in row}) + with path.open("w", encoding="utf-8", newline="") as file: + writer = csv.DictWriter(file, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(rows) + + +def write_markdown_report( + path: Path, + *, + aggregate: dict[str, Any], + quantization_reports: dict[str, dict[str, Any]], +) -> None: + lines = [ + "# RIFE PTQ/QAT Accuracy Report", + "", + "## Run", + "", + f"- Triples: {aggregate['triple_count']}", + f"- Input shape: `1x3x{aggregate['height']}x{aggregate['width']}`", + f"- Mode: `{aggregate['mode']}`", + f"- IO quantization: `{aggregate['io_quantization']}`", + f"- Data source: `{aggregate['data_source']}`", + f"- Random seed: `{aggregate['random_seed']}`", + "", + "## Aggregate", + "", + "| Metric | Mean | Min | Max |", + "| --- | ---: | ---: | ---: |", + ] + for key, label in ( + ("eager_psnr", "Eager PSNR"), + ("eager_ssim", "Eager SSIM"), + ("ptq_psnr", "PTQ PSNR"), + ("ptq_ssim", "PTQ SSIM"), + ("ptq_vs_eager_mae", "PTQ vs eager MAE"), + ("ptq_vs_eager_sqnr", "PTQ vs eager SQNR"), + ("qat_psnr", "QAT PSNR"), + ("qat_ssim", "QAT SSIM"), + ("qat_vs_eager_mae", "QAT vs eager MAE"), + ("qat_vs_eager_sqnr", "QAT vs eager SQNR"), + ): + if f"{key}_mean" not in aggregate: + continue + lines.append( + "| " + + " | ".join( + [ + label, + format_metric(aggregate.get(f"{key}_mean")), + format_metric(aggregate.get(f"{key}_min")), + format_metric(aggregate.get(f"{key}_max")), + ] + ) + + " |" + ) + + if quantization_reports: + lines.extend( + [ + "", + "## Quantization Coverage", + "", + "| Mode | Nodes | Quantized Ops | Quantized Context Ops | " + "FP32/Fallback Ops | Quantize Boundaries | " + "Dequantize Boundaries |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + ) + for name, report in sorted(quantization_reports.items()): + summary = report.get("summary", {}) + counts = summary.get("status_counts", {}) + lines.append( + "| " + + " | ".join( + [ + name, + format_metric(summary.get("node_count")), + format_metric(counts.get("quantized_op", 0)), + format_metric(counts.get("quantized_context_op", 0)), + format_metric(counts.get("fp32_or_fallback_op", 0)), + format_metric(counts.get("quantize_boundary", 0)), + format_metric(counts.get("dequantize_boundary", 0)), + ] + ) + + " |" + ) + + lines.extend( + [ + "", + "## Files", + "", + "- `metrics.json`: full aggregate and per-triple metrics.", + "- `metrics.csv`: tabular per-triple metrics.", + "- `quantization_coverage/`: PTQ/QAT graph coverage reports.", + "", + ] + ) + path.write_text("\n".join(lines), encoding="utf-8") + + +def write_quantization_coverage( + output_dir: Path, + name: str, + coverage: dict[str, Any], +) -> None: + coverage_dir = output_dir / "quantization_coverage" + coverage_dir.mkdir(parents=True, exist_ok=True) + write_json(coverage_dir / f"{name}.json", coverage) + if coverage.get("status") != "ok": + return + write_csv(coverage_dir / f"{name}_nodes.csv", coverage["nodes"]) + write_csv( + coverage_dir / f"{name}_op_status_counts.csv", coverage["op_status_counts"] + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + source = parser.add_mutually_exclusive_group(required=True) + source.add_argument( + "--triples-root", + type=Path, + help=( + "Directory containing frame sequences. Each directory with at " + "least three images is interpreted as frame0, frame2, target " + "frame1 triples." + ), + ) + source.add_argument( + "--triples-list", + type=Path, + help=( + "CSV or text file containing triples. Header form: " + "name,input0,input1,target. Text form: input0 input1 target." + ), + ) + parser.add_argument( + "--model-root", + type=Path, + required=True, + help=( + "Path to a standard Practical-RIFE repo containing train_log/ " + "and train_log/flownet.pkl." + ), + ) + parser.add_argument( + "--checkpoint", + type=Path, + default=None, + help=( + "Path to the Practical-RIFE flownet.pkl checkpoint. Defaults " + "to /train_log/flownet.pkl." + ), + ) + parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT) + parser.add_argument("--width", type=int, default=DEFAULT_WIDTH) + parser.add_argument( + "--preprocess", + choices=("resize", "center-crop"), + default="resize", + help="How to produce the fixed model input shape.", + ) + parser.add_argument( + "--sequence-stride", + type=int, + default=2, + help="Stride between discovered sliding-window triples.", + ) + parser.add_argument( + "--max-triples", + type=int, + default=30, + help="Maximum number of triples to evaluate.", + ) + parser.add_argument( + "--random-seed", + type=int, + default=2026, + help="Seed used when explicit random smoke-test samples are created.", + ) + parser.add_argument( + "--allow-random-data", + action="store_true", + help=( + "Allow deterministic random data when --triples-list is missing. " + "Only use this for smoke testing, not accuracy measurement." + ), + ) + parser.add_argument( + "--mode", + choices=("eager", "ptq", "both", "qat", "all"), + default="both", + help=( + "eager=FP32 only, ptq=PTQ only, both=eager+PTQ, " + "qat=eager+QAT, all=eager+PTQ+QAT." + ), + ) + parser.add_argument( + "--io-quantization", + choices=("uint8", "int8", "none"), + default="uint8", + help=( + "Model IO quantization for PTQ/QAT. uint8 preserves external " + "image IO; int8 uses symmetric IO; none leaves model IO FP32." + ), + ) + parser.add_argument( + "--calibration-samples", + type=int, + default=8, + help="Number of triples used to calibrate PTQ.", + ) + parser.add_argument( + "--qat-samples", + type=int, + default=8, + help="Number of triples used for QAT training.", + ) + parser.add_argument( + "--qat-steps", + type=int, + default=3, + help="Number of QAT epochs over the selected samples.", + ) + parser.add_argument( + "--qat-lr", + type=float, + default=1.0e-5, + help="Learning rate for the QAT optimizer.", + ) + parser.add_argument( + "--export-vgf", + choices=("none", "ptq", "qat", "all"), + default="none", + help=( + "Optionally lower the selected quantized model to VGF after " + "accuracy evaluation. This writes artifacts under --output-dir." + ), + ) + parser.add_argument( + "--save-pte", + action="store_true", + help="Also save an ExecuTorch .pte after VGF lowering.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=EXECUTORCH_ROOT / "out" / "rife_qat_example_768x384", + help="Directory for metrics and coverage reports.", + ) + return parser.parse_args() + + +def load_requested_triples(args: argparse.Namespace) -> list[FrameTriple]: + if args.triples_list is not None: + if args.triples_list.exists(): + triples = load_triples_list(args.triples_list) + elif args.allow_random_data: + print( + "WARNING: --triples-list was not found: " + f"{args.triples_list}. Creating and using random data. " + "Random data is only suitable for flow smoke testing, not " + "real accuracy assessment." + ) + triples = make_random_triples(args.max_triples) + else: + raise FileNotFoundError( + f"--triples-list was not found: {args.triples_list}. " + "Pass an existing triples file for accuracy measurement, " + "or add --allow-random-data for smoke testing." + ) + else: + triples = discover_triples(args.triples_root, args.sequence_stride) + return validate_triples(triples)[: args.max_triples] + + +def data_source(triples: list[FrameTriple]) -> str: + return "random" if all(triple.synthetic for triple in triples) else "frames" + + +def validate_args(args: argparse.Namespace) -> None: + validate_input_shape(args.height, args.width) + if args.max_triples <= 0: + raise ValueError("--max-triples must be positive.") + if args.calibration_samples <= 0: + raise ValueError("--calibration-samples must be positive.") + if args.qat_samples <= 0: + raise ValueError("--qat-samples must be positive.") + if args.qat_steps <= 0: + raise ValueError("--qat-steps must be positive.") + if args.qat_lr <= 0.0: + raise ValueError("--qat-lr must be positive.") + if args.export_vgf in ("ptq", "all") and args.mode not in ( + "ptq", + "both", + "all", + ): + raise ValueError("--export-vgf ptq requires --mode ptq, both, or all.") + if args.export_vgf in ("qat", "all") and args.mode not in ("qat", "all"): + raise ValueError("--export-vgf qat requires --mode qat or all.") + + +def build_requested_quantized_models( + args: argparse.Namespace, + model: torch.nn.Module, + triples: list[FrameTriple], +) -> tuple[torch.nn.Module | None, torch.nn.Module | None, dict[str, dict[str, Any]]]: + run_ptq = args.mode in ("ptq", "both", "all") + run_qat = args.mode in ("qat", "all") + ptq_model = None + quantization_reports: dict[str, dict[str, Any]] = {} + + if run_ptq: + calibration_inputs = [ + load_triple_tensors( + triple, + height=args.height, + width=args.width, + preprocess=args.preprocess, + random_seed=args.random_seed, + )[0] + for triple in triples[: args.calibration_samples] + ] + print(f"Building PTQ model with {len(calibration_inputs)} samples.") + ptq = build_ptq_model( + model, + calibration_inputs, + io_quantization=args.io_quantization, + ) + ptq_model = ptq.model + quantization_reports["ptq"] = ptq.coverage + + qat_model = None + if run_qat: + training_samples = [ + load_triple_tensors( + triple, + height=args.height, + width=args.width, + preprocess=args.preprocess, + random_seed=args.random_seed, + ) + for triple in triples[: args.qat_samples] + ] + print( + "Building QAT model with " + f"{len(training_samples)} samples and {args.qat_steps} steps." + ) + qat = build_qat_model( + model, + training_samples, + steps=args.qat_steps, + lr=args.qat_lr, + io_quantization=args.io_quantization, + ) + qat_model = qat.model + quantization_reports["qat"] = qat.coverage + + return ptq_model, qat_model, quantization_reports + + +def export_requested_vgf_artifacts( + args: argparse.Namespace, + triples: list[FrameTriple], + ptq_model: torch.nn.Module | None, + qat_model: torch.nn.Module | None, +) -> None: + if args.export_vgf == "none": + return + + inputs, _target = load_triple_tensors( + triples[0], + height=args.height, + width=args.width, + preprocess=args.preprocess, + random_seed=args.random_seed, + ) + 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( + ptq_model, + inputs, + output_dir=args.output_dir, + name="ptq", + 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( + qat_model, + inputs, + output_dir=args.output_dir, + name="qat", + save_pte=args.save_pte, + ) + + +def update_row_with_prediction_metrics( + row: dict[str, Any], + *, + prefix: str, + prediction: torch.Tensor, + target: torch.Tensor, +) -> None: + for key, value in quality_metrics(prediction, target).items(): + row[f"{prefix}_{key}"] = value + + +def update_row_with_delta_metrics( + row: dict[str, Any], + *, + prefix: str, + prediction: torch.Tensor, + reference: torch.Tensor, +) -> None: + for key, value in tensor_metrics(prediction, reference).items(): + row[f"{prefix}_{key}"] = value + + +def print_progress(index: int, total: int, row: dict[str, Any]) -> None: + progress = f"{index + 1}/{total} {row['name']}" + if "eager_psnr" in row: + progress += f" eager_psnr={row['eager_psnr']:.3f}" + if "ptq_psnr" in row: + progress += f" ptq_psnr={row['ptq_psnr']:.3f}" + if "qat_psnr" in row: + progress += f" qat_psnr={row['qat_psnr']:.3f}" + print(progress) + + +def evaluate_triples( + args: argparse.Namespace, + triples: list[FrameTriple], + model: torch.nn.Module, + ptq_model: torch.nn.Module | None, + qat_model: torch.nn.Module | None, +) -> list[dict[str, Any]]: + run_eager = args.mode in ("eager", "both", "qat", "all") + rows: list[dict[str, Any]] = [] + 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, + ) + row: dict[str, Any] = { + "index": index, + "name": triple.name, + "input0": str(triple.input0) if triple.input0 is not None else "", + "input1": str(triple.input1) if triple.input1 is not None else "", + "target": str(triple.target) if triple.target is not None else "", + "synthetic": triple.synthetic, + } + eager_output = None + if run_eager: + eager_output = run_model(model, inputs) + update_row_with_prediction_metrics( + row, + prefix="eager", + prediction=eager_output, + target=target, + ) + + if ptq_model is not None: + ptq_output = run_model(ptq_model, inputs) + update_row_with_prediction_metrics( + row, + prefix="ptq", + prediction=ptq_output, + target=target, + ) + if eager_output is None: + eager_output = run_model(model, inputs) + update_row_with_delta_metrics( + row, + prefix="ptq_vs_eager", + prediction=ptq_output, + reference=eager_output, + ) + + if qat_model is not None: + qat_output = run_model(qat_model, inputs) + update_row_with_prediction_metrics( + row, + prefix="qat", + prediction=qat_output, + target=target, + ) + if eager_output is None: + eager_output = run_model(model, inputs) + update_row_with_delta_metrics( + row, + prefix="qat_vs_eager", + prediction=qat_output, + reference=eager_output, + ) + + rows.append(row) + print_progress(index, len(triples), row) + + return rows + + +def build_aggregate( + args: argparse.Namespace, rows: list[dict[str, Any]] +) -> dict[str, Any]: + aggregate: dict[str, Any] = { + "triple_count": len(rows), + "height": args.height, + "width": args.width, + "preprocess": args.preprocess, + "mode": args.mode, + "io_quantization": args.io_quantization, + "data_source": "random" if rows and rows[0].get("synthetic") else "frames", + "random_seed": args.random_seed, + } + for prefix in ("eager", "ptq", "ptq_vs_eager", "qat", "qat_vs_eager"): + aggregate.update(aggregate_metrics(rows, prefix)) + return aggregate + + +def write_outputs( + args: argparse.Namespace, + aggregate: dict[str, Any], + rows: list[dict[str, Any]], + quantization_reports: dict[str, dict[str, Any]], +) -> None: + args.output_dir.mkdir(parents=True, exist_ok=True) + write_json( + args.output_dir / "metrics.json", + { + "aggregate": aggregate, + "rows": rows, + "quantization_coverage": { + name: report.get("summary", report) + for name, report in quantization_reports.items() + }, + }, + ) + write_csv(args.output_dir / "metrics.csv", rows) + for name, report in quantization_reports.items(): + write_quantization_coverage(args.output_dir, name, report) + write_markdown_report( + args.output_dir / "report.md", + aggregate=aggregate, + quantization_reports=quantization_reports, + ) + + print(f"Wrote {args.output_dir / 'metrics.json'}") + print(f"Wrote {args.output_dir / 'metrics.csv'}") + print(f"Wrote {args.output_dir / 'report.md'}") + + +def main() -> int: + args = parse_args() + validate_args(args) + + triples = load_requested_triples(args) + source = data_source(triples) + print( + f"Loaded {len(triples)} " + f"{'random samples' if source == 'random' else 'frame triples'}." + ) + print(f"Input shape: 1x3x{args.height}x{args.width}") + print(f"Mode: {args.mode}") + print(f"IO quantization: {args.io_quantization}") + print(f"Data source: {source}") + + model = load_rife_model(args.model_root, args.checkpoint) + ptq_model, qat_model, quantization_reports = build_requested_quantized_models( + args, + model, + triples, + ) + rows = evaluate_triples(args, triples, model, ptq_model, qat_model) + write_outputs(args, build_aggregate(args, rows), rows, quantization_reports) + export_requested_vgf_artifacts(args, triples, ptq_model, qat_model) + return 0 + + +if __name__ == "__main__": + sys.exit(main())