Trainer rewrite: AMP, convergence-aware stopping, Runner.save(), clean adata separation

[jackytamkc]

Depends on #29. Until #29 merges, the diff shown here includes its two-line num_nodes() patch (commit 4c811f9). Once #29 merges, that commit becomes a no-op and only the three trainer commits remain.

Motivation

Two problems with the existing trainer:

1. No convergence signal, no way to stop early. The loop runs a fixed number of epochs with no record of what happened inside. Because the total loss is dominated by a BCE floor on sparse adjacency reconstruction, loss curves plateau long before the embedding stops changing — and small late-epoch loss improvements don't necessarily correspond to any change in the downstream clustering. Users end up either training far too long or stopping too early based on a signal that doesn't reflect how the output is actually used.
2. Training scaffolding bleeds into the data object. prepare_data{,_batch} writes DGL graphs, the GraphDataLoader, and batch_idx into adata.uns, and fit() then writes loss / lr_history / early_stop there too. The result: adata.copy() pulls in multiprocess dataloaders, write_h5ad() carries scaffolding it can't serialise cleanly, and the "trained AnnData" is not actually downstream-ready.

This PR addresses both: stop on a signal that matches how the embedding is used, persist enough training metadata to reconstruct what happened, and cleanly separate the model (state dicts + diagnostics, saved as model.pt) from the data (embedding only, saved as adata.h5ad).

What changes

Commit 2 — Trainer rewrite (9555191)

• Mixed precision: torch.autocast(cuda, bf16) around the GCN forward and adjacency-reconstruction losses. The MIM / discriminator block stays fp32 — log(x + 1e-6) near 1.0 is unsafe in bf16.
• Loss-based early stopping with patience, min_delta, and optional restore_best (rolls weights back to the best-loss epoch before producing the final embedding).
• ReduceLROnPlateau scheduler (lr_patience, lr_factor, min_lr) with verbose logging when the LR drops.
• ARI-based convergence probe (Runner_batch only): every ari_every epochs, snapshot the embedding, KMeans-cluster it (ari_k), and compute ARI vs. the previous snapshot. Stop when ari_stop_patience consecutive probes stay above ari_stop_threshold. This is the signal that matches downstream use, not the raw loss.
• Batch caching: move every batch to device once up front and reuse the same graph/feat/adj tensors across epochs — eliminates the per-epoch H2D transfer.
• _infer_embedding helper shared by the ARI probe and the final embedding pass.
• _model.py / _build.py: MGAE.forward builds its fused graph GPU-natively (no scipy.csr roundtrip); the batch GraphDataLoader gets num_workers=4 for the one-time upfront caching pass.

Commit 3 — Keep model scaffolding off adata (7006515)

• Runner.__init__ pops g_<view> into self.graph.
• Runner_batch.__init__ pops dataloader into self.dataloader and batch_idx into self.batch_idx.
• After Runner construction, adata.uns is clean of model plumbing — adata.copy() / write_h5ad() work without needing workarounds like adata.uns.pop('dataloader', None).
• New Runner.save(output_dir) / Runner_batch.save(output_dir) writes:
  • {output_dir}/adata.h5ad — results-only AnnData
  • {output_dir}/model.pt — {model_state, discriminator_state, config} with enough config to rebuild MGAE + Discriminator for later inference.

Commit 4 — Move training diagnostics off adata too (ddfcf09)

• loss_history, lr_history, ari_history, early_stop now live on the Runner (self.*) instead of adata.uns.
• save() bundles them into model.pt under a training key, alongside the state dicts.
• Result: the AnnData returned from fit() contains only obsm['X_scniche'] as a training-related addition — nothing in adata.uns is there because of training bookkeeping. It is genuinely downstream-ready (niche clustering, enrichment plots, etc.) without any scaffolding or diagnostics noise.

Backward compatibility

Defaults reproduce the old training behaviour as closely as possible:

• patience=None, ari_every=None, ari_stop_threshold=None → early stopping and ARI probe are off by default.
• lr_patience=5, lr_factor=0.5 → an LR schedule is on by default. Pass lr_patience=None to hold the LR fixed.
• use_amp=True by default. Pass use_amp=False for the original fp32 path.
• restore_best=True — weights revert to best-loss epoch before the final embedding pass.

Breaking change (intentional): training diagnostics are no longer written to adata.uns. Code that reads adata.uns['loss'] / adata.uns['lr_history'] / adata.uns['ari_history'] / adata.uns['early_stop'] needs to read from model.loss_history / model.lr_history / model.ari_history / model.early_stop instead, or from the corresponding keys inside model.pt after Runner.save(). This is the point of the separation — mixing them in previously was the bug.

Test plan

• Runner.fit() with all defaults on a small dataset: finishes in epochs, adata.obsm['X_scniche'] sane, adata.uns untouched by training.
• Runner.fit(use_amp=False, lr_patience=None, restore_best=False) produces the same embedding as before this PR on a fixed seed.
• Runner_batch.fit(ari_every=10, ari_stop_threshold=0.98, ari_stop_patience=2) on a large dataset: stops at the first pair of probes hitting ARI ≥ 0.98; model.early_stop['stopped_by'] == 'ari'.
• Runner_batch.fit(patience=10, min_delta=1e-4): stops on loss plateau; model.early_stop['stopped_by'] == 'loss'.
• Runner.save('out/') / Runner_batch.save('out/') writes adata.h5ad + model.pt; model.pt round-trips back into a rebuildable MGAE + Discriminator, and the training block contains loss/lr/ari/early_stop.
• adata.copy() works immediately after Runner_batch(...) construction (was previously broken by the multiprocess dataloader sitting in adata.uns).

🤖 Generated with Claude Code