Make training reproducible: seeded per-loader shuffles, no global RNG side effects

docs/api/datasets.rst

@@ -230,6 +230,20 @@ Each event in the returned list has:
 - ``event.icd_code``, ``event["icd_code"]``, ``event.attr_dict`` — different
   ways to access the other attributes. All attribute names are lowercase.
 
+Shuffling and Reproducibility
+-----------------------------
+
+``get_dataloader(dataset, batch_size, shuffle=True, seed=None)`` gives each
+loader its own shuffle order: the same ``seed`` gives the same sequence of
+epoch orders, a different seed a different one, and each epoch is reshuffled.
+With ``seed=None`` the seed is drawn from torch's random generator when the
+loader is created, so ``torch.manual_seed(...)`` alone makes runs
+reproducible. This also holds with ``num_workers > 0``. Each loader applies
+its own settings to the dataset at the start of every pass, so an unshuffled
+evaluation loader on a training dataset leaves the training loader shuffling;
+just don't iterate two loaders over one dataset at the same time. See the
+Trainer docs for the full checklist.
+
 Things to Watch Out For
 ------------------------
 

docs/api/models/pyhealth.models.GNN.rst

@@ -3,6 +3,13 @@
 
 The GNN model (pyhealth trainer does not apply to GNN, refer to the example/ChestXray-image-generation-GAN.ipynb for examples of using GNN model).
 
+.. note::
+   ``pyhealth.models.gnn`` no longer seeds the global ``torch``/``numpy``
+   random generators at import time. Previously ``torch.manual_seed(3)`` and
+   ``np.random.seed(1)`` ran when ``pyhealth.models`` was imported,
+   overwriting any seed set before the import. For reproducible GCN/GAT
+   weights, seed right before constructing the model.
+
 .. autoclass:: pyhealth.models.GAT
     :members:
     :undoc-members:

docs/api/trainer.rst

@@ -93,6 +93,48 @@ object. To capture results for downstream use, call ``evaluate()`` separately:
     with open("results.json", "w") as f:
         json.dump(scores, f, indent=2)
 
+Reproducible Runs
+-----------------
+
+On CPU, one seed makes a training run bit-for-bit reproducible, in the same
+process or a new one:
+
+.. code-block:: python
+
+    import torch
+    from pyhealth.datasets import get_dataloader
+
+    torch.manual_seed(42)        # or pyhealth.utils.set_seed(42)
+    train_loader = get_dataloader(train_ds, batch_size=64, shuffle=True)
+    model = RNN(dataset=train_ds)
+    Trainer(model=model).train(train_loader, val_loader, epochs=10, monitor="pr_auc")
+
+How the pieces fit:
+
+- **Shuffle order** belongs to each loader. ``get_dataloader(..., shuffle=True)``
+  draws its seed from torch's generator when it is created (or takes
+  ``seed=...``), and each epoch gets a new order derived from that seed and
+  the epoch number. Shuffling never touches Python's, NumPy's or torch's
+  global random state, and works the same for in-memory and disk-backed
+  datasets.
+- **Model initialization** uses torch's generator, as usual.
+- **Cache writes** (``set_task``, ``create_sample_dataset(in_memory=False)``)
+  leave the global random state as it was, so a run behaves the same whether
+  or not the cache already existed. PyHealth never seeds the global
+  generators itself; ``pyhealth.utils.set_seed`` is the only place, and only
+  when you call it.
+- **The best epoch** is restored at the end of ``train`` when ``monitor`` is
+  set, with or without logging. Without an output directory
+  (``enable_logging=False``) the best weights are kept as a CPU copy in
+  memory; pass ``keep_best_in_memory=False`` to skip that copy for very large
+  models (the last epoch's weights are then kept, with a warning).
+
+On GPU, some kernels are non-deterministic. For bit-identical GPU runs also
+call ``torch.use_deterministic_algorithms(True)``, set
+``torch.backends.cudnn.benchmark = False`` and the environment variable
+``CUBLAS_WORKSPACE_CONFIG=:4096:8``. These are process-wide settings that can
+slow training, so PyHealth leaves them to you.
+
 API Reference
 -------------
 

examples/reproducible_training.py

@@ -0,0 +1,54 @@
+"""Reproducible training: one seed, identical predictions.
+
+Trains the same model twice with the same seed and once with another, on an
+in-memory and a disk-backed dataset, and compares the predictions. With
+``torch.manual_seed`` set before the loaders and the model are created, the
+shuffle order, the initial weights and therefore the predictions are
+bit-identical; a different seed gives different ones.
+
+Runs in under a minute on CPU with synthetic data; no download.
+
+Usage:
+    python examples/reproducible_training.py
+"""
+
+import hashlib
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset, get_dataloader
+from pyhealth.models import RNN
+from pyhealth.trainer import Trainer
+
+samples = [
+    {
+        "patient_id": f"p{i}",
+        "codes": [f"c{(i + j) % 9}" for j in range(1 + i % 4)],
+        "label": int(i % 3 == 0),
+    }
+    for i in range(60)
+]
+
+
+def run(seed: int, in_memory: bool) -> str:
+    torch.manual_seed(seed)
+    dataset = create_sample_dataset(
+        samples, {"codes": "sequence"}, {"label": "binary"}, in_memory=in_memory
+    )
+    train_loader = get_dataloader(dataset, batch_size=8, shuffle=True)
+    trainer = Trainer(model=RNN(dataset=dataset), enable_logging=False)
+    trainer.train(train_loader, get_dataloader(dataset, 16), epochs=3, monitor="roc_auc")
+    _, y_prob, _ = trainer.inference(get_dataloader(dataset, 16))
+    return hashlib.sha256(y_prob.tobytes()).hexdigest()[:12]
+
+
+def main():
+    for in_memory in (True, False):
+        kind = "in-memory" if in_memory else "disk-backed"
+        a, b, c = run(0, in_memory), run(0, in_memory), run(1, in_memory)
+        print(f"{kind:12s} seed 0: {a}  seed 0 again: {b}  seed 1: {c}")
+        print(f"{'':12s} same seed identical: {a == b}, different seed differs: {a != c}")
+
+
+if __name__ == "__main__":
+    main()

pyhealth/calib/predictionset/scrib/__init__.py

@@ -15,6 +15,7 @@
 from pyhealth.calib.base_classes import SetPredictor
 from pyhealth.calib.utils import prepare_numpy_dataset
 from pyhealth.models import BaseModel
+from pyhealth.utils import preserve_rng_state
 
 from . import quicksearch as qs
 
@@ -108,6 +109,7 @@ def _sample_new_loc(self, old_p, restart_range=0.1):
             new_p[k] = max(min(int(new_p[k] + diff[k] * self.N), self.N - 1), 0)
         return new_p
 
+    @preserve_rng_state()
     def search_once(self, seed=7):
         def print_(s):
             if self.verbose:
@@ -163,7 +165,6 @@ def search(
         best_loss, best_ts = np.inf, None
         searcher = cls(prob, label, rks, loss_func=loss_func, **kwargs)
         for seed in range(B):
-            np.random.seed(seed)
             ts, _l = searcher.search_once(seed + 1)
             print(f"{seed}: loss={_l}")
             if _l < best_loss:

pyhealth/datasets/base_dataset.py

@@ -45,7 +45,7 @@
 from .configs import load_yaml_config
 from .sample_dataset import STORAGE_FORMAT, SampleDataset, SampleBuilder
 from .splitter import Split
-from ..utils import download_file, set_env
+from ..utils import download_file, preserve_rng_state, set_env
 
 # Set logging level for distributed to ERROR to reduce verbosity
 logging.getLogger("distributed").setLevel(logging.ERROR)
@@ -1123,6 +1123,7 @@ def _proc_transform(
         finally:
             self.clean_tmpdir()
 
+    @preserve_rng_state()
     def set_task(
         self,
         task: Optional[BaseTask] = None,

pyhealth/datasets/sample_dataset.py

@@ -6,9 +6,9 @@
 from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type
 from typing import TYPE_CHECKING
 import inspect
-import random
 from bisect import bisect_right
 import litdata
+from ..utils import preserve_rng_state
 from litdata.utilities.train_test_split import deepcopy_dataset
 import copy
 import logging
@@ -48,6 +48,15 @@ def _decode_record(item: Any) -> Any:
     return {k: values[k] if k in values else item[k] for k in payload["order"]}
 
 
+def _draw_seed() -> int:
+    """A shuffle seed drawn from torch's global generator.
+
+    So ``torch.manual_seed(...)``, which users already call, makes shuffling
+    reproducible, and different torch seeds give different orders.
+    """
+    return int(torch.randint(0, 2**31 - 1, ()).item())
+
+
 class _MaskedSamples:
     """A re-iterable view of a sample stream that keeps only masked positions.
 
@@ -528,6 +537,28 @@ def __getitem__(self, index: Any) -> Any:
         """
         return _decode_record(super().__getitem__(index))
 
+    def set_shuffle(self, shuffle: bool, seed: int | None = None) -> None:
+        """Sets whether iterating the dataset shuffles, and the shuffle seed.
+
+        The order of each pass is determined by ``seed`` and the epoch (see
+        ``set_epoch``); litdata shuffles with its own generator, so the global
+        random state is not used. Loaders from
+        :func:`~pyhealth.datasets.get_dataloader` call this for you at the
+        start of every pass.
+
+        Args:
+            shuffle: Whether ``iter(dataset)`` shuffles.
+            seed: Seed of the shuffle order. Default: drawn from torch's
+                random generator, so ``torch.manual_seed`` controls it.
+
+        Examples:
+            >>> dataset.set_shuffle(True, seed=0)  # doctest: +SKIP
+        """
+        super().set_shuffle(shuffle)
+        if shuffle:
+            # litdata rebuilds its shuffler from ``seed`` on every iteration.
+            self.seed = _draw_seed() if seed is None else int(seed)
+
     def _remove_ignored_processors(self):
         """Remove any processors that are IgnoreProcessor instances."""
         for key in [
@@ -718,9 +749,44 @@ def __init__(
         self._data = [builder.transform({"sample": pickle.dumps(s)}) for s in samples]
 
         self._shuffle = False
+        self._shuffle_seed = 0
+        self._epoch = 1  # counted from 1, as litdata does
+
+    def set_shuffle(self, shuffle: bool, seed: int | None = None) -> None:
+        """Sets whether iterating the dataset shuffles, and the shuffle seed.
 
-    def set_shuffle(self, shuffle: bool) -> None:
+        The order of each pass is determined by ``seed`` and the epoch (see
+        ``set_epoch``), using a private generator, so the global random state
+        is not used. Loaders from :func:`~pyhealth.datasets.get_dataloader`
+        call this for you at the start of every pass.
+
+        Args:
+            shuffle: Whether ``iter(dataset)`` shuffles.
+            seed: Seed of the shuffle order. Default: drawn from torch's
+                random generator, so ``torch.manual_seed`` controls it.
+
+        Examples:
+            >>> dataset.set_shuffle(True, seed=0)  # doctest: +SKIP
+        """
         self._shuffle = shuffle
+        if shuffle:
+            self._shuffle_seed = _draw_seed() if seed is None else int(seed)
+        self._epoch = 1
+
+    def set_epoch(self, epoch: int) -> None:
+        """Sets the epoch whose shuffle order the next pass uses.
+
+        Iterating advances the epoch by one, so direct iteration gives a new
+        order each pass. With worker processes, iteration happens on copies of
+        the dataset, so the loader sets the epoch in the main process instead.
+
+        Args:
+            epoch: The epoch number (counted from 1, as litdata does).
+
+        Examples:
+            >>> dataset.set_epoch(2)  # doctest: +SKIP
+        """
+        self._epoch = int(epoch)
 
     def __len__(self) -> int:
         """Returns the number of samples in the dataset.
@@ -748,11 +814,19 @@ def __iter__(self) -> Iterable[Dict[str, Any]]:  # type: ignore
             An iterator yielding processed sample dictionaries.
         """
         if self._shuffle:
-            shuffled_data = self._data[:]
-            random.shuffle(shuffled_data)
-            return iter(shuffled_data)
+            rng = np.random.default_rng([self._shuffle_seed, self._epoch])
+            order = rng.permutation(len(self._data))
         else:
-            return iter(self._data)
+            order = np.arange(len(self._data))
+        self._epoch += 1
+        # With DataLoader worker processes, each worker iterates its own copy
+        # of the dataset; give each a disjoint share so no sample repeats.
+        worker = torch.utils.data.get_worker_info()
+        if worker is not None:
+            order = order[worker.id :: worker.num_workers]
+        # The order is fixed here, not lazily, so later changes to the
+        # dataset's settings don't affect a pass already started.
+        return (self._data[i] for i in order)
 
     def subset(self, indices: Union[Sequence[int], slice]) -> SampleDataset:
         if isinstance(indices, slice):
@@ -800,6 +874,7 @@ def _write_samples(builder: "SampleBuilder", samples: list[dict], path: Path) ->
     _litdata_merge(path)
 
 
+@preserve_rng_state()
 def create_sample_dataset(
     samples: List[Dict[str, Any]],
     input_schema: Dict[str, Any],