Trainer: subsample train coverage metrics (train_metrics_every, default 10)

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` 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:
 
 - `r/<group>`, `r2/<group>`: mean per-track Pearson r and R² over the group's tracks.

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)
+        # 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")
         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 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(
@@ -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 = 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