From d9c795d9ae9cb4842b60b0988b0dcca8a7e85b98 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Fri, 24 Apr 2026 12:55:36 +0300 Subject: [PATCH] Add train recall/AP to train.csv (merged with loss metrics per epoch) 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) --- src/training/train_gtauav.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 53a34f3..ab3dd92 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -733,16 +733,13 @@ def train(cfg: TrainConfigGTAUAV) -> None: "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 --- if torch.cuda.is_available(): vram_gb = torch.cuda.max_memory_allocated() / 1e9 tracker.log_scalar("system/vram_peak_gb", vram_gb, step=global_step) # Evaluation. + train_recall = {} if (epoch + 1) % cfg.eval_every == 0 or epoch == cfg.epochs - 1: # Train R@K (subset — same size as test set for speed). train_eval_batches = len(test_loader) @@ -754,6 +751,17 @@ def train(cfg: TrainConfigGTAUAV) -> None: 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) + + # 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( "train-recall epoch=%d R@1=%.4f R@5=%.4f R@10=%.4f AP=%.4f loss=%.4f", epoch,