diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index f594ca6..525d60e 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -16,6 +16,7 @@ from dataclasses import dataclass, field from pathlib import Path import coloredlogs +import pandas as pd import torch import torch.nn as nn from torch.amp import GradScaler, autocast @@ -200,6 +201,36 @@ def _evaluate( return metrics +class CSVLogger: + """Log train/val metrics to CSV files using pandas. + + Creates: + {output_dir}/logs/train.csv — all train epochs + {output_dir}/logs/val.csv — all val epochs + {output_dir}/logs/epoch_{N}.csv — per-epoch details + """ + + def __init__(self, output_dir: Path) -> None: + self.log_dir = output_dir / "logs" + self.log_dir.mkdir(parents=True, exist_ok=True) + self.train_rows: list[dict] = [] + self.val_rows: list[dict] = [] + + def log_train(self, epoch: int, metrics: dict, lr: float, elapsed: float) -> None: + row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics} + self.train_rows.append(row) + # Append to cumulative CSV. + pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False) + # Per-epoch CSV. + pd.DataFrame([row]).to_csv(self.log_dir / f"epoch_{epoch:03d}_train.csv", index=False) + + def log_val(self, epoch: int, metrics: dict) -> None: + row = {"epoch": epoch, **metrics} + self.val_rows.append(row) + pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False) + pd.DataFrame([row]).to_csv(self.log_dir / f"epoch_{epoch:03d}_val.csv", index=False) + + def _clear_vram() -> None: """Free VRAM from previous runs before starting.""" import gc @@ -354,6 +385,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: LOGGER.info("🔄 Resuming from epoch %d", start_epoch) history: list[dict] = [] + csv_logger = CSVLogger(output_dir) LOGGER.info("🚀 Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch) @@ -442,10 +474,14 @@ def train(cfg: TrainConfigGTAUAV) -> None: "train": means, } + # Log train metrics to CSV. + csv_logger.log_train(epoch, means, optimizer.param_groups[0]["lr"], elapsed) + # Evaluation. if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1: val_metrics = _evaluate(model, test_loader, cfg.device) epoch_record["val"] = val_metrics + csv_logger.log_val(epoch, val_metrics) LOGGER.info( "🎯 val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f", epoch,