diff --git a/src/training/plot_metrics.py b/src/training/plot_metrics.py index bbe39ec..37c45ae 100644 --- a/src/training/plot_metrics.py +++ b/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] diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index c57e60a..f4f8358 100644 --- a/src/training/train_gtauav.py +++ b/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,