| """ |
| Unified MVTec training loop with tqdm, AMP, early stopping, and resumable checkpoints. |
| """ |
|
|
| import os |
| import sys |
| import time |
| from typing import Any, Dict, Optional, Tuple |
|
|
| import torch |
| from torch.amp import GradScaler, autocast |
| from torch.optim import Adam |
| from torch.optim.lr_scheduler import ReduceLROnPlateau |
| from tqdm.auto import tqdm |
|
|
| sys.path.append(os.path.join(os.path.dirname(__file__), "..")) |
|
|
| from models.model import OmniAL |
| from scripts.evaluate import evaluate_dataset |
| from utils.checkpointing import ( |
| append_epoch_metrics, |
| ensure_run_dir, |
| find_latest_checkpoint, |
| load_checkpoint, |
| restore_scheduler, |
| save_checkpoint, |
| save_metrics_json, |
| ) |
| from utils.metrics import OmniALLoss |
|
|
|
|
| def build_model_and_optimizer(cfg: dict, device: str = None): |
| device = device or cfg["train"]["device"] |
| device = device if torch.cuda.is_available() else "cpu" |
|
|
| model = OmniAL(cfg).to(device) |
| tcfg = cfg["train"] |
| optimizer = Adam( |
| model.parameters(), |
| lr=float(tcfg.get("lr", 1e-4)), |
| weight_decay=float(tcfg.get("weight_decay", 1e-4)), |
| ) |
| criterion = OmniALLoss(cfg) |
| return model, optimizer, criterion, device |
|
|
|
|
| def build_scheduler(optimizer, cfg: dict): |
| """ReduceLROnPlateau on the eval monitor (mode=max). Used by unified and separate.""" |
| tcfg = cfg.get("train", {}) |
| if str(tcfg.get("scheduler", "plateau")).lower() in ("none", "off", "false"): |
| return None |
| return ReduceLROnPlateau( |
| optimizer, |
| mode=tcfg.get("scheduler_mode", "max"), |
| factor=float(tcfg.get("scheduler_factor", 0.5)), |
| patience=int(tcfg.get("scheduler_patience", 3)), |
| min_lr=float(tcfg.get("scheduler_min_lr", 1e-7)), |
| ) |
|
|
|
|
| def run_one_epoch( |
| model, |
| dataloader, |
| optimizer, |
| criterion, |
| device: str, |
| train: bool = True, |
| use_amp: bool = False, |
| scaler: Optional[GradScaler] = None, |
| desc: str = "train", |
| ) -> Dict[str, float]: |
| model.train() if train else model.eval() |
| epoch_losses = { |
| "total": 0.0, |
| "l2_img": 0.0, |
| "ssim_img": 0.0, |
| "l2_jnd": 0.0, |
| "ssim_jnd": 0.0, |
| "focal": 0.0, |
| } |
| n_batches = 0 |
| use_amp = bool(use_amp and device != "cpu" and torch.cuda.is_available()) |
| amp_device = "cuda" if (device != "cpu" and torch.cuda.is_available()) else "cpu" |
| if scaler is None: |
| scaler = GradScaler("cuda", enabled=use_amp) if amp_device == "cuda" else GradScaler(enabled=False) |
|
|
| context = torch.enable_grad() if train else torch.no_grad() |
| pbar = tqdm(dataloader, desc=desc, leave=False, dynamic_ncols=True) |
|
|
| with context: |
| for batch in pbar: |
| anomaly_img = batch["anomaly_img"].to(device, non_blocking=True) |
| targets = { |
| "image": batch["image"].to(device, non_blocking=True), |
| "jnd": batch["jnd"].to(device, non_blocking=True), |
| "severity": batch["severity"].to(device, non_blocking=True), |
| } |
|
|
| with autocast(amp_device, enabled=use_amp): |
| outputs = model(anomaly_img) |
| losses = criterion(outputs, targets) |
|
|
| if train: |
| loss = losses["total"] |
| if not torch.isfinite(loss): |
| print(f"[WARN] non-finite loss={loss.item()}; skipping step", flush=True) |
| optimizer.zero_grad(set_to_none=True) |
| continue |
| optimizer.zero_grad(set_to_none=True) |
| scaler.scale(loss).backward() |
| |
| scaler.unscale_(optimizer) |
| torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) |
| scaler.step(optimizer) |
| scaler.update() |
|
|
| for k in epoch_losses: |
| v = losses[k] |
| epoch_losses[k] += v.item() if torch.is_tensor(v) else v |
| n_batches += 1 |
| pbar.set_postfix(loss=f"{losses['total'].item():.4f}") |
|
|
| del anomaly_img, targets, outputs, losses |
|
|
| for k in epoch_losses: |
| epoch_losses[k] /= max(n_batches, 1) |
| return epoch_losses |
|
|
|
|
| def _monitor_value(eval_results: Dict[str, Any], monitor: str) -> float: |
| if not eval_results or "average" not in eval_results: |
| return float("-inf") |
| avg = eval_results["average"] |
| key_map = { |
| "avg_I-AUROC": "I-AUROC", |
| "avg_P-AUROC": "P-AUROC", |
| "avg_P-AP": "P-AP", |
| "I-AUROC": "I-AUROC", |
| "P-AUROC": "P-AUROC", |
| "P-AP": "P-AP", |
| } |
| k = key_map.get(monitor, monitor) |
| return float(avg.get(k, float("-inf"))) |
|
|
|
|
| def _write_latest_per_category(run_dir: str, epoch: int, per_cat: Dict[str, Any], average: Dict[str, float]) -> None: |
| """Overwrite a human-readable table for the latest epoch only.""" |
| path = os.path.join(run_dir, "metrics", "latest_per_category.txt") |
| lines = [f"Epoch {epoch} per-category results", "-" * 56] |
| for cat in sorted(per_cat.keys()): |
| m = per_cat[cat] |
| lines.append( |
| f"{cat:15s} | I-AUROC={m['I-AUROC']:.3f} P-AUROC={m['P-AUROC']:.3f} P-AP={m['P-AP']:.3f}" |
| ) |
| if average: |
| lines.append("-" * 56) |
| lines.append( |
| f"{'AVERAGE':15s} | I-AUROC={average.get('I-AUROC', 0):.3f} " |
| f"P-AUROC={average.get('P-AUROC', 0):.3f} P-AP={average.get('P-AP', 0):.3f}" |
| ) |
| with open(path, "w", encoding="utf-8") as f: |
| f.write("\n".join(lines) + "\n") |
|
|
|
|
| def train_unified( |
| model, |
| train_loader, |
| optimizer, |
| criterion, |
| device: str, |
| cfg: dict, |
| run_dir: str, |
| start_epoch: int = 1, |
| history: Optional[Dict[str, Any]] = None, |
| scaler: Optional[GradScaler] = None, |
| scheduler=None, |
| best_metric: float = float("-inf"), |
| hf_repo_id: Optional[str] = None, |
| hf_token: Optional[str] = None, |
| hf_run_name: Optional[str] = None, |
| ) -> Tuple[Dict[str, Any], int, float]: |
| """ |
| Train unified OmniAL on MVTec with per-epoch eval, checkpointing, and early stopping. |
| Returns (history, last_epoch, best_metric). |
| """ |
| run_dir = str(ensure_run_dir(run_dir)) |
| ckpt_dir = os.path.join(run_dir, "checkpoints") |
| 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", "avg_I-AUROC") |
| save_every = int(cfg["logging"].get("save_every", 2)) |
| hf_push_every = int(cfg["logging"].get("hf_push_every", 0)) |
| use_amp = bool(cfg["train"].get("use_amp", True)) |
| dataset_name = cfg["train"].get("dataset_name", "mvtec") |
|
|
| if scaler is None: |
| 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) |
| if scheduler is None: |
| scheduler = build_scheduler(optimizer, cfg) |
|
|
| if history is None: |
| history = { |
| "train": [], |
| "eval": [], |
| "per_category": [], |
| } |
|
|
| epochs_without_improvement = 0 |
| last_per_category = {} |
|
|
| print( |
| f"Run dir: {run_dir} | epochs {start_epoch}-{max_epochs} | " |
| f"batch={train_loader.batch_size} | amp={use_amp} | patience={patience} | " |
| f"monitor={monitor} | wd={cfg['train'].get('weight_decay', 1e-4)} | " |
| f"lr={optimizer.param_groups[0]['lr']}" |
| ) |
| print("Per-category metrics are written to metrics/ (not printed each epoch).") |
|
|
| for epoch in range(start_epoch, max_epochs + 1): |
| t0 = time.time() |
| train_losses = run_one_epoch( |
| model, train_loader, optimizer, criterion, device, |
| train=True, use_amp=use_amp, scaler=scaler, |
| desc=f"Epoch {epoch}/{max_epochs} train", |
| ) |
| history["train"].append({"epoch": epoch, **train_losses}) |
|
|
| |
| eval_results = evaluate_dataset( |
| model, cfg, dataset_name=dataset_name, device=device, verbose=False |
| ) |
| monitor_val = _monitor_value(eval_results, monitor) |
| elapsed = time.time() - t0 |
| if scheduler is not None: |
| scheduler.step(monitor_val) |
| current_lr = optimizer.param_groups[0]["lr"] |
|
|
| per_cat = {k: v for k, v in eval_results.items() if k != "average"} |
| last_per_category = per_cat |
| epoch_record = { |
| "epoch": epoch, |
| "train": train_losses, |
| "eval_average": eval_results.get("average", {}), |
| "eval_per_category": per_cat, |
| "monitor": monitor, |
| "monitor_value": monitor_val, |
| "lr": current_lr, |
| "elapsed_sec": round(elapsed, 2), |
| } |
| history["eval"].append(epoch_record["eval_average"]) |
| history["per_category"].append({ |
| "epoch": epoch, |
| "categories": per_cat, |
| }) |
| append_epoch_metrics(run_dir, epoch_record) |
| save_metrics_json(run_dir, history) |
| |
| _write_latest_per_category(run_dir, epoch, per_cat, epoch_record["eval_average"]) |
|
|
| |
| print( |
| f"[Epoch {epoch}/{max_epochs}] " |
| f"train_total={train_losses['total']:.4f} | " |
| f"I-AUROC={epoch_record['eval_average'].get('I-AUROC', 0):.4f} | " |
| f"P-AUROC={epoch_record['eval_average'].get('P-AUROC', 0):.4f} | " |
| f"P-AP={epoch_record['eval_average'].get('P-AP', 0):.4f} | " |
| f"{monitor}={monitor_val:.4f} | lr={current_lr:.2e} | {elapsed:.1f}s", |
| flush=True, |
| ) |
|
|
| improved = monitor_val > best_metric |
| if improved: |
| best_metric = monitor_val |
| epochs_without_improvement = 0 |
| save_checkpoint( |
| os.path.join(ckpt_dir, "best.pt"), |
| epoch, model, optimizer, scaler, history, best_metric, cfg, tag="best", |
| scheduler=scheduler, |
| ) |
| print(f" -> new best {monitor}={best_metric:.4f} (saved best.pt)", flush=True) |
| else: |
| epochs_without_improvement += 1 |
| print(f" -> no improvement ({epochs_without_improvement}/{patience})", flush=True) |
| save_checkpoint( |
| os.path.join(ckpt_dir, "last.pt"), |
| epoch, model, optimizer, scaler, history, best_metric, cfg, tag="last", |
| scheduler=scheduler, |
| ) |
|
|
| if epoch % save_every == 0: |
| ep_path = os.path.join(ckpt_dir, f"epoch_{epoch:04d}.pt") |
| save_checkpoint( |
| ep_path, epoch, model, optimizer, scaler, history, best_metric, cfg, |
| tag=f"epoch_{epoch}", scheduler=scheduler, |
| ) |
| print(f" -> checkpoint saved: {ep_path}") |
|
|
| if ( |
| hf_repo_id |
| and hf_push_every > 0 |
| and epoch % hf_push_every == 0 |
| ): |
| try: |
| from utils.hf_sync import push_folder_to_hub |
| path_in_repo = f"runs/{hf_run_name or os.path.basename(run_dir)}" |
| url = push_folder_to_hub( |
| run_dir, hf_repo_id, path_in_repo=path_in_repo, |
| token=hf_token, |
| commit_message=f"Checkpoint epoch {epoch}", |
| ) |
| print(f" -> HF push OK: {url}") |
| except Exception as e: |
| print(f" -> HF push failed: {e}") |
|
|
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
|
|
| if epochs_without_improvement >= patience: |
| print(f"Early stopping at epoch {epoch} (no {monitor} improvement for {patience} epochs).") |
| break |
|
|
| |
| if last_per_category: |
| print("\nFinal per-category results:") |
| for cat, m in sorted(last_per_category.items()): |
| print(f" {cat:15s} I-AUROC={m['I-AUROC']:.3f} P-AUROC={m['P-AUROC']:.3f} P-AP={m['P-AP']:.3f}") |
|
|
| return history, epoch, best_metric |
|
|
|
|
| def resume_from_run( |
| model, |
| optimizer, |
| scaler, |
| device: str, |
| run_dir: str, |
| checkpoint_path: Optional[str] = None, |
| scheduler=None, |
| ) -> Tuple[int, Dict[str, Any], float]: |
| """Load last/best checkpoint if present. Returns (start_epoch, history, best_metric).""" |
| path = checkpoint_path or find_latest_checkpoint(run_dir) |
| if not path or not os.path.isfile(path): |
| return 1, None, float("-inf") |
|
|
| ckpt = load_checkpoint(path, model, optimizer, scaler, device=device) |
| restore_scheduler(scheduler, ckpt) |
| start_epoch = int(ckpt["epoch"]) + 1 |
| history = ckpt.get("history") |
| best_metric = float(ckpt.get("best_metric", float("-inf"))) |
| lr = optimizer.param_groups[0]["lr"] |
| print(f"Resumed from {path} -> start_epoch={start_epoch}, best_metric={best_metric:.4f}, lr={lr:.2e}") |
| return start_epoch, history, best_metric |
|
|
