diff --git a/docs/training_configuration.md b/docs/training_configuration.md index 6546167..bce2d93 100644 --- a/docs/training_configuration.md +++ b/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 diff --git a/scripts/example_train_config.json b/scripts/example_train_config.json index 7f7194d..d2aa9b9 100644 --- a/scripts/example_train_config.json +++ b/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 } diff --git a/src/transcriptml/training/trainer.py b/src/transcriptml/training/trainer.py index 026dbcd..00ac296 100644 --- a/src/transcriptml/training/trainer.py +++ b/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") diff --git a/tests/test_training_losses.py b/tests/test_training_losses.py index 8f2b78b..5a733e3 100644 --- a/tests/test_training_losses.py +++ b/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):