Add training debug functionality

docs/training_configuration.md

@@ -129,6 +129,7 @@ above.
 | `mmap_mode` | string or `null` | `"r"` | NumPy memory-map mode used when loading bundle arrays. Use `"r"` to avoid reading the complete input into memory, or `null` to load it normally. |
 | `seed` | integer | `123` | Seeds Python, NumPy, and PyTorch. It also seeds a config-defined random split unless `split.seed` is set. |
 | `progress` | boolean | `true` | Whether to print data-processing, batch, epoch, and evaluation progress. |
+| `debug_epoch_predictions` | boolean | `false` | Save deterministic end-of-epoch train and validation predictions to `debug_epoch_predictions.csv`. This adds one evaluation pass over the training split per epoch. |
 | `sequence_controls` | mapping, list, or `null` | `null` | Optional sequence ablations applied before split selection. |
 | `split_source` | string | `"auto"` | Whether splits come from the bundle or from the `split` block. |
 | `split` | mapping | random 80/10/10 | Config-defined split settings, used according to `split_source`. |
@@ -175,6 +176,27 @@ an epoch is considered improved when validation loss decreases or validation
 Pearson correlation increases. That epoch replaces `best.pt` and resets
 patience. `last.pt` is written after every epoch.
 
+### Epoch Prediction Debugging
+
+Set `"debug_epoch_predictions": true` to write
+`debug_epoch_predictions.csv` in the training output directory. The file has
+one row per train or validation example per completed epoch. Predictions are
+made in evaluation mode using the model at the end of the epoch, so dropout
+and stochastic augmentation are disabled and all training examples are
+included.
+
+The columns include `epoch`, `split`, original dataset `index`, `id`, `target`,
+`prediction`, `residual`, `squared_error`, configured split-level `loss` and
+`pearson`, the corresponding online `history_loss` and `history_pearson`,
+`loss_name`, and checkpoint-monitoring context. Split-level metrics are
+repeated on each example row to keep the CSV self-contained.
+
+For the training split, `loss` and `pearson` can differ from the
+`history_loss` and `history_pearson` columns. History metrics are collected
+while batches are being trained and the model parameters are changing, with
+augmentation and dropout active when configured. The debug metrics instead
+evaluate the final model state for that epoch deterministically.
+
 ## Saluki Model Parameters
 
 Saluki dataset bundles normally have six channels: A, C, G, U, CDS codon

scripts/example_train_config.json

@@ -22,5 +22,6 @@
   "device": "auto",
   "num_workers": 0,
   "mmap_mode": "r",
-  "seed": 42
+  "seed": 42,
+  "debug_epoch_predictions": false
 }

src/transcriptml/training/trainer.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import csv
 import json
 import random
 from dataclasses import asdict, dataclass, field
@@ -41,6 +42,7 @@ class TrainConfig:
     mmap_mode: str | None = "r"
     seed: int = 123
     progress: bool = True
+    debug_epoch_predictions: bool = False
     sequence_controls: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None
     split_source: str = "auto"
     split: Mapping[str, Any] = field(
@@ -256,7 +258,8 @@ def _run_loader(
     target_metrics: bool = True,
     progress: bool = True,
     progress_label: str | None = None,
-) -> dict[str, float]:
+    return_predictions: bool = False,
+) -> dict[str, float | np.ndarray | None]:
     """Run one train or evaluation pass over a loader.
 
     Args:
@@ -271,10 +274,18 @@ def _run_loader(
             Pearson correlation.
         progress: Whether to emit progress messages while iterating.
         progress_label: Optional label shown in progress messages.
+        return_predictions: Whether to include concatenated predictions and
+            targets in the returned mapping.
     """
 
     if loader is None:
-        return {"loss": float("nan"), "pearson": float("nan")}
+        result: dict[str, float | np.ndarray | None] = {
+            "loss": float("nan"),
+            "pearson": float("nan"),
+        }
+        if return_predictions:
+            result.update({"predictions": np.array([]), "targets": None})
+        return result
     training = optimizer is not None
     model.train(training)
     loss_numerator = 0.0
@@ -312,10 +323,106 @@ def _run_loader(
     reporter.close()
     y_pred = np.concatenate(preds) if preds else np.array([])
     y_true = np.concatenate(targets) if targets else np.array([])
-    return {
+    result: dict[str, float | np.ndarray | None] = {
         "loss": float(loss_numerator / max(loss_denominator, 1e-12)),
         "pearson": pearson_corr(y_true, y_pred) if target_metrics else float("nan"),
     }
+    if return_predictions:
+        result.update(
+            {
+                "predictions": y_pred,
+                "targets": y_true if target_metrics else None,
+            }
+        )
+    return result
+
+
+_DEBUG_PREDICTION_FIELDS = (
+    "epoch",
+    "split",
+    "index",
+    "id",
+    "target",
+    "prediction",
+    "residual",
+    "squared_error",
+    "loss",
+    "pearson",
+    "history_loss",
+    "history_pearson",
+    "loss_name",
+    "evaluation_mode",
+    "monitor_improved",
+)
+
+
+def _initialize_debug_predictions_csv(path: str | Path) -> None:
+    """Create an empty epoch-prediction CSV with its header."""
+
+    with Path(path).open("w", newline="", encoding="utf-8") as handle:
+        csv.DictWriter(handle, fieldnames=_DEBUG_PREDICTION_FIELDS).writeheader()
+
+
+def _append_debug_predictions(
+    path: str | Path,
+    *,
+    epoch: int,
+    split: str,
+    indices: Sequence[int],
+    ids: Sequence[str],
+    metrics: Mapping[str, float | np.ndarray | None],
+    history_loss: float,
+    history_pearson: float,
+    loss_name: str,
+    monitor_improved: bool,
+) -> None:
+    """Append deterministic end-of-epoch predictions for one dataset split."""
+
+    predictions = np.asarray(metrics.get("predictions"), dtype=np.float64).reshape(-1)
+    targets_value = metrics.get("targets")
+    targets = (
+        None
+        if targets_value is None
+        else np.asarray(targets_value, dtype=np.float64).reshape(-1)
+    )
+    split_indices = [int(index) for index in indices]
+    if predictions.size != len(split_indices):
+        raise ValueError(
+            f"Debug predictions for split '{split}' have {predictions.size} rows; "
+            f"expected {len(split_indices)}"
+        )
+    if targets is not None and targets.size != predictions.size:
+        raise ValueError(
+            f"Debug targets for split '{split}' have {targets.size} rows; "
+            f"expected {predictions.size}"
+        )
+
+    loss = float(metrics["loss"])
+    pearson = float(metrics["pearson"])
+    with Path(path).open("a", newline="", encoding="utf-8") as handle:
+        writer = csv.DictWriter(handle, fieldnames=_DEBUG_PREDICTION_FIELDS)
+        for position, (index, prediction) in enumerate(zip(split_indices, predictions)):
+            target = None if targets is None else float(targets[position])
+            residual = None if target is None else target - float(prediction)
+            writer.writerow(
+                {
+                    "epoch": int(epoch),
+                    "split": split,
+                    "index": index,
+                    "id": str(ids[index]),
+                    "target": "" if target is None else target,
+                    "prediction": float(prediction),
+                    "residual": "" if residual is None else residual,
+                    "squared_error": "" if residual is None else residual**2,
+                    "loss": loss,
+                    "pearson": pearson,
+                    "history_loss": float(history_loss),
+                    "history_pearson": float(history_pearson),
+                    "loss_name": loss_name,
+                    "evaluation_mode": True,
+                    "monitor_improved": bool(monitor_improved),
+                }
+            )
 
 
 def _is_better(value: float, best: float | None, monitor: str) -> bool:
@@ -459,6 +566,21 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any])
         num_workers=cfg.num_workers,
         pin_memory=pin_memory,
     )
+    debug_train_loader = (
+        _loader(
+            dataset,
+            splits["train"],
+            cfg.batch_size,
+            shuffle=False,
+            num_workers=cfg.num_workers,
+            pin_memory=pin_memory,
+        )
+        if cfg.debug_epoch_predictions
+        else None
+    )
+    debug_predictions_path = out / "debug_epoch_predictions.csv"
+    if cfg.debug_epoch_predictions:
+        _initialize_debug_predictions_csv(debug_predictions_path)
     history: list[dict[str, float | int]] = []
     monitors = _monitor_names(cfg.monitor)
     best_metrics: dict[str, float | None] = {name: None for name in monitors}
@@ -486,6 +608,7 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any])
             target_metrics=has_targets,
             progress=cfg.progress,
             progress_label=f"epoch {epoch} val",
+            return_predictions=cfg.debug_epoch_predictions,
         )
         row = {
             "epoch": epoch,
@@ -496,6 +619,42 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any])
         }
         history.append(row)
         improved, monitor_values = _monitor_improved(row, monitors, best_metrics)
+        if cfg.debug_epoch_predictions:
+            debug_train_metrics = _run_loader(
+                model,
+                debug_train_loader,
+                device=device,
+                loss_fn=loss_fn,
+                optimizer=None,
+                target_metrics=has_targets,
+                progress=cfg.progress,
+                progress_label=f"epoch {epoch} debug train",
+                return_predictions=True,
+            )
+            _append_debug_predictions(
+                debug_predictions_path,
+                epoch=epoch,
+                split="train",
+                indices=splits["train"],
+                ids=bundle.ids,
+                metrics=debug_train_metrics,
+                history_loss=float(row["train_loss"]),
+                history_pearson=float(row["train_pearson"]),
+                loss_name=str(normalized_loss_config["name"]),
+                monitor_improved=improved,
+            )
+            _append_debug_predictions(
+                debug_predictions_path,
+                epoch=epoch,
+                split="val",
+                indices=splits["val"],
+                ids=bundle.ids,
+                metrics=val_metrics,
+                history_loss=float(row["val_loss"]),
+                history_pearson=float(row["val_pearson"]),
+                loss_name=str(normalized_loss_config["name"]),
+                monitor_improved=improved,
+            )
         if improved:
             best_metrics = dict(monitor_values)
             best_epoch = epoch
@@ -596,6 +755,8 @@ def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any])
         "test_mse": test_result.get("loss"),
         "test_pearson": test_result.get("pearson"),
     }
+    if cfg.debug_epoch_predictions:
+        summary["debug_epoch_predictions"] = str(debug_predictions_path)
     if sequence_control_stats is not None:
         summary["sequence_controls"] = sequence_control_stats
     (out / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

tests/test_training_losses.py

@@ -1,3 +1,4 @@
+import csv
 import math
 
 import numpy as np
@@ -134,6 +135,70 @@ def test_train_model_default_mse_remains_compatible(tmp_path):
     assert "best_monitors" not in result["summary"]
     assert result["summary"]["test_loss"] == pytest.approx(result["summary"]["test_mse"])
     assert (tmp_path / "best.pt").exists()
+    assert not (tmp_path / "debug_epoch_predictions.csv").exists()
+
+
+def test_train_model_debug_epoch_predictions_csv(tmp_path):
+    y = np.linspace(-1.0, 1.0, 8, dtype=np.float32)
+    bundle = DatasetBundle(
+        X=_tiny_x(8),
+        y=y,
+        ids=[f"tx-{i}" for i in range(8)],
+        schema="rna4",
+        splits={"train": [0, 1, 2, 3], "val": [4, 5], "test": [6, 7]},
+    )
+
+    result = train_model(
+        bundle,
+        {
+            "dataset": "unused",
+            "output_dir": str(tmp_path),
+            "model": _tiny_model_config(),
+            "batch_size": 3,
+            "epochs": 2,
+            "patience": -1,
+            "progress": False,
+            "debug_epoch_predictions": True,
+        },
+    )
+
+    debug_path = tmp_path / "debug_epoch_predictions.csv"
+    assert result["summary"]["debug_epoch_predictions"] == str(debug_path)
+    with debug_path.open(newline="", encoding="utf-8") as handle:
+        rows = list(csv.DictReader(handle))
+
+    assert len(rows) == 2 * (4 + 2)
+    assert {row["epoch"] for row in rows} == {"1", "2"}
+    assert {row["split"] for row in rows} == {"train", "val"}
+    assert {row["loss_name"] for row in rows} == {"mse"}
+    assert {row["evaluation_mode"] for row in rows} == {"True"}
+    assert {row["id"] for row in rows if row["split"] == "val"} == {"tx-4", "tx-5"}
+
+    for epoch in (1, 2):
+        for split, expected_indices in (("train", [0, 1, 2, 3]), ("val", [4, 5])):
+            group = [
+                row
+                for row in rows
+                if int(row["epoch"]) == epoch and row["split"] == split
+            ]
+            assert [int(row["index"]) for row in group] == expected_indices
+            targets = np.asarray([float(row["target"]) for row in group])
+            predictions = np.asarray([float(row["prediction"]) for row in group])
+            squared_errors = np.asarray([float(row["squared_error"]) for row in group])
+            assert squared_errors == pytest.approx((targets - predictions) ** 2)
+            assert float(group[0]["loss"]) == pytest.approx(float(np.mean(squared_errors)))
+            assert float(group[0]["pearson"]) == pytest.approx(
+                float(np.corrcoef(targets, predictions)[0, 1])
+            )
+            assert len({row["loss"] for row in group}) == 1
+            assert len({row["pearson"] for row in group}) == 1
+
+    val_rows_epoch_1 = [
+        row for row in rows if row["epoch"] == "1" and row["split"] == "val"
+    ]
+    assert float(val_rows_epoch_1[0]["history_loss"]) == pytest.approx(
+        float(val_rows_epoch_1[0]["loss"])
+    )
 
 
 def test_train_model_drops_singleton_training_batch_for_batchnorm(tmp_path, monkeypatch):