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

Add train recall/AP to train.csv (merged with loss metrics per epoch)

Browse Source 
train.csv now includes eval_loss, r@1_q2g, r@5_q2g, r@10_q2g, ap_q2g
alongside training loss/temperature/gates when eval runs that epoch.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
pikaliov
2026-04-24 12:55:36 +03:00
parent b6dccbba7b
commit d9c795d9ae
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/training/train_gtauav.py
View File

@@ -733,16 +733,13 @@ def train(cfg: TrainConfigGTAUAV) -> None:
"train": means, "train": means,
} }
# Log train metrics to CSV + generate plots every epoch.
csv_logger.log_train(epoch, means, optimizer.param_groups[0]["lr"], elapsed)
generate_plots(csv_logger.log_dir)
# --- Log VRAM usage --- # --- Log VRAM usage ---
if torch.cuda.is_available(): if torch.cuda.is_available():
vram_gb = torch.cuda.max_memory_allocated() / 1e9 vram_gb = torch.cuda.max_memory_allocated() / 1e9
tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step) tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step)
# Evaluation. # Evaluation.
train_recall = {}
if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1: if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1:
# Train R@K (subset — same size as test set for speed). # Train R@K (subset — same size as test set for speed).
train_eval_batches = len(test_loader) train_eval_batches = len(test_loader)
@@ -754,6 +751,17 @@ def train(cfg: TrainConfigGTAUAV) -> None:
epoch_record["train_recall"] = train_recall epoch_record["train_recall"] = train_recall
csv_logger.log_train_recall(epoch, 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) tracker.log_train(epoch, {f"recall/{k}": v for k, v in train_recall.items() if k.startswith("r@")}, step=global_step)
# Log train metrics to CSV (includes recall/AP if eval ran this epoch).
train_row = {**means}
if train_recall:
train_row["eval_loss"] = train_recall.get("loss", 0.0)
train_row["r@1_q2g"] = train_recall.get("r@1_q2g", 0.0)
train_row["r@5_q2g"] = train_recall.get("r@5_q2g", 0.0)
train_row["r@10_q2g"] = train_recall.get("r@10_q2g", 0.0)
train_row["ap_q2g"] = train_recall.get("ap_q2g", 0.0)
csv_logger.log_train(epoch, train_row, optimizer.param_groups[0]["lr"], elapsed)
generate_plots(csv_logger.log_dir)
LOGGER.info( LOGGER.info(
"train-recall epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f", "train-recall epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f",
epoch, epoch,