From 7b13a4c4dbf77edeb5a49bc79f5abaf125afbbad Mon Sep 17 00:00:00 2001
From: pikaliov <pikaliov@gmail.com>
Date: Tue, 21 Apr 2026 20:54:48 +0300
Subject: [PATCH] Write per-batch CSV immediately (append mode, no buffering)

train_batches.csv and epoch_N_batches.csv now update after every batch
instead of flushing at epoch end. Uses file append mode for efficiency.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 src/training/train_gtauav.py | 43 +++++++++++++++++++++++++++++++-----
 1 file changed, 38 insertions(+), 5 deletions(-)

diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py
index 82fd3eb..f0ecd2e 100644
--- a/src/training/train_gtauav.py
+++ b/src/training/train_gtauav.py
@@ -228,9 +228,12 @@ class CSVLogger:
     """Log train/val metrics to CSV files using pandas.
 
     Creates:
-        {output_dir}/logs/train.csv       — all train epochs
-        {output_dir}/logs/val.csv         — all val epochs
-        {output_dir}/logs/epoch_{N}.csv   — per-epoch details
+        {output_dir}/logs/train.csv               — epoch-level train averages
+        {output_dir}/logs/val.csv                  — epoch-level val metrics
+        {output_dir}/logs/train_batches.csv        — per-batch train metrics (all epochs)
+        {output_dir}/logs/epoch_{N}_train.csv      — per-epoch summary
+        {output_dir}/logs/epoch_{N}_val.csv        — per-epoch val
+        {output_dir}/logs/epoch_{N}_batches.csv    — per-batch for single epoch
     """
 
     def __init__(self, output_dir: Path) -> None:
@@ -238,13 +241,42 @@ class CSVLogger:
         self.log_dir.mkdir(parents=True, exist_ok=True)
         self.train_rows: list[dict] = []
         self.val_rows: list[dict] = []
+        self._current_epoch: int = -1
+        self._batch_columns: list[str] | None = None
+        self._cumulative_batch_path = self.log_dir / "train_batches.csv"
+        self._epoch_batch_path: Path | None = None
+
+    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."""
+        row = {"epoch": epoch, "batch": batch_idx, "global_step": global_step, **metrics}
+
+        # On new epoch, start a fresh per-epoch CSV.
+        if epoch != self._current_epoch:
+            self._current_epoch = epoch
+            self._epoch_batch_path = self.log_dir / f"epoch_{epoch:03d}_batches.csv"
+
+        # Determine columns on first call (consistent order).
+        if self._batch_columns is None:
+            self._batch_columns = list(row.keys())
+
+        row_df = pd.DataFrame([row], columns=self._batch_columns)
+        write_header = not self._cumulative_batch_path.exists()
+
+        # Append to cumulative CSV.
+        row_df.to_csv(
+            self._cumulative_batch_path, mode="a", header=write_header, index=False,
+        )
+        # Append to per-epoch CSV.
+        write_epoch_header = not self._epoch_batch_path.exists()
+        row_df.to_csv(
+            self._epoch_batch_path, mode="a", header=write_epoch_header, index=False,
+        )
 
     def log_train(self, epoch: int, metrics: dict, lr: float, elapsed: float) -> None:
+        """Log epoch-level train averages."""
         row = {"epoch": epoch, "lr": lr, "elapsed_s": round(elapsed, 1), **metrics}
         self.train_rows.append(row)
-        # Append to cumulative CSV.
         pd.DataFrame(self.train_rows).to_csv(self.log_dir / "train.csv", index=False)
-        # Per-epoch CSV.
         pd.DataFrame([row]).to_csv(self.log_dir / f"epoch_{epoch:03d}_train.csv", index=False)
 
     def log_val(self, epoch: int, metrics: dict) -> None:
@@ -516,6 +548,7 @@ def train(cfg: TrainConfigGTAUAV) -> None:
                 "lr": optimizer.param_groups[0]["lr"],
             }
             tracker.log_train(epoch, step_metrics, step=global_step)
+            csv_logger.log_batch(epoch, n_batches, global_step, step_metrics)
 
             for key, val in loss_dict.items():
                 agg[key] = agg.get(key, 0.0) + float(val.item())