diff --git a/src/training/plot_metrics.py b/src/training/plot_metrics.py index 37c45ae..df97c62 100644 --- a/src/training/plot_metrics.py +++ b/src/training/plot_metrics.py @@ -70,9 +70,9 @@ def plot_train_metrics(train_df: pd.DataFrame, out_dir: Path) -> None: def plot_val_metrics(val_df: pd.DataFrame, out_dir: Path, train_recall_df: pd.DataFrame | None = None) -> None: - """Plot recall metrics: train vs val R@K.""" - fig, axes = plt.subplots(1, 2, figsize=(14, 5)) - fig.suptitle("Recall Metrics (train vs val)", fontsize=16, fontweight="bold") + """Plot recall + AP metrics: train vs val.""" + fig, axes = plt.subplots(1, 3, figsize=(20, 5)) + fig.suptitle("Retrieval Metrics (train vs val)", fontsize=16, fontweight="bold") # 1. Recall@K (q→g): train + val. ax = axes[0] @@ -102,6 +102,23 @@ def plot_val_metrics(val_df: pd.DataFrame, out_dir: Path, train_recall_df: pd.Da ax.set_ylim(0, 1) ax.legend(fontsize=8) + # 3. AP (train vs val, both directions). + ax = axes[2] + if "ap_q2g" in val_df.columns: + sns.lineplot(data=val_df, x="epoch", y="ap_q2g", ax=ax, marker="o", linewidth=2, color="royalblue", label="val AP q→g") + if "ap_g2q" in val_df.columns: + sns.lineplot(data=val_df, x="epoch", y="ap_g2q", ax=ax, marker="s", linewidth=2, color="coral", label="val AP g→q") + if train_recall_df is not None: + if "ap_q2g" in train_recall_df.columns: + sns.lineplot(data=train_recall_df, x="epoch", y="ap_q2g", ax=ax, marker="x", linewidth=1.5, linestyle="--", color="royalblue", label="train AP q→g") + if "ap_g2q" in train_recall_df.columns: + sns.lineplot(data=train_recall_df, x="epoch", y="ap_g2q", ax=ax, marker="x", linewidth=1.5, linestyle="--", color="coral", label="train AP g→q") + ax.set_title("Average Precision") + ax.set_xlabel("Epoch") + ax.set_ylabel("AP") + ax.set_ylim(0, 1) + ax.legend(fontsize=8) + plt.tight_layout() path = out_dir / "val_metrics.png" fig.savefig(path, dpi=150, bbox_inches="tight") diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index f4f8358..7950bf9 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -228,18 +228,27 @@ def _evaluate( if batch_losses: metrics["loss"] = sum(batch_losses) / len(batch_losses) + # R@K and AP (q→g). sorted_idx = sim.argsort(dim=1, descending=True) for k in k_values: top_k = sorted_idx[:, :k] hit = (top_k == targets.unsqueeze(1)).any(dim=1).float() metrics[f"r@{k}_q2g"] = float(hit.mean().item()) + # AP q→g: for each query, rank of the correct gallery = 1/(rank+1). + ranks_q2g = (sorted_idx == targets.unsqueeze(1)).nonzero(as_tuple=True)[1].float() + metrics["ap_q2g"] = float((1.0 / (ranks_q2g + 1)).mean().item()) + + # R@K and AP (g→q). sorted_idx_g2q = sim.t().argsort(dim=1, descending=True) for k in k_values: top_k = sorted_idx_g2q[:, :k] hit = (top_k == targets.unsqueeze(1)).any(dim=1).float() metrics[f"r@{k}_g2q"] = float(hit.mean().item()) + ranks_g2q = (sorted_idx_g2q == targets.unsqueeze(1)).nonzero(as_tuple=True)[1].float() + metrics["ap_g2q"] = float((1.0 / (ranks_g2q + 1)).mean().item()) + metrics["gate_q"] = model.fusion_query.gate_value metrics["gate_g"] = model.fusion_gallery.gate_value return metrics @@ -696,11 +705,13 @@ def train(cfg: TrainConfigGTAUAV) -> None: csv_logger.log_train_recall(epoch, train_recall) tracker.log_train(epoch, {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, step=global_step) LOGGER.info( - "train-recall epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f", + "train-recall epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f", epoch, train_recall.get("r@1_q2g", 0.0), train_recall.get("r@5_q2g", 0.0), train_recall.get("r@10_q2g", 0.0), + train_recall.get("ap_q2g", 0.0), + train_recall.get("loss", 0.0), ) # Val R@K (full test set). @@ -721,13 +732,14 @@ def train(cfg: TrainConfigGTAUAV) -> None: tracker.log_scalar("val/best_r@1_q2g", best_r1, step=global_step) LOGGER.info( - "val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f", + "val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f gate_q=%.4f", epoch, val_metrics.get("r@1_q2g", 0.0), val_metrics.get("r@5_q2g", 0.0), val_metrics.get("r@10_q2g", 0.0), + val_metrics.get("ap_q2g", 0.0), + val_metrics.get("loss", 0.0), val_metrics.get("gate_q", 1.0), - val_metrics.get("gate_g", 1.0), ) # --- Grad-CAM visualization ---