"""Phase 5B — improved accident-detector training and experiment tracking. Builds on the Phase 5 pipeline with the verified fixes: - 3-class model head (background + accident + non_accident) with explicit label-space mapping (`accident.labels`). - Per-component loss tracking. - Per-epoch checkpointing, validation evaluation, best-checkpoint selection on validation accident AP@50 (primary) with mAP@50 tie-break, and optional early stopping. - CSV experiment tracking under reports/accident/phase5b_experiments.csv. The Phase 5 baseline checkpoint (models/accident/best_detector.pth) is never touched. Phase 5B artifacts go to models/accident/phase5b_*.pth. """ from __future__ import annotations import csv import json import time from pathlib import Path from typing import Any import torch import torch.optim as optim from torch.utils.data import DataLoader from .dataset import AccidentParquetDataset, collate_detection from .metrics import evaluate_coco, threshold_analysis from .model import NUM_CLASSES, build_detector, count_parameters from .train import gather_predictions from .utils import DATASET_DIR, MODEL_DIR, REPORT_DIR, get_device, seed_everything EXPERIMENT_COLUMNS = [ "experiment_id", "model", "pretrained", "epochs", "batch_size", "learning_rate", "augmentation", "class_handling", "val_map50", "val_map50_95", "val_accident_ap50", "val_accident_precision", "val_accident_recall", "training_time", "notes", ] THRESHOLD_SWEEP = [0.05, 0.10, 0.20, 0.30, 0.40, 0.50, 0.60, 0.70, 0.80, 0.90] PHASE5B_CONFIG = { "seed": 42, "model": "fasterrcnn_mobilenet_v3_large_320_fpn", "num_classes": NUM_CLASSES, "image_size": [320, 320], "batch_size": 4, "epochs": 25, "learning_rate": 0.005, "momentum": 0.9, "weight_decay": 0.0005, "lr_scheduler_step": 8, "lr_scheduler_gamma": 0.1, "num_workers": 0, "pretrained": True, "device": "cpu", "early_stopping_patience": 6, "min_improvement": 0.002, "score_threshold": 0.05, } def checkpoint_score(metrics: dict[str, float]) -> tuple[float, float]: """Model-selection score: accident AP@50 primary, mAP@50 tie-break.""" return (round(metrics["accident_ap50"], 4), round(metrics["map50"], 4)) def is_improvement(score: float, best_score: float, min_improvement: float = 0.0) -> bool: """True when the new score beats the best by at least min_improvement.""" return score > best_score + min_improvement def append_experiment_row(csv_path: Path, row: dict[str, Any]) -> None: """Append one experiment row, creating the CSV with a header if needed.""" csv_path.parent.mkdir(parents=True, exist_ok=True) file_exists = csv_path.exists() with open(csv_path, "a", newline="", encoding="utf-8") as f: writer = csv.DictWriter(f, fieldnames=EXPERIMENT_COLUMNS) if not file_exists: writer.writeheader() writer.writerow({col: row.get(col, "") for col in EXPERIMENT_COLUMNS}) def build_loaders( config: dict[str, Any], dataset_dir: Path = DATASET_DIR, ) -> tuple[DataLoader, DataLoader]: """Create train and validation loaders for one experiment.""" target_size = tuple(config["image_size"]) common = { "batch_size": config["batch_size"], "num_workers": config.get("num_workers", 0), "collate_fn": collate_detection, } train_ds = AccidentParquetDataset(split="train", target_size=target_size, dataset_dir=dataset_dir) val_ds = AccidentParquetDataset(split="validation", target_size=target_size, dataset_dir=dataset_dir) train_loader = DataLoader(train_ds, shuffle=True, **common) val_loader = DataLoader(val_ds, shuffle=False, **common) return train_loader, val_loader def evaluate_metrics( model: torch.nn.Module, val_loader: DataLoader, device: torch.device, image_size: tuple[int, int], score_threshold: float = 0.05, ) -> dict[str, Any]: """Return flattened validation metrics used for selection and tracking.""" gt_records, pred_records = gather_predictions( model, val_loader, device, score_threshold=score_threshold ) metrics = evaluate_coco(gt_records, pred_records, image_size) accident = metrics["per_class"]["accident"] return { "map50": metrics["overall"]["mAP_50"], "map50_95": metrics["overall"]["mAP_50_95"], "accident_ap50": accident["AP_50"], "accident_ap50_95": accident["AP_50_95"], "accident_precision": accident["precision"], "accident_recall": accident["recall"], "non_accident_ap50": metrics["per_class"]["non_accident"]["AP_50"], "gt_records": gt_records, "pred_records": pred_records, } def run_experiment( config: dict[str, Any], experiment_id: str, notes: str, augmentation: str = "none", class_handling: str = "background_offset", dataset_dir: Path = DATASET_DIR, model_dir: Path = MODEL_DIR, report_dir: Path = REPORT_DIR, csv_path: Path | None = None, best_model_name: str = "phase5b_fasterrcnn_best.pth", ) -> dict[str, Any]: """Train one experiment with per-epoch checkpointing and validation eval. The best checkpoint (by validation accident AP@50, mAP@50 tie-break) is saved to `model_dir / best_model_name`. A rolling last checkpoint is kept so long CPU runs can resume. """ config = dict(config) seed_everything(config["seed"]) device = get_device() model_dir.mkdir(parents=True, exist_ok=True) report_dir.mkdir(parents=True, exist_ok=True) if csv_path is None: csv_path = report_dir / "phase5b_experiments.csv" train_loader, val_loader = build_loaders(config, dataset_dir=dataset_dir) model = build_detector( num_classes=config["num_classes"], pretrained=config.get("pretrained", True) ) model.to(device) # Optional resume from an interrupted-run checkpoint (see # `save_interrupted_checkpoint`). SGD momentum buffers cannot be # serialized out of a killed process; they re-warm within a few steps. start_epoch = 1 resume_path = config.get("resume_from") if resume_path: with open(resume_path, "rb") as f: resume_ckpt = torch.load(f, map_location="cpu", weights_only=False) model.load_state_dict(resume_ckpt["model_state_dict"]) start_epoch = int(resume_ckpt["epoch"]) + 1 print( f"[{experiment_id}] resuming from {resume_path} " f"(epoch {int(resume_ckpt['epoch'])}, best score {resume_ckpt.get('best_score_tuple')})" ) params = [p for p in model.parameters() if p.requires_grad] optimizer = optim.SGD( params, lr=config["learning_rate"], momentum=config["momentum"], weight_decay=config["weight_decay"], ) lr_scheduler = optim.lr_scheduler.StepLR( optimizer, step_size=config.get("lr_scheduler_step", 8), gamma=config.get("lr_scheduler_gamma", 0.1), last_epoch=start_epoch - 1, # fast-forward the schedule on resume ) image_size = tuple(config["image_size"]) score_threshold = config.get("score_threshold", 0.05) patience = config.get("early_stopping_patience", 0) min_improvement = config.get("min_improvement", 0.002) best_score_tuple = (-1.0, -1.0) best_epoch = -1 best_metrics: dict[str, Any] | None = None epochs_without_improvement = 0 history: list[dict[str, Any]] = [] if resume_path: best_score_tuple = tuple(resume_ckpt.get("best_score_tuple", (-1.0, -1.0))) best_epoch = int(resume_ckpt.get("best_epoch", -1)) history = list(resume_ckpt.get("history", [])) best_path = model_dir / best_model_name last_path = model_dir / best_model_name.replace(".pth", "_last.pth") history_path = report_dir / f"phase5b_{experiment_id.lower()}_history.json" print(f"[{experiment_id}] training {config['model']} on {device}") print(f"[{experiment_id}] trainable parameters: {count_parameters(model):,}") start_time = time.time() for epoch in range(start_epoch, config["epochs"] + 1): epoch_start = time.time() losses = train_one_epoch_components(model, optimizer, train_loader, device, epoch) lr_scheduler.step() metrics = evaluate_metrics(model, val_loader, device, image_size, score_threshold) score_tuple = checkpoint_score(metrics) history.append( { "epoch": epoch, "learning_rate": optimizer.param_groups[0]["lr"], **{k: round(v, 6) for k, v in losses.items()}, "val_map50": round(metrics["map50"], 4), "val_map50_95": round(metrics["map50_95"], 4), "val_accident_ap50": round(metrics["accident_ap50"], 4), "val_accident_precision": round(metrics["accident_precision"], 4), "val_accident_recall": round(metrics["accident_recall"], 4), "val_non_accident_ap50": round(metrics["non_accident_ap50"], 4), "epoch_seconds": round(time.time() - epoch_start, 1), } ) print( f"[{experiment_id}] epoch {epoch}/{config['epochs']} — " f"loss={losses['total']:.4f} — val accident AP@50={metrics['accident_ap50']:.4f} — " f"val mAP@50={metrics['map50']:.4f} — recall={metrics['accident_recall']:.4f}" ) # Rolling checkpoint so long CPU runs are not lost. torch.save(model.state_dict(), last_path) if score_tuple > best_score_tuple: best_score_tuple = score_tuple best_epoch = epoch best_metrics = {k: v for k, v in metrics.items() if not k.endswith("_records")} torch.save(model.state_dict(), best_path) epochs_without_improvement = 0 print(f"[{experiment_id}] new best checkpoint (epoch {epoch})") else: epochs_without_improvement += 1 # Flush history every epoch so long CPU runs are observable. with open(history_path, "w", encoding="utf-8") as f: json.dump( { "experiment_id": experiment_id, "config": config, "history": history, }, f, indent=2, ) if patience and epochs_without_improvement >= patience: print(f"[{experiment_id}] early stopping after {epoch} epochs (no improvement for {patience})") break total_time = time.time() - start_time # Reload the best checkpoint for final reporting. if best_metrics is None: raise RuntimeError(f"[{experiment_id}] no checkpoint was selected") model.load_state_dict(torch.load(best_path, map_location=device, weights_only=True)) final_metrics = evaluate_metrics(model, val_loader, device, image_size, score_threshold) thresholds = threshold_analysis( final_metrics["gt_records"], final_metrics["pred_records"], thresholds=THRESHOLD_SWEEP, image_size=image_size, ) summary = { "experiment_id": experiment_id, "notes": notes, "config": config, "augmentation": augmentation, "class_handling": class_handling, "best_epoch": best_epoch, "epochs_completed": history[-1]["epoch"], "trainable_parameters": count_parameters(model), "total_training_seconds": round(total_time, 1), "best_model_path": str(best_path), "last_model_path": str(last_path), "val_map50": final_metrics["map50"], "val_map50_95": final_metrics["map50_95"], "val_accident_ap50": final_metrics["accident_ap50"], "val_accident_precision": final_metrics["accident_precision"], "val_accident_recall": final_metrics["accident_recall"], "val_non_accident_ap50": final_metrics["non_accident_ap50"], "threshold_analysis": thresholds, "history": history, } # Persist per-experiment artifacts. with open(history_path, "w", encoding="utf-8") as f: json.dump(summary, f, indent=2) append_experiment_row( csv_path, { "experiment_id": experiment_id, "model": config["model"], "pretrained": config.get("pretrained", True), "epochs": summary["epochs_completed"], "batch_size": config["batch_size"], "learning_rate": config["learning_rate"], "augmentation": augmentation, "class_handling": class_handling, "val_map50": round(final_metrics["map50"], 4), "val_map50_95": round(final_metrics["map50_95"], 4), "val_accident_ap50": round(final_metrics["accident_ap50"], 4), "val_accident_precision": round(final_metrics["accident_precision"], 4), "val_accident_recall": round(final_metrics["accident_recall"], 4), "training_time": round(total_time, 1), "notes": notes, }, ) return summary def train_one_epoch_components( model: torch.nn.Module, optimizer: torch.optim.Optimizer, loader: DataLoader, device: torch.device, epoch: int, ) -> dict[str, float]: """One training epoch; dataset labels are offset to model space (+1).""" from .labels import dataset_to_model_labels model.train() totals = { "loss_classifier": 0.0, "loss_box_reg": 0.0, "loss_objectness": 0.0, "loss_rpn_box_reg": 0.0, "total": 0.0, } count = 0 for images, targets in loader: images = [img.to(device) for img in images] targets = [{k: v.to(device) for k, v in t.items()} for t in targets] targets = [{**t, "labels": dataset_to_model_labels(t["labels"])} for t in targets] optimizer.zero_grad() loss_dict = model(images, targets) loss = sum(loss_dict.values()) loss.backward() optimizer.step() totals["total"] += loss.item() for key in ("loss_classifier", "loss_box_reg", "loss_objectness", "loss_rpn_box_reg"): totals[key] += loss_dict[key].item() count += 1 return {k: v / count for k, v in totals.items()} if count else totals def save_interrupted_checkpoint( model_checkpoint_path: Path, history_path: Path, config: dict[str, Any], out_path: Path | None = None, ) -> Path: """Bundle the latest valid state of an interrupted run for resuming. Includes: model state_dict, completed epoch, best validation score so far, best checkpoint path, training configuration, seed, and the per-epoch history. SGD momentum buffers cannot be recovered from a stopped process (optimizer state was held in process memory); this is documented in the checkpoint and the momentum buffers re-warm within a few optimizer steps. StepLR state is reconstructed deterministically from epoch + schedule. """ if out_path is None: out_path = MODEL_DIR / "phase5b_interrupted_checkpoint.pth" state = torch.load(model_checkpoint_path, map_location="cpu", weights_only=True) epoch = 0 history: list[dict[str, Any]] = [] best_score_tuple = (-1.0, -1.0) best_epoch = -1 if history_path.exists(): with open(history_path, encoding="utf-8") as f: payload = json.load(f) entries = payload.get("history", payload if isinstance(payload, list) else []) history = entries if entries: epoch = int(entries[-1]["epoch"]) best_entry = max(entries, key=lambda h: (h["val_accident_ap50"], h["val_map50"])) best_score_tuple = (best_entry["val_accident_ap50"], best_entry["val_map50"]) best_epoch = int(best_entry["epoch"]) ckpt = { "model_state_dict": state, "epoch": epoch, "best_score_tuple": best_score_tuple, "best_epoch": best_epoch, "config": config, "history": history, "scheduler": { "type": "StepLR", "step_size": config.get("lr_scheduler_step", 8), "gamma": config.get("lr_scheduler_gamma", 0.1), "last_epoch": epoch, }, "optimizer_state": None, # SGD momentum not recoverable after process stop "interrupted": True, "notes": "Saved at an epoch boundary during intentional Phase 5B " "interruption. Resume via run_experiment config 'resume_from'.", } MODEL_DIR.mkdir(parents=True, exist_ok=True) torch.save(ckpt, out_path) metadata = { "checkpoint_path": str(out_path), "model_checkpoint_source": str(model_checkpoint_path), "completed_epoch": epoch, "best_score_tuple": list(best_score_tuple), "best_epoch": best_epoch, "history_path": str(history_path), "config": config, "optimizer_state_included": False, "optimizer_note": "SGD momentum buffers lost on process stop; " "scheduler reconstructed from epoch. Resume re-warms momentum.", "baseline_model_untouched": "models/accident/best_detector.pth", } with open(MODEL_DIR / "phase5b_checkpoint_metadata.json", "w", encoding="utf-8") as f: json.dump(metadata, f, indent=2) return out_path def load_phase5b_detector( checkpoint_path: Path | None = None, image_size: tuple[int, int] = (320, 320), ): """Load the Phase 5B detector through the inference contract wrapper.""" from .inference import AccidentDetector if checkpoint_path is None: checkpoint_path = MODEL_DIR / "phase5b_fasterrcnn_best.pth" return AccidentDetector(checkpoint_path, num_classes=NUM_CLASSES, image_size=image_size)