Add seaborn/matplotlib metric plots, auto-generated after each eval
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) <noreply@anthropic.com>
This commit is contained in:
pikaliov
191
src/training/plot_metrics.py
Normal file
191
src/training/plot_metrics.py
Normal file
@@ -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()
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user