From 83ce04150d01ebf222f225ccf359b650f38f1b63 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 21 Apr 2026 19:54:18 +0300 Subject: [PATCH] Add seaborn/matplotlib metric plots, auto-generated after each eval MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit New: src/training/plot_metrics.py - train_metrics.png: loss, temperature, gates, lr - val_metrics.png: R@K q→g and g→q - overview.png: combined loss + R@1 + gates/tau Auto-generates plots in {output_dir}/logs/ after each validation epoch. Also callable standalone: python -m src.training.plot_metrics --log-dir ... Co-Authored-By: Claude Opus 4.6 (1M context) --- src/training/plot_metrics.py | 191 +++++++++++++++++++++++++++++++++++ src/training/train_gtauav.py | 2 + 2 files changed, 193 insertions(+) create mode 100644 src/training/plot_metrics.py diff --git a/src/training/plot_metrics.py b/src/training/plot_metrics.py new file mode 100644 index 0000000..a9787bf --- /dev/null +++ b/src/training/plot_metrics.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +"""Generate training/validation metric plots from CSV logs. + +Reads train.csv and val.csv, produces publication-quality plots +using seaborn + matplotlib. Called automatically after each eval epoch +or standalone via CLI. + +Usage: + python -m src.training.plot_metrics --log-dir out/gtauav/with_text/logs +""" + +import argparse +import logging +from pathlib import Path + +import matplotlib.pyplot as plt +import pandas as pd +import seaborn as sns + +LOGGER = logging.getLogger("caption_test.plot") + +# Seaborn style. +sns.set_theme(style="whitegrid", palette="deep", font_scale=1.1) + + +def plot_train_metrics(train_df: pd.DataFrame, out_dir: Path) -> None: + """Plot training metrics: loss, temperature, gates, lr.""" + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + fig.suptitle("Training Metrics", fontsize=16, fontweight="bold") + + # 1. Loss. + ax = axes[0, 0] + sns.lineplot(data=train_df, x="epoch", y="total", ax=ax, marker="o", linewidth=2) + ax.set_title("Loss") + ax.set_xlabel("Epoch") + ax.set_ylabel("InfoNCE Loss") + + # 2. Temperature (tau). + ax = axes[0, 1] + sns.lineplot(data=train_df, x="epoch", y="temperature", ax=ax, marker="s", linewidth=2, color="orange") + ax.set_title("Temperature (τ)") + ax.set_xlabel("Epoch") + ax.set_ylabel("τ") + + # 3. Gate values. + ax = axes[1, 0] + if "gate_q" in train_df.columns and "gate_g" in train_df.columns: + sns.lineplot(data=train_df, x="epoch", y="gate_q", ax=ax, marker="o", linewidth=2, label="gate_q (drone)") + sns.lineplot(data=train_df, x="epoch", y="gate_g", ax=ax, marker="s", linewidth=2, label="gate_g (sat)") + ax.legend() + ax.set_title("Gate Values (σ(α))") + ax.set_xlabel("Epoch") + ax.set_ylabel("Image weight") + ax.set_ylim(0, 1) + + # 4. Learning rate. + ax = axes[1, 1] + sns.lineplot(data=train_df, x="epoch", y="lr", ax=ax, marker="^", linewidth=2, color="green") + ax.set_title("Learning Rate") + ax.set_xlabel("Epoch") + ax.set_ylabel("LR") + ax.ticklabel_format(axis="y", style="scientific", scilimits=(0, 0)) + + plt.tight_layout() + path = out_dir / "train_metrics.png" + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + LOGGER.info("📊 Saved %s", path) + + +def plot_val_metrics(val_df: pd.DataFrame, out_dir: Path) -> None: + """Plot validation metrics: R@K, gates.""" + fig, axes = plt.subplots(1, 2, figsize=(14, 5)) + fig.suptitle("Validation Metrics", fontsize=16, fontweight="bold") + + # 1. Recall@K (q→g). + ax = axes[0] + for k in [1, 5, 10]: + col = f"r@{k}_q2g" + if col in val_df.columns: + sns.lineplot(data=val_df, x="epoch", y=col, ax=ax, marker="o", linewidth=2, label=f"R@{k}") + ax.set_title("Recall@K (drone → satellite)") + ax.set_xlabel("Epoch") + ax.set_ylabel("Recall") + ax.set_ylim(0, 1) + ax.legend() + + # 2. Recall@K (g→q). + ax = axes[1] + for k in [1, 5, 10]: + col = f"r@{k}_g2q" + if col in val_df.columns: + sns.lineplot(data=val_df, x="epoch", y=col, ax=ax, marker="s", linewidth=2, label=f"R@{k}") + ax.set_title("Recall@K (satellite → drone)") + ax.set_xlabel("Epoch") + ax.set_ylabel("Recall") + ax.set_ylim(0, 1) + ax.legend() + + plt.tight_layout() + path = out_dir / "val_metrics.png" + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + LOGGER.info("📊 Saved %s", path) + + +def plot_combined(train_df: pd.DataFrame, val_df: pd.DataFrame, out_dir: Path) -> None: + """Combined overview: loss + R@1 on same figure.""" + fig, axes = plt.subplots(1, 3, figsize=(18, 5)) + fig.suptitle("Training Overview", fontsize=16, fontweight="bold") + + # 1. Train 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") + ax.set_xlabel("Epoch") + ax.set_ylabel("InfoNCE Loss") + + # 2. Val R@1 q→g. + ax = axes[1] + if "r@1_q2g" in val_df.columns: + sns.lineplot(data=val_df, x="epoch", y="r@1_q2g", ax=ax, marker="o", linewidth=2, color="royalblue", label="R@1 q→g") + if "r@1_g2q" in val_df.columns: + sns.lineplot(data=val_df, x="epoch", y="r@1_g2q", ax=ax, marker="s", linewidth=2, color="coral", label="R@1 g→q") + ax.set_title("Validation R@1") + ax.set_xlabel("Epoch") + ax.set_ylabel("Recall@1") + ax.set_ylim(0, None) + ax.legend() + + # 3. Gates + Temperature. + ax = axes[2] + if "gate_q" in train_df.columns: + sns.lineplot(data=train_df, x="epoch", y="gate_q", ax=ax, linewidth=2, label="gate_q") + sns.lineplot(data=train_df, x="epoch", y="gate_g", ax=ax, linewidth=2, label="gate_g") + ax2 = ax.twinx() + sns.lineplot(data=train_df, x="epoch", y="temperature", ax=ax2, linewidth=2, color="orange", linestyle="--", label="τ") + ax.set_title("Gates & Temperature") + ax.set_xlabel("Epoch") + ax.set_ylabel("Gate value") + ax2.set_ylabel("τ") + lines1, labels1 = ax.get_legend_handles_labels() + lines2, labels2 = ax2.get_legend_handles_labels() + ax.legend(lines1 + lines2, labels1 + labels2, loc="best") + ax2.get_legend().remove() if ax2.get_legend() else None + + plt.tight_layout() + path = out_dir / "overview.png" + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + LOGGER.info("📊 Saved %s", path) + + +def generate_plots(log_dir: str | Path) -> None: + """Generate all plots from CSV logs in log_dir.""" + log_dir = Path(log_dir) + + train_csv = log_dir / "train.csv" + val_csv = log_dir / "val.csv" + + if train_csv.exists(): + train_df = pd.read_csv(train_csv) + plot_train_metrics(train_df, log_dir) + else: + LOGGER.warning("No train.csv found in %s", log_dir) + train_df = None + + if val_csv.exists(): + val_df = pd.read_csv(val_csv) + plot_val_metrics(val_df, log_dir) + else: + LOGGER.warning("No val.csv found in %s", log_dir) + val_df = None + + if train_df is not None and val_df is not None: + plot_combined(train_df, val_df, log_dir) + + +def main() -> None: + import coloredlogs + coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s") + + parser = argparse.ArgumentParser(description="Generate metric plots from CSV logs.") + parser.add_argument("--log-dir", type=str, required=True, help="Path to logs/ directory.") + args = parser.parse_args() + generate_plots(args.log_dir) + + +if __name__ == "__main__": + main() diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 525d60e..c2736b5 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -27,6 +27,7 @@ from tqdm import tqdm from src.datasets.gtauav_dataset import GTAUAVDataset, collate_gtauav_batch from src.losses.multi_infonce import InfoNCELoss +from src.training.plot_metrics import generate_plots from src.models.asymmetric_encoder import ( AsymmetricEncoder, get_dino_transform, @@ -482,6 +483,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: val_metrics = _evaluate(model, test_loader, cfg.device) epoch_record["val"] = val_metrics csv_logger.log_val(epoch, val_metrics) + generate_plots(csv_logger.log_dir) LOGGER.info( "🎯 val epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f gate_q=%.4f gate_g=%.4f", epoch,