| """ |
| Separate (one-model-per-class) MVTec training. |
| |
| Saves metrics only — no checkpoints. After each class finishes, the model is deleted. |
| Compatible with the unified path: does not change train_unified behavior. |
| """ |
|
|
| import gc |
| import json |
| import os |
| import sys |
| import time |
| from pathlib import Path |
| from typing import Any, Dict, List, Optional |
|
|
| import numpy as np |
| import torch |
| from torch.amp import GradScaler |
|
|
| sys.path.append(os.path.join(os.path.dirname(__file__), "..")) |
|
|
| from data.data_loader import build_dataloader |
| from scripts.evaluate import evaluate_category |
| from scripts.train_unified import build_model_and_optimizer, build_scheduler, run_one_epoch |
| from utils.checkpointing import append_epoch_metrics, ensure_run_dir, save_metrics_json |
| from utils.metrics import compute_image_auroc, compute_pixel_ap, compute_pixel_auroc |
|
|
|
|
| def _eval_one_category(model, cfg: dict, category: str, device: str) -> Dict[str, float]: |
| root = cfg["data"]["mvtec_root"] |
| size = tuple(cfg["data"]["mvtec_size"]) |
| dataset_name = cfg["train"].get("dataset_name", "mvtec") |
| smooth_k = cfg["model"].get("anomaly_smooth_kernel", 7) |
| smooth_s = cfg["model"].get("anomaly_smooth_sigma", 1.5) |
|
|
| res = evaluate_category( |
| model, root, category, size, device, dataset_name, |
| smooth_kernel=smooth_k, smooth_sigma=smooth_s, |
| ) |
| if len(np.unique(res["image_labels"])) < 2: |
| return {"I-AUROC": 0.0, "P-AUROC": 0.0, "P-AP": 0.0} |
|
|
| return { |
| "I-AUROC": compute_image_auroc(res["image_scores"], res["image_labels"]), |
| "P-AUROC": compute_pixel_auroc(res["pixel_scores"], res["pixel_labels"]), |
| "P-AP": compute_pixel_ap(res["pixel_scores"], res["pixel_labels"]), |
| } |
|
|
|
|
| def _free_model(*objs) -> None: |
| for o in objs: |
| del o |
| gc.collect() |
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
|
|
| def train_one_category( |
| category: str, |
| cfg: dict, |
| device: str, |
| run_dir: str, |
| ) -> Dict[str, Any]: |
| """ |
| Train a fresh model on one category. Metrics only (no checkpoints). |
| Deletes the model before returning. |
| """ |
| cat_dir = str(ensure_run_dir(os.path.join(run_dir, category))) |
| max_epochs = int(cfg["train"].get("max_epochs", cfg["train"].get("epochs", 400))) |
| patience = int(cfg["train"].get("early_stopping_patience", 10)) |
| monitor = cfg["train"].get("early_stopping_monitor", "I-AUROC") |
| if monitor.startswith("avg_"): |
| monitor = monitor[len("avg_"):] |
| use_amp = bool(cfg["train"].get("use_amp", True)) |
|
|
| train_loader, train_dataset = build_dataloader( |
| cfg, dataset_name="mvtec", shuffle=True, categories=[category] |
| ) |
| if len(train_dataset) == 0: |
| print(f"[WARN] {category}: no training images; skipping.") |
| return {"category": category, "skipped": True} |
|
|
| model, optimizer, criterion, device = build_model_and_optimizer(cfg, device=device) |
| scheduler = build_scheduler(optimizer, cfg) |
|
|
| amp_on = use_amp and device != "cpu" and torch.cuda.is_available() |
| scaler = GradScaler("cuda", enabled=amp_on) if amp_on else GradScaler(enabled=False) |
|
|
| history: Dict[str, Any] = {"train": [], "eval": []} |
| best_metric = float("-inf") |
| best_eval: Dict[str, float] = {} |
| epochs_without_improvement = 0 |
| last_epoch = 0 |
|
|
| print( |
| f"\n=== Separate train: {category} | samples={len(train_dataset)} | " |
| f"batch={train_loader.batch_size} | max_epochs={max_epochs} | patience={patience} | " |
| f"wd={cfg['train'].get('weight_decay', 1e-4)} | lr={optimizer.param_groups[0]['lr']} ===", |
| flush=True, |
| ) |
|
|
| for epoch in range(1, max_epochs + 1): |
| last_epoch = epoch |
| t0 = time.time() |
| train_losses = run_one_epoch( |
| model, train_loader, optimizer, criterion, device, |
| train=True, use_amp=use_amp, scaler=scaler, |
| desc=f"{category} {epoch}/{max_epochs}", |
| ) |
| eval_m = _eval_one_category(model, cfg, category, device) |
| elapsed = time.time() - t0 |
| monitor_val = float(eval_m.get(monitor, eval_m.get("I-AUROC", float("-inf")))) |
| if scheduler is not None: |
| scheduler.step(monitor_val) |
| current_lr = optimizer.param_groups[0]["lr"] |
|
|
| history["train"].append({"epoch": epoch, **train_losses}) |
| history["eval"].append({"epoch": epoch, **eval_m}) |
|
|
| epoch_record = { |
| "epoch": epoch, |
| "category": category, |
| "train": train_losses, |
| "eval": eval_m, |
| "monitor": monitor, |
| "monitor_value": monitor_val, |
| "lr": current_lr, |
| "elapsed_sec": round(elapsed, 2), |
| } |
| append_epoch_metrics(cat_dir, epoch_record) |
| save_metrics_json(cat_dir, history) |
|
|
| print( |
| f"[{category} Epoch {epoch}/{max_epochs}] " |
| f"train_total={train_losses['total']:.4f} | " |
| f"I-AUROC={eval_m.get('I-AUROC', 0):.4f} | " |
| f"P-AUROC={eval_m.get('P-AUROC', 0):.4f} | " |
| f"P-AP={eval_m.get('P-AP', 0):.4f} | " |
| f"{monitor}={monitor_val:.4f} | lr={current_lr:.2e} | {elapsed:.1f}s", |
| flush=True, |
| ) |
|
|
| if monitor_val > best_metric: |
| best_metric = monitor_val |
| best_eval = dict(eval_m) |
| epochs_without_improvement = 0 |
| print(f" -> new best {monitor}={best_metric:.4f}", flush=True) |
| else: |
| epochs_without_improvement += 1 |
| print(f" -> no improvement ({epochs_without_improvement}/{patience})", flush=True) |
|
|
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
| if epochs_without_improvement >= patience: |
| print(f"Early stopping {category} at epoch {epoch}.", flush=True) |
| break |
|
|
| result = { |
| "category": category, |
| "skipped": False, |
| "last_epoch": last_epoch, |
| "best_metric": best_metric, |
| "best_eval": best_eval, |
| "history": history, |
| } |
| with open(os.path.join(cat_dir, "metrics", "best_summary.json"), "w", encoding="utf-8") as f: |
| json.dump( |
| {k: v for k, v in result.items() if k != "history"}, |
| f, indent=2, ensure_ascii=False, |
| ) |
|
|
| _free_model(model, optimizer, criterion, scaler, scheduler, train_loader, train_dataset) |
| return result |
|
|
|
|
| def train_all_separate( |
| cfg: dict, |
| run_dir: str, |
| device: str = None, |
| categories: Optional[List[str]] = None, |
| skip_done: bool = True, |
| hf_repo_id: Optional[str] = None, |
| hf_token: Optional[str] = None, |
| hf_run_name: Optional[str] = None, |
| hf_push_every_category: bool = True, |
| ) -> Dict[str, Any]: |
| """ |
| Train one model per MVTec category. Metrics-only; models are deleted after each class. |
| """ |
| device = device or cfg["train"].get("device", "cuda") |
| device = device if torch.cuda.is_available() else "cpu" |
| run_dir = str(ensure_run_dir(run_dir)) |
| cats = categories or list(cfg["data"]["mvtec_categories"]) |
|
|
| summary_path = os.path.join(run_dir, "summary.json") |
| summary: Dict[str, Any] = {"categories": {}, "average": {}} |
| if os.path.isfile(summary_path): |
| try: |
| with open(summary_path, "r", encoding="utf-8") as f: |
| summary = json.load(f) |
| except Exception: |
| pass |
|
|
| print( |
| f"Separate run dir: {run_dir} | classes={len(cats)} | " |
| f"device={device} | amp={cfg['train'].get('use_amp', True)}", |
| flush=True, |
| ) |
|
|
| for cat in cats: |
| done_file = os.path.join(run_dir, cat, "metrics", "best_summary.json") |
| if skip_done and os.path.isfile(done_file): |
| with open(done_file, "r", encoding="utf-8") as f: |
| prev = json.load(f) |
| summary["categories"][cat] = prev.get("best_eval", prev) |
| print(f"=== Skip {cat} (already has metrics) ===", flush=True) |
| continue |
|
|
| result = train_one_category(cat, cfg, device, run_dir) |
| if not result.get("skipped"): |
| summary["categories"][cat] = result["best_eval"] |
|
|
| |
| finished = [v for v in summary["categories"].values() if isinstance(v, dict) and "I-AUROC" in v] |
| if finished: |
| summary["average"] = { |
| k: float(np.mean([c[k] for c in finished])) |
| for k in ["I-AUROC", "P-AUROC", "P-AP"] |
| } |
| with open(summary_path, "w", encoding="utf-8") as f: |
| json.dump(summary, f, indent=2, ensure_ascii=False) |
|
|
| print( |
| f"Running average over {len(finished)} classes: " |
| f"I-AUROC={summary['average'].get('I-AUROC', 0):.4f} " |
| f"P-AUROC={summary['average'].get('P-AUROC', 0):.4f} " |
| f"P-AP={summary['average'].get('P-AP', 0):.4f}", |
| flush=True, |
| ) |
|
|
| if hf_repo_id and hf_push_every_category: |
| try: |
| from utils.hf_sync import push_folder_to_hub |
| url = push_folder_to_hub( |
| run_dir, hf_repo_id, |
| path_in_repo=f"runs/{hf_run_name or Path(run_dir).name}", |
| token=hf_token, |
| commit_message=f"Separate metrics after {cat}", |
| ) |
| print(f" -> HF metrics push OK: {url}", flush=True) |
| except Exception as e: |
| print(f" -> HF push failed: {e}", flush=True) |
|
|
| print("\n=== Separate training finished ===") |
| for cat, m in sorted(summary.get("categories", {}).items()): |
| if isinstance(m, dict) and "I-AUROC" in m: |
| print( |
| f" {cat:15s} I-AUROC={m['I-AUROC']:.3f} " |
| f"P-AUROC={m['P-AUROC']:.3f} P-AP={m['P-AP']:.3f}" |
| ) |
| if summary.get("average"): |
| a = summary["average"] |
| print( |
| f" {'AVERAGE':15s} I-AUROC={a.get('I-AUROC', 0):.3f} " |
| f"P-AUROC={a.get('P-AUROC', 0):.3f} P-AP={a.get('P-AP', 0):.3f}" |
| ) |
|
|
| return summary |
|
|
