from __future__ import annotations
import csv
import json
import random
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any, Mapping, Sequence
import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, Dataset, Subset
from transcriptml.data.bundle import DatasetBundle, load_bundle
from transcriptml.data.controls import apply_sequence_controls_to_bundle
from transcriptml.devices import resolve_device
from transcriptml.models.common import squeeze_prediction
from transcriptml.models.registry import build_model, normalize_model_config, save_checkpoint
from transcriptml.training.evaluation import evaluate_model, predict_to_csv
from transcriptml.training.losses import TrainingLoss, build_training_loss
from transcriptml.training.metrics import pearson_corr
from transcriptml.training.splits import normalize_splits, predefined_split_indices, random_split_indices
from transcriptml.progress import ProgressReporter, log_progress
[docs]
@dataclass
class TrainConfig:
dataset: str
output_dir: str
model: Mapping[str, Any] = field(default_factory=lambda: {"name": "small_cnn", "params": {}})
batch_size: int = 64
epochs: int = 20
learning_rate: float = 1e-3
weight_decay: float = 0.0
optimizer: str | Mapping[str, Any] = "adamw"
lr_scheduler: str | Mapping[str, Any] | None = None
mixed_precision: bool = False
gradient_clip_norm: float | None = 0.5
patience: int = 5
monitor: str | Sequence[str] = "val_loss"
loss: str | Mapping[str, Any] | None = field(default_factory=lambda: {"name": "mse"})
device: str = "cpu"
num_workers: int = 0
mmap_mode: str | None = "r"
seed: int = 123
progress: bool = True
debug_epoch_predictions: bool = False
head_layernorm: bool = False
sequence_controls: Mapping[str, Any] | Sequence[Mapping[str, Any]] | None = None
split_source: str = "auto"
cv_plan: str | None = None
fold: int | None = None
max_train_jitter: int = 0
allow_random_window_split: bool = False
deduplicate_loci: bool = True
split: Mapping[str, Any] = field(
default_factory=lambda: {"method": "random", "val_frac": 0.1, "test_frac": 0.1}
)
def _seed_everything(seed: int) -> None:
"""Seed Python, NumPy, and PyTorch random number generators.
Args:
seed: Integer seed applied across supported random number generators.
"""
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
def _load_config(path: str | Path) -> dict[str, Any]:
"""Load a JSON or TOML training configuration file.
Args:
path: Path to a ``.json`` or ``.toml`` config file.
"""
p = Path(path)
if p.suffix.lower() == ".toml":
try:
import tomllib
except ModuleNotFoundError:
import tomli as tomllib
return tomllib.loads(p.read_text(encoding="utf-8"))
return json.loads(p.read_text(encoding="utf-8"))
def _as_train_config(config: TrainConfig | Mapping[str, Any]) -> TrainConfig:
"""Normalize mapping-like training config to ``TrainConfig``.
Args:
config: Existing ``TrainConfig`` or mapping of constructor fields.
"""
if isinstance(config, TrainConfig):
return config
return TrainConfig(**dict(config))
def _make_config_splits(bundle: DatasetBundle, cfg: TrainConfig) -> dict[str, list[int]]:
"""Choose dataset splits from the training config.
Args:
bundle: Dataset bundle used for size and metadata.
cfg: Training configuration containing split strategy settings.
"""
split_cfg = dict(cfg.split or {})
method = split_cfg.get("method", "random")
if method == "random":
return random_split_indices(
bundle.X.shape[0],
val_frac=float(split_cfg.get("val_frac", 0.1)),
test_frac=float(split_cfg.get("test_frac", 0.1)),
seed=int(split_cfg.get("seed", cfg.seed)),
)
if method == "metadata":
if bundle.metadata is None:
raise ValueError("metadata split requested but bundle has no metadata")
return predefined_split_indices(bundle.metadata, split_col=str(split_cfg.get("split_col", "split")))
if method == "predefined":
return normalize_splits(split_cfg["splits"])
raise ValueError(f"Unknown split method '{method}'")
def _select_splits(bundle: DatasetBundle, cfg: TrainConfig) -> tuple[dict[str, list[int]], str]:
"""Choose dataset splits and report which source was used.
Args:
bundle: Dataset bundle that may already contain predefined splits.
cfg: Training configuration containing split source and strategy.
"""
if cfg.cv_plan is not None or cfg.fold is not None:
if cfg.cv_plan is None or cfg.fold is None:
raise ValueError("cv_plan and fold must be provided together")
if bundle.metadata is None:
raise ValueError("chromosome CV requires dataset bundle metadata")
from transcriptml.workflows.chromosome_cv import (
load_chromosome_cv_plan,
resolve_chromosome_cv_plan,
)
plan = load_chromosome_cv_plan(cfg.cv_plan)
resolution = resolve_chromosome_cv_plan(
plan, bundle.metadata, fold=cfg.fold
)
return normalize_splits(resolution.indices), "cv_plan"
source = str(cfg.split_source or "auto").strip().lower()
if source not in {"auto", "bundle", "config"}:
raise ValueError("split_source must be one of: auto, bundle, config")
has_bundle_splits = bundle.splits is not None
if source == "bundle":
if not has_bundle_splits:
raise ValueError("split_source='bundle' requested but dataset bundle has no splits")
return normalize_splits(bundle.splits), "bundle"
if source == "config":
return _make_config_splits(bundle, cfg), "config"
if has_bundle_splits:
return normalize_splits(bundle.splits), "bundle"
return _make_config_splits(bundle, cfg), "config"
def _make_splits(bundle: DatasetBundle, cfg: TrainConfig) -> dict[str, list[int]]:
"""Choose dataset splits from the bundle or training config."""
splits, _ = _select_splits(bundle, cfg)
return splits
def _cv_plan_id(cfg: TrainConfig) -> str | None:
"""Return the validated plan identifier recorded by CV training artifacts."""
if cfg.cv_plan is None:
return None
from transcriptml.workflows.chromosome_cv import load_chromosome_cv_plan
return load_chromosome_cv_plan(cfg.cv_plan).plan_id
class _ArrayRegressionDataset(Dataset):
def __init__(self, X: np.ndarray, y: np.ndarray, aux_arrays: Mapping[str, np.ndarray] | None = None):
"""Wrap NumPy arrays as a PyTorch regression dataset.
Args:
X: Encoded input array with examples on axis 0.
y: Scalar target array aligned to ``X``.
aux_arrays: Optional per-example auxiliary arrays used by
metadata-aware losses.
"""
self.X = X
self.y = y
self.aux_arrays = dict(aux_arrays or {})
def __len__(self) -> int:
"""Return the number of examples."""
return int(self.X.shape[0])
def __getitem__(self, idx: int) -> tuple[np.ndarray, np.float32, dict[str, np.float32]]:
"""Return one input array and scalar target.
Args:
idx: Integer example index to retrieve.
"""
i = int(idx)
aux = {name: np.float32(values[i]) for name, values in self.aux_arrays.items()}
return np.asarray(self.X[i]), np.float32(self.y[i]), aux
def _collate_regression(
batch: list[tuple[np.ndarray, np.float32, Mapping[str, np.float32]]],
) -> tuple[torch.Tensor, torch.Tensor, dict[str, torch.Tensor]]:
"""Stack NumPy regression examples into tensors.
Args:
batch: List of ``(input_array, scalar_target, aux)`` examples from the
dataset.
"""
xs, ys, auxs = zip(*batch)
aux_keys = auxs[0].keys() if auxs else []
aux = {
key: torch.as_tensor(np.asarray([item[key] for item in auxs], dtype=np.float32))
for key in aux_keys
}
return torch.as_tensor(np.stack(xs, axis=0)), torch.as_tensor(np.asarray(ys, dtype=np.float32)), aux
def _loader(
dataset: Dataset,
indices: list[int],
batch_size: int,
*,
shuffle: bool,
num_workers: int = 0,
pin_memory: bool = False,
drop_last: bool = False,
) -> DataLoader | None:
"""Create a DataLoader for a split, or ``None`` for empty splits.
Args:
dataset: PyTorch dataset containing all examples.
indices: Example indices assigned to the split.
batch_size: Number of examples per batch.
shuffle: Whether to shuffle the split each epoch.
num_workers: Number of worker processes used by the DataLoader.
pin_memory: Whether the DataLoader should pin host memory.
drop_last: Whether to drop the final incomplete batch.
"""
if not indices:
return None
return DataLoader(
Subset(dataset, indices),
batch_size=batch_size,
shuffle=shuffle,
num_workers=int(num_workers),
pin_memory=pin_memory,
persistent_workers=int(num_workers) > 0,
collate_fn=_collate_regression,
drop_last=drop_last,
)
def _would_create_singleton_batch(n_examples: int, batch_size: int) -> bool:
"""Return whether batching would leave one training example by itself."""
n = int(n_examples)
bs = int(batch_size)
if bs <= 1 or n <= bs:
return False
return n % bs == 1
def _has_batch_normalization(model: nn.Module) -> bool:
"""Return whether a model contains a PyTorch batch-normalization module."""
return any(isinstance(module, nn.modules.batchnorm._BatchNorm) for module in model.modules())
def _run_loader(
model: nn.Module,
loader: DataLoader | None,
*,
device: torch.device,
loss_fn: TrainingLoss,
optimizer: torch.optim.Optimizer | None = None,
gradient_clip_norm: float | None = None,
target_metrics: bool = True,
progress: bool = True,
progress_label: str | None = None,
return_predictions: bool = False,
) -> dict[str, float | np.ndarray | None]:
"""Run one train or evaluation pass over a loader.
Args:
model: PyTorch model to train or evaluate.
loader: DataLoader for a split, or ``None`` for an empty split.
device: Torch device used for tensors and model execution.
loss_fn: Loss module used to compare predictions and targets.
optimizer: Optional optimizer. When provided, gradients are updated.
gradient_clip_norm: Optional positive norm for gradient clipping during
training.
target_metrics: Whether to compute target-based metrics such as
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:
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
loss_denominator = 0.0
preds: list[np.ndarray] = []
targets: list[np.ndarray] = []
reporter = ProgressReporter(
progress_label or ("train batches" if training else "eval batches"),
total=len(loader),
unit="batches",
enabled=progress,
percent_step=25.0,
)
for xb, yb, auxb in loader:
xb = xb.to(device)
yb = yb.to(device).float().reshape(-1)
auxb = {name: values.to(device).float().reshape(-1) for name, values in auxb.items()}
if training:
optimizer.zero_grad(set_to_none=True)
with torch.set_grad_enabled(training):
yhat = squeeze_prediction(model(xb)).reshape(-1)
loss_output = loss_fn(yhat, yb, auxb)
loss = loss_output.loss
if training:
loss.backward()
if gradient_clip_norm is not None and float(gradient_clip_norm) > 0:
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=float(gradient_clip_norm))
optimizer.step()
loss_numerator += float(loss_output.numerator.detach().cpu().item())
loss_denominator += float(loss_output.denominator.detach().cpu().item())
preds.append(yhat.detach().cpu().numpy())
if target_metrics:
targets.append(yb.detach().cpu().numpy())
reporter.update()
reporter.close()
y_pred = np.concatenate(preds) if preds else np.array([])
y_true = np.concatenate(targets) if targets else np.array([])
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:
"""Return whether a monitored metric improved.
Args:
value: Current metric value.
best: Previous best metric value, or ``None`` if unset.
monitor: Metric name used to decide whether lower or higher is better.
"""
if np.isnan(value):
return False
if best is None:
return True
if monitor.endswith("loss") or monitor in {"loss", "mse", "val_mse"}:
return value < best
return value > best
def _monitor_names(monitor: str | Sequence[str]) -> tuple[str, ...]:
"""Normalize one or more monitored metric names.
Args:
monitor: Comma-separated metric string or sequence of metric names.
"""
if isinstance(monitor, str):
names = [part.strip() for part in monitor.split(",") if part.strip()]
else:
names = [str(part).strip() for part in monitor if str(part).strip()]
if not names:
raise ValueError("monitor must name at least one metric")
return tuple(names)
def _format_best_metrics(best_metrics: Mapping[str, float | None]) -> str:
"""Format monitored best values for progress output.
Args:
best_metrics: Mapping from monitor names to their current best values.
"""
return ", ".join(f"best_{name}={value}" for name, value in best_metrics.items())
def _monitor_improved(
row: Mapping[str, float | int],
monitors: Sequence[str],
best_metrics: Mapping[str, float | None],
) -> tuple[bool, dict[str, float]]:
"""Return whether any monitored metric improved over the current best epoch.
Args:
row: Current epoch metrics keyed by metric name.
monitors: Metric names to compare against ``best_metrics``.
best_metrics: Previous best values for each monitored metric.
"""
missing_monitors = [name for name in monitors if name not in row]
if missing_monitors:
raise ValueError(f"Unknown monitor metric(s): {', '.join(missing_monitors)}")
values = {name: float(row[name]) for name in monitors}
improved = any(_is_better(values[name], best_metrics[name], name) for name in monitors)
return improved, values
[docs]
def train_model(bundle: DatasetBundle, config: TrainConfig | Mapping[str, Any]) -> dict[str, Any]:
"""Train a model from an in-memory dataset bundle and config.
Args:
bundle: Dataset bundle containing encoded inputs and regression targets.
config: Training configuration object or mapping of config fields.
"""
cfg = _as_train_config(config)
requested_model = normalize_model_config(cfg.model)
if requested_model.name == "rbpnet":
if cfg.sequence_controls:
raise ValueError(
"generic sequence_controls are not supported for structured RBPNet training"
)
from transcriptml.rbpnet.training import train_rbpnet_model
return train_rbpnet_model(bundle, cfg)
_seed_everything(cfg.seed)
device = resolve_device(cfg.device)
out = Path(cfg.output_dir)
out.mkdir(parents=True, exist_ok=True)
sequence_control_stats: dict[str, Any] | None = None
if cfg.sequence_controls:
bundle, sequence_control_stats = apply_sequence_controls_to_bundle(
bundle,
cfg.sequence_controls,
default_save_dir=out / "sequence_controlled_dataset",
progress=cfg.progress,
)
loss_fn, aux_arrays, normalized_loss_config = build_training_loss(
cfg.loss,
metadata=bundle.metadata,
n_examples=int(bundle.X.shape[0]),
)
has_targets = bundle.y is not None
if bundle.y is None:
if loss_fn.requires_target:
raise ValueError(f"Training loss '{loss_fn.name}' requires bundle.y")
y_train = np.zeros(int(bundle.X.shape[0]), dtype=np.float32)
else:
y_train = bundle.y
splits, split_source_used = _select_splits(bundle, cfg)
cv_plan_id = _cv_plan_id(cfg)
split_counts = {name: len(splits.get(name, [])) for name in ("train", "val", "test")}
model_config = normalize_model_config(cfg.model)
if cfg.head_layernorm and model_config.name != "saluki_exact":
raise ValueError("head_layernorm is only supported for model 'saluki_exact'")
if model_config.name == "saluki_exact":
model_config.params = dict(model_config.params or {})
model_config.params["head_layernorm"] = bool(cfg.head_layernorm)
log_progress(
(
"training: "
f"device={device}, output={out}, loss={normalized_loss_config['name']}, "
f"head_layernorm={cfg.head_layernorm}, "
f"split_source={split_source_used} requested={cfg.split_source}, "
f"train={split_counts['train']}, val={split_counts['val']}, "
f"test={split_counts['test']}"
),
enabled=cfg.progress,
)
model = build_model(model_config).to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
dataset = _ArrayRegressionDataset(bundle.X, y_train, aux_arrays)
pin_memory = device.type == "cuda"
drop_last_train = _has_batch_normalization(model) and _would_create_singleton_batch(
len(splits["train"]), cfg.batch_size
)
if drop_last_train:
log_progress(
(
"training: train split would end with a singleton batch; "
"dropping the final shuffled training example each epoch"
),
enabled=cfg.progress,
)
train_loader = _loader(
dataset,
splits["train"],
cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers,
pin_memory=pin_memory,
drop_last=drop_last_train,
)
val_loader = _loader(
dataset,
splits["val"],
cfg.batch_size,
shuffle=False,
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}
best_epoch = -1
stale = 0
for epoch in range(1, cfg.epochs + 1):
log_progress(f"epoch {epoch}/{cfg.epochs}: starting", enabled=cfg.progress)
train_metrics = _run_loader(
model,
train_loader,
device=device,
loss_fn=loss_fn,
optimizer=optimizer,
gradient_clip_norm=cfg.gradient_clip_norm,
target_metrics=has_targets,
progress=cfg.progress,
progress_label=f"epoch {epoch} train",
)
val_metrics = _run_loader(
model,
val_loader,
device=device,
loss_fn=loss_fn,
optimizer=None,
target_metrics=has_targets,
progress=cfg.progress,
progress_label=f"epoch {epoch} val",
return_predictions=cfg.debug_epoch_predictions,
)
row = {
"epoch": epoch,
"train_loss": train_metrics["loss"],
"train_pearson": train_metrics["pearson"],
"val_loss": val_metrics["loss"],
"val_pearson": val_metrics["pearson"],
}
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
stale = 0
save_checkpoint(
out / "best.pt",
model,
model_config,
epoch=epoch,
metrics=row,
optimizer_state=optimizer.state_dict(),
extra={
"splits": splits,
"split_source_used": split_source_used,
"cv_plan_id": cv_plan_id,
"train_config": asdict(cfg),
"loss_config": normalized_loss_config,
},
)
else:
stale += 1
save_checkpoint(
out / "last.pt",
model,
model_config,
epoch=epoch,
metrics=row,
optimizer_state=optimizer.state_dict(),
extra={
"splits": splits,
"split_source_used": split_source_used,
"cv_plan_id": cv_plan_id,
"train_config": asdict(cfg),
"loss_config": normalized_loss_config,
},
)
log_progress(
(
f"epoch {epoch}/{cfg.epochs}: "
f"train_loss={row['train_loss']:.6g}, train_pearson={row['train_pearson']:.4g}, "
f"val_loss={row['val_loss']:.6g}, val_pearson={row['val_pearson']:.4g}, "
f"{_format_best_metrics(best_metrics)}"
),
enabled=cfg.progress,
)
if cfg.patience >= 0 and stale > cfg.patience:
log_progress(f"early stopping after epoch {epoch}: patience exceeded", enabled=cfg.progress)
break
(out / "history.json").write_text(json.dumps(history, indent=2), encoding="utf-8")
(out / "splits.json").write_text(json.dumps(splits, indent=2), encoding="utf-8")
# Well there's a chance there is none, so kind of weird...
# Not that a user should ever not want to include a test split
log_progress("training: evaluating test split", enabled=cfg.progress)
test_loader = _loader(
dataset,
splits.get("test", []),
cfg.batch_size,
shuffle=False,
num_workers=cfg.num_workers,
pin_memory=pin_memory,
)
test_loss_metrics = _run_loader(
model,
test_loader,
device=device,
loss_fn=loss_fn,
optimizer=None,
target_metrics=has_targets,
progress=cfg.progress,
progress_label="test loss",
)
test_result = evaluate_model(
model,
bundle,
indices=splits.get("test", []),
batch_size=cfg.batch_size,
device=device,
progress=cfg.progress,
)
if splits.get("test"):
ids = [bundle.ids[int(i)] for i in test_result["indices"]]
predict_to_csv(
out / "test_predictions.csv",
ids=ids,
predictions=test_result["predictions"],
targets=test_result.get("targets"),
indices=test_result["indices"],
)
summary = {
"best_epoch": best_epoch,
"monitor": list(monitors),
"best_monitor_values": best_metrics,
"epochs_run": len(history),
"loss": normalized_loss_config,
"head_layernorm": bool(cfg.head_layernorm),
"split_source_requested": cfg.split_source,
"split_source_used": split_source_used,
"cv_plan": cfg.cv_plan,
"cv_plan_id": cv_plan_id,
"fold": cfg.fold,
"split_counts": split_counts,
"test_loss": test_loss_metrics.get("loss"),
"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")
log_progress(
f"training: done; best_epoch={best_epoch}, summary={out / 'summary.json'}",
enabled=cfg.progress,
)
# Not sure this return value ever gets used anywhere
return {"model": model, "history": history, "splits": splits, "summary": summary}
[docs]
def train_from_config(
config_path: str | Path,
*,
progress: bool | None = None,
cv_plan: str | Path | None = None,
fold: int | None = None,
dataset: str | Path | None = None,
output_dir: str | Path | None = None,
) -> dict[str, Any]:
"""Load a training config and train its requested model.
Args:
config_path: Path to a JSON or TOML training configuration file.
progress: Optional override for whether progress messages are emitted.
cv_plan: Optional chromosome CV plan overriding the config.
fold: Optional zero-based CV test fold overriding the config.
dataset: Optional dataset-directory override.
output_dir: Optional output-directory override.
"""
cfg = TrainConfig(**_load_config(config_path))
if cv_plan is not None:
cfg.cv_plan = str(cv_plan)
if fold is not None:
cfg.fold = int(fold)
if dataset is not None:
cfg.dataset = str(dataset)
if output_dir is not None:
cfg.output_dir = str(output_dir)
if progress is not None:
cfg.progress = bool(progress)
log_progress(f"training: loading dataset {cfg.dataset}", enabled=cfg.progress)
bundle = load_bundle(cfg.dataset, mmap_mode=cfg.mmap_mode)
return train_model(bundle, cfg)