| |
| """Compare legacy and safetensors model predictions on one prepared sample.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import yaml |
|
|
|
|
| def sha256(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as stream: |
| while block := stream.read(16 * 1024 * 1024): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--dataset-root", required=True, type=Path) |
| parser.add_argument("--legacy-checkpoint", required=True, type=Path) |
| parser.add_argument("--release-checkpoint", required=True, type=Path) |
| parser.add_argument("--model-config", required=True, type=Path) |
| parser.add_argument("--timestamp", required=True) |
| parser.add_argument("--device", default="auto") |
| parser.add_argument("--output", type=Path, default=None) |
| args = parser.parse_args() |
|
|
| modeling_root = Path(__file__).resolve().parents[2] / "modeling" |
| if str(modeling_root) not in sys.path: |
| sys.path.insert(0, str(modeling_root)) |
| from src.data_pipeline.fast_dataset import FastDataset |
| from src.training_validation.common import build_model, load_model_checkpoint |
|
|
| requested = args.device |
| if requested == "auto": |
| requested = "cuda" if torch.cuda.is_available() else "cpu" |
| device = torch.device(requested) |
| with args.model_config.open("r", encoding="utf-8") as stream: |
| model_config = yaml.safe_load(stream) |
|
|
| dataset_config = { |
| "output_root": str(args.dataset_root.resolve()), |
| "catalog_path": str((args.dataset_root / "catalog.csv").resolve()), |
| "input_window": {"past_minutes": 50, "interval_minutes": 10}, |
| "splits": {"check": [{"start": args.timestamp, "end": args.timestamp}]}, |
| } |
| dataset = FastDataset( |
| dataset_config, |
| split="check", |
| required_inputs=["concat", "hsr"], |
| required_labels=[], |
| ) |
| if dataset.samples != [args.timestamp]: |
| raise ValueError(f"timestamp is not a complete inference sample: {args.timestamp}") |
| item = dataset[0] |
| x = torch.cat([item["inputs"][name] for name in ("concat", "hsr")], dim=1).unsqueeze(0).to(device) |
|
|
| model = build_model(model_config).to(device).eval() |
| with torch.inference_mode(): |
| load_model_checkpoint(model, args.legacy_checkpoint, device) |
| legacy = {name: value.detach().cpu() for name, value in model(x).items()} |
| load_model_checkpoint(model, args.release_checkpoint, device) |
| released = {name: value.detach().cpu() for name, value in model(x).items()} |
|
|
| errors = { |
| name: float(torch.max(torch.abs(legacy[name] - released[name])).item()) |
| for name in legacy |
| } |
| result = { |
| "timestamp": args.timestamp, |
| "device": str(device), |
| "legacy_checkpoint": {"filename": args.legacy_checkpoint.name, "sha256": sha256(args.legacy_checkpoint)}, |
| "release_checkpoint": {"filename": args.release_checkpoint.name, "sha256": sha256(args.release_checkpoint)}, |
| "max_absolute_error": errors, |
| "tolerance": 1e-5, |
| "passed": all(value <= 1e-5 for value in errors.values()), |
| } |
| if not result["passed"]: |
| raise AssertionError(f"prediction mismatch: {errors}") |
| output = args.output or args.release_checkpoint.parent / "prediction_equivalence.json" |
| output.parent.mkdir(parents=True, exist_ok=True) |
| temporary = output.with_name(f".{output.name}.partial") |
| with temporary.open("w", encoding="utf-8") as stream: |
| json.dump(result, stream, indent=2) |
| stream.write("\n") |
| os.replace(temporary, output) |
| print(json.dumps(result, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
