diff --git a/docs/train.md b/docs/train.md index 69eeba6..44bfc31 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` 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/`, `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..a6e104d 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) + # 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