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>
This commit is contained in:
pikaliov
2026-04-21 20:54:48 +03:00
parent 69857e0ade
commit 7b13a4c4db

View File

@@ -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())