Pikaliov/caption-test
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add train recall evaluation (R@K on train subset each epoch)

Browse Source 
- Evaluate R@K on train set (subset matching test size) alongside val
- New train_recall.csv with per-epoch train R@1/R@5/R@10
- Plot train vs val recall on same chart (solid=val, dashed=train)
- Helps detect overfitting: train R@1 up + val R@1 flat = overfit
- Train eval uses clean transforms (no augmentation)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-22 08:01:09 +03:00
parent 66c83469ef
commit 93ad66810d
2 changed files with 72 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
src/training/plot_metrics.py
View File

@@ -69,34 +69,38 @@ def plot_train_metrics(train_df: pd.DataFrame, out_dir: Path) -> None:
LOGGER.info("📊 Saved %s", path)
def plot_val_metrics(val_df: pd.DataFrame, out_dir: Path) -> None:
"""Plot validation metrics: R@K, gates."""
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("Validation Metrics", fontsize=16, fontweight="bold")
fig.suptitle("Recall Metrics (train vs val)", fontsize=16, fontweight="bold")
# 1. Recall@K (q→g).
# 1. Recall@K (q→g): train + val.
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}")
sns.lineplot(data=val_df, x="epoch", y=col, ax=ax, marker="o", linewidth=2, label=f"val R@{k}")
if train_recall_df is not None and col in train_recall_df.columns:
sns.lineplot(data=train_recall_df, x="epoch", y=col, ax=ax, marker="x", linewidth=1.5, linestyle="--", label=f"train R@{k}")
ax.set_title("Recall@K (drone → satellite)")
ax.set_xlabel("Epoch")
ax.set_ylabel("Recall")
ax.set_ylim(0, 1)
ax.legend()
ax.legend(fontsize=8)
# 2. Recall@K (g→q).
# 2. Recall@K (g→q): train + val.
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}")
sns.lineplot(data=val_df, x="epoch", y=col, ax=ax, marker="s", linewidth=2, label=f"val R@{k}")
if train_recall_df is not None and col in train_recall_df.columns:
sns.lineplot(data=train_recall_df, x="epoch", y=col, ax=ax, marker="x", linewidth=1.5, linestyle="--", label=f"train R@{k}")
ax.set_title("Recall@K (satellite → drone)")
ax.set_xlabel("Epoch")
ax.set_ylabel("Recall")
ax.set_ylim(0, 1)
ax.legend()
ax.legend(fontsize=8)
plt.tight_layout()
path = out_dir / "val_metrics.png"
@@ -166,9 +170,14 @@ def generate_plots(log_dir: str | Path) -> None:
LOGGER.warning("No train.csv found in %s", log_dir)
train_df = None
train_recall_csv = log_dir / "train_recall.csv"
train_recall_df = None
if train_recall_csv.exists():
train_recall_df = pd.read_csv(train_recall_csv)
if val_csv.exists():
val_df = pd.read_csv(val_csv)
plot_val_metrics(val_df, log_dir)
plot_val_metrics(val_df, log_dir, train_recall_df=train_recall_df)
else:
LOGGER.warning("No val.csv found in %s", log_dir)
val_df = None

56
src/training/train_gtauav.py
View File

@@ -176,13 +176,17 @@ def _evaluate(
loader: DataLoader,
device: str,
k_values: tuple[int, ...] = (1, 5, 10),
max_batches: int | None = None,
desc: str = "eval",
) -> dict[str, float]:
"""Compute R@K on validation set."""
"""Compute R@K. Use max_batches to limit for train set (faster)."""
model.eval()
all_query: list[torch.Tensor] = []
all_gallery: list[torch.Tensor] = []
for batch in tqdm(loader, desc=" eval", unit="batch", leave=False):
for i, batch in enumerate(tqdm(loader, desc=f" {desc}", unit="batch", leave=False)):
if max_batches is not None and i >= max_batches:
break
drone_img = batch["drone_img"].to(device, non_blocking=True)
sat_img = batch["sat_img"].to(device, non_blocking=True)
@@ -248,6 +252,7 @@ class CSVLogger:
# Load existing CSV data on resume (so plots show full history).
train_csv = self.log_dir / "train.csv"
val_csv = self.log_dir / "val.csv"
train_recall_csv = self.log_dir / "train_recall.csv"
if train_csv.exists():
self.train_rows = pd.read_csv(train_csv).to_dict("records")
LOGGER.info("CSVLogger: loaded %d previous train epochs", len(self.train_rows))
@@ -258,6 +263,10 @@ class CSVLogger:
LOGGER.info("CSVLogger: loaded %d previous val epochs", len(self.val_rows))
else:
self.val_rows = []
if train_recall_csv.exists():
self.train_recall_rows = pd.read_csv(train_recall_csv).to_dict("records")
else:
self.train_recall_rows = []
def log_batch(self, epoch: int, batch_idx: int, global_step: int, metrics: dict) -> None:
"""Log metrics for a single training batch. Writes to disk immediately."""
@@ -300,6 +309,13 @@ class CSVLogger:
self.val_rows.append(row)
pd.DataFrame(self.val_rows).to_csv(self.log_dir / "val.csv", index=False)
def log_train_recall(self, epoch: int, metrics: dict) -> None:
"""Log train recall metrics. Replaces existing entry for same epoch."""
row = {"epoch": epoch, **metrics}
self.train_recall_rows = [r for r in self.train_recall_rows if r.get("epoch") != epoch]
self.train_recall_rows.append(row)
pd.DataFrame(self.train_recall_rows).to_csv(self.log_dir / "train_recall.csv", index=False)
def _clear_vram() -> None:
"""Free VRAM from previous runs before starting."""
@@ -445,6 +461,22 @@ def train(cfg: TrainConfigGTAUAV) -> None:
collate_fn=collate_gtauav_batch,
pin_memory=True,
)
# Train eval loader: clean transforms (no augmentation), for R@K on train set.
train_eval_ds = GTAUAVDataset(
pair_json=cfg.train_json,
rgb_root=cfg.rgb_root,
caption_root=cfg.caption_root,
image_transform=eval_tf,
filter_meta=cfg.filter_meta,
)
train_eval_loader = DataLoader(
train_eval_ds,
batch_size=cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers,
collate_fn=collate_gtauav_batch,
pin_memory=True,
)
effective_batch = cfg.batch_size * cfg.grad_accum_steps
LOGGER.info(
@@ -639,7 +671,25 @@ def train(cfg: TrainConfigGTAUAV) -> None:
# Evaluation.
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
val_metrics = _evaluate(model, test_loader, cfg.device)
# Train R@K (subset — same size as test set for speed).
train_eval_batches = len(test_loader)
train_recall = _evaluate(
model, train_eval_loader, cfg.device,
max_batches=train_eval_batches, desc="eval-train",
)
epoch_record["train_recall"] = train_recall
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",
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),
)
# Val R@K (full test set).
val_metrics = _evaluate(model, test_loader, cfg.device, desc="eval-val")
epoch_record["val"] = val_metrics
csv_logger.log_val(epoch, val_metrics)
generate_plots(csv_logger.log_dir)