File size: 3,846 Bytes
7da2ecb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 | #!/usr/bin/env python3
"""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()
|