From ae755f0f88ffcaa74a884c7fd49589d44f9d7dc3 Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 15 Oct 2026 21:04:41 -0700 Subject: [PATCH 1/2] Trainer: train r/r2/spec on every Nth batch per dataset (train_metrics_every, default 10) Train loss still covers every batch. At 8361 tracks the per-step metric update costs ~9% of an A100 training step; every 10th batch cuts it to <1%. Co-Authored-By: Claude Opus 5.5 --- docs/train.md | 2 ++ src/baskerville/trainer.py | 14 +++++++++++--- 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/docs/train.md b/docs/train.md index 69eeba6..626852b 100644 --- a/docs/train.md +++ b/docs/train.md @@ -123,6 +123,8 @@ The loss is a weighted mean over tracks, so only the ratios matter: scaling ever Each epoch, `log.txt` gets a summary line per dataset (`loss`, `r`, `r2`, `spec`), and `metrics.tsv` gets every metric as a long table with columns `epoch, split, dataset, group, metric, value`. Dataset-level rows have `group` = `all`. +Train `loss` covers every step, but train `r`, `r2` and `spec` are computed on every `train.train_metrics_every`th batch of each dataset (default 10; 1 scores every batch). At thousands of tracks, scoring every step costs ~9% of training time. + Tracks are grouped by the targets `group` column (else the `description` prefix). Each group with at least `train.spec_group_min` tracks (default 20, after merging strand pairs) gets: - `r/`, `r2/`: mean per-track Pearson r and R² over the group's tracks. diff --git a/src/baskerville/trainer.py b/src/baskerville/trainer.py index dd8b4ad..0b3d42e 100644 --- a/src/baskerville/trainer.py +++ b/src/baskerville/trainer.py @@ -284,6 +284,8 @@ def __init__( self.stop_stat = parse_stop_stat(stop_stat) self.stop_stat_loss_fallback = isinstance(stop_stat, str) self.spec_group_min = self.params.get("spec_group_min", 20) + # train r, r2, spec on every Nth batch per dataset; train loss on every batch + self.train_metrics_every = self.params.get("train_metrics_every", 10) # MLM parameters (only loaded when loss == "mlm") if self.loss == "mlm": @@ -1160,8 +1162,8 @@ def _setup_optimization(self): # graph on some torch versions. Falls back to self.model when not compiled. self.model_eval = getattr(self.model, "_orig_mod", self.model) - def _train_batch_cov(self, di_tensor, batch_data): - """Coverage-supervised training step.""" + def _train_batch_cov(self, di_tensor, batch_data, score=True): + """Coverage-supervised training step; score=False updates only the loss.""" di = di_tensor[0].item() x, y, yg, gene_presence, gene_out_mask = self._unpack_batch(batch_data) sw = torch.tensor(self.species_weights[di]).to( @@ -1179,6 +1181,8 @@ def _train_batch_cov(self, di_tensor, batch_data): ) loss = self._compute_loss(yh, y, yg, gene_presence, sw, di) + if not score: + y = yg = None self.train_metrics[di].update(yh, y, yg, gene_presence, loss, x.shape[0]) return loss @@ -1207,6 +1211,7 @@ def _train_batch_mlm(self, batch): def _train_epoch(self, ei): """Run one training epoch over ``train_dataload``.""" train_batch_i = 0 + dataset_batches = [0] * self.num_datasets t0 = time.time() last_progress = t0 for di_tensor, batch_data in tqdm( @@ -1219,7 +1224,10 @@ def _train_epoch(self, ei): if self.is_mlm: loss = self._train_batch_mlm(batch_data) else: - loss = self._train_batch_cov(di_tensor, batch_data) + di = di_tensor[0].item() + score = dataset_batches[di] % self.train_metrics_every == 0 + dataset_batches[di] += 1 + loss = self._train_batch_cov(di_tensor, batch_data, score) self._optimizer_step(loss) train_batch_i += 1 From d3218bf0010e1b887bd92d218b56bb9ea0f2044e Mon Sep 17 00:00:00 2001 From: David Kelley Date: Thu, 15 Oct 2026 21:15:21 -0700 Subject: [PATCH 2/2] Trainer: keep gene metrics on every batch under train_metrics_every Sampled batches can hold fewer than two valid genes, and gene R2 raises on fewer than two observations at epoch end. Gene metrics have no length axis, so they cost little; only coverage r/r2/spec are subsampled. Co-Authored-By: Claude Opus 5.5 --- docs/train.md | 2 +- src/baskerville/trainer.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/train.md b/docs/train.md index 626852b..44bfc31 100644 --- a/docs/train.md +++ b/docs/train.md @@ -123,7 +123,7 @@ The loss is a weighted mean over tracks, so only the ratios matter: scaling ever Each epoch, `log.txt` gets a summary line per dataset (`loss`, `r`, `r2`, `spec`), and `metrics.tsv` gets every metric as a long table with columns `epoch, split, dataset, group, metric, value`. Dataset-level rows have `group` = `all`. -Train `loss` covers every step, but train `r`, `r2` and `spec` are computed on every `train.train_metrics_every`th batch of each dataset (default 10; 1 scores every batch). At thousands of tracks, scoring every step costs ~9% of training time. +Train `loss` and gene metrics cover every batch, but train coverage `r`, `r2` and `spec` are computed on every `train.train_metrics_every`th batch of each dataset (default 10; 1 scores every batch). At thousands of tracks, scoring every batch costs ~9% of training time. Tracks are grouped by the targets `group` column (else the `description` prefix). Each group with at least `train.spec_group_min` tracks (default 20, after merging strand pairs) gets: diff --git a/src/baskerville/trainer.py b/src/baskerville/trainer.py index 0b3d42e..a6e104d 100644 --- a/src/baskerville/trainer.py +++ b/src/baskerville/trainer.py @@ -284,7 +284,7 @@ def __init__( self.stop_stat = parse_stop_stat(stop_stat) self.stop_stat_loss_fallback = isinstance(stop_stat, str) self.spec_group_min = self.params.get("spec_group_min", 20) - # train r, r2, spec on every Nth batch per dataset; train loss on every batch + # coverage r, r2, spec on every Nth batch per dataset; loss, gene metrics always self.train_metrics_every = self.params.get("train_metrics_every", 10) # MLM parameters (only loaded when loss == "mlm") @@ -1163,7 +1163,7 @@ def _setup_optimization(self): self.model_eval = getattr(self.model, "_orig_mod", self.model) def _train_batch_cov(self, di_tensor, batch_data, score=True): - """Coverage-supervised training step; score=False updates only the loss.""" + """Coverage-supervised training step; score=False skips coverage metrics.""" di = di_tensor[0].item() x, y, yg, gene_presence, gene_out_mask = self._unpack_batch(batch_data) sw = torch.tensor(self.species_weights[di]).to( @@ -1182,7 +1182,7 @@ def _train_batch_cov(self, di_tensor, batch_data, score=True): loss = self._compute_loss(yh, y, yg, gene_presence, sw, di) if not score: - y = yg = None + y = None self.train_metrics[di].update(yh, y, yg, gene_presence, loss, x.shape[0]) return loss