Add val/train loss to evaluation and plots

- _evaluate() now computes per-batch loss when loss_fn is provided
- Val loss and train recall loss saved in val.csv and train_recall.csv
- Overview plot shows train vs val loss curves side by side
- Helps detect overfitting: val loss diverging from train loss

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/training/plot_metrics.py

@@ -114,12 +114,15 @@ def plot_combined(train_df: pd.DataFrame, val_df: pd.DataFrame, out_dir: Path) -
     fig, axes = plt.subplots(1, 3, figsize=(18, 5))
     fig.suptitle("Training Overview", fontsize=16, fontweight="bold")
 
-    # 1. Train loss.
+    # 1. Train + Val loss.
     ax = axes[0]
-    sns.lineplot(data=train_df, x="epoch", y="total", ax=ax, marker="o", linewidth=2, color="crimson")
-    ax.set_title("Train Loss")
+    sns.lineplot(data=train_df, x="epoch", y="total", ax=ax, marker="o", linewidth=2, color="crimson", label="train")
+    if "loss" in val_df.columns:
+        sns.lineplot(data=val_df, x="epoch", y="loss", ax=ax, marker="s", linewidth=2, color="royalblue", label="val")
+    ax.set_title("Loss (train vs val)")
     ax.set_xlabel("Epoch")
     ax.set_ylabel("InfoNCE Loss")
+    ax.legend()
 
     # 2. Val R@1 q→g.
     ax = axes[1]

src/training/train_gtauav.py

@@ -175,14 +175,18 @@ def _evaluate(
     model: AsymmetricEncoder,
     loader: DataLoader,
     device: str,
+    loss_fn: nn.Module | None = None,
+    epoch: int = 0,
+    total_epochs: int = 1,
     k_values: tuple[int, ...] = (1, 5, 10),
     max_batches: int | None = None,
     desc: str = "eval",
 ) -> dict[str, float]:
-    """Compute R@K. Use max_batches to limit for train set (faster)."""
+    """Compute R@K and optional loss. Use max_batches to limit for train set."""
     model.eval()
     all_query: list[torch.Tensor] = []
     all_gallery: list[torch.Tensor] = []
+    batch_losses: list[float] = []
 
     for i, batch in enumerate(tqdm(loader, desc=f"  {desc}", unit="batch", leave=False)):
         if max_batches is not None and i >= max_batches:
@@ -206,6 +210,11 @@ def _evaluate(
         all_query.append(embeddings["query"].cpu())
         all_gallery.append(embeddings["gallery"].cpu())
 
+        # Per-batch loss (if loss_fn provided).
+        if loss_fn is not None:
+            loss_dict = loss_fn(embeddings, epoch=epoch, total_epochs=total_epochs)
+            batch_losses.append(float(loss_dict["total"].item()))
+
     query = torch.cat(all_query, dim=0)
     gallery = torch.cat(all_gallery, dim=0)
 
@@ -214,6 +223,11 @@ def _evaluate(
     targets = torch.arange(n)
 
     metrics: dict[str, float] = {}
+
+    # Average loss across batches.
+    if batch_losses:
+        metrics["loss"] = sum(batch_losses) / len(batch_losses)
+
     sorted_idx = sim.argsort(dim=1, descending=True)
     for k in k_values:
         top_k = sorted_idx[:, :k]
@@ -675,6 +689,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             train_eval_batches = len(test_loader)
             train_recall = _evaluate(
                 model, train_eval_loader, cfg.device,
+                loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs,
                 max_batches=train_eval_batches, desc="eval-train",
             )
             epoch_record["train_recall"] = train_recall
@@ -689,7 +704,11 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             )
 
             # Val R@K (full test set).
-            val_metrics = _evaluate(model, test_loader, cfg.device, desc="eval-val")
+            val_metrics = _evaluate(
+                model, test_loader, cfg.device,
+                loss_fn=loss_fn, epoch=epoch, total_epochs=cfg.epochs,
+                desc="eval-val",
+            )
             epoch_record["val"] = val_metrics
             csv_logger.log_val(epoch, val_metrics)
             generate_plots(csv_logger.log_dir)
@@ -753,7 +772,10 @@ def train(cfg: TrainConfigGTAUAV) -> None:
 
     # Save final eval report.
     LOGGER.info("Running final evaluation...")
-    final_metrics = _evaluate(model, test_loader, cfg.device)
+    final_metrics = _evaluate(
+        model, test_loader, cfg.device,
+        loss_fn=loss_fn, epoch=cfg.epochs - 1, total_epochs=cfg.epochs,
+    )
     report = {
         "config": vars(cfg),
         "metrics": final_metrics,