From f8e4c07e42e5bf588cfd2e36da09ff7af393f46f Mon Sep 17 00:00:00 2001 From: solarsys Date: Fri, 9 Oct 2026 02:19:30 -0400 Subject: [PATCH 1/2] Make training reproducible: seeded per-loader shuffles, no global RNG side effects A fixed torch/numpy seed did not make PyHealth training reproducible: - InMemorySampleDataset shuffled with Python's global `random`, which torch.manual_seed and np.random.seed do not touch. - litdata's DataProcessor.run (used by create_sample_dataset(in_memory=False)) reseeds random/numpy/torch to 42 in the calling process, and set_task's cache write consumed Python's `random`; runs differed by whether the cache already existed. - The disk-backed shuffle ignored the user's seed. - get_dataloader changed the shared dataset's shuffle flag, so an evaluation loader turned off shuffling for a training loader. - With enable_logging=False the best epoch was never restored. - Importing pyhealth.models reseeded numpy and torch (gnn.py). Changes: - get_dataloader(..., seed=None): each loader owns its shuffle flag, seed and epoch counter (a view; the dataset is not modified). Each pass is a new deterministic order from (seed, epoch): NumPy default_rng in memory, litdata's own (seed, epoch) shuffle on disk via a shallow copy. With seed=None the seed is drawn from torch's generator at creation, so torch.manual_seed alone makes runs reproducible. A private torch.Generator keeps DataLoader's per-pass base-seed draw off the global state. set_shuffle(shuffle, seed=None) on datasets is seeded the same way. - pyhealth.utils.preserve_rng_state() saves/restores random, numpy, torch and CUDA states; set_task and create_sample_dataset are wrapped in it. - Trainer.train keeps a CPU copy of the best weights when there is no output directory and restores them (keep_best_in_memory=False opts out, with a warning). - No library code seeds global RNGs any more: removed gnn.py's import-time seeding (from #1227); interpret SHAP/LIME/RandomBaseline, SCRIB search, train_lstm_model and the KG splitter now seed only for the call. set_seed documents GPU flags and that PYTHONHASHSEED only affects subprocesses. - Docs: "Reproducible Runs" (trainer), "Shuffling and Reproducibility" (datasets), GNN note; example examples/reproducible_training.py. Tests (tests/core/test_reproducibility.py): seeded orders on both backends (same seed, different seeds, per-epoch reshuffle, restart, global RNG untouched, torch seed alone), loader independence, no RNG change from set_task (fresh and cached) or disk dataset creation, import side effects, seeded interpreter, best-epoch restore with and without logging plus the opt-out, and bit-identical MLP/RNN predictions across separate processes in memory and on disk (with and without an existing cache). Co-Authored-By: Claude Opus 5.5 --- docs/api/datasets.rst | 12 + docs/api/models/pyhealth.models.GNN.rst | 7 + docs/api/trainer.rst | 42 +++ examples/reproducible_training.py | 54 +++ .../calib/predictionset/scrib/__init__.py | 3 +- pyhealth/datasets/base_dataset.py | 3 +- pyhealth/datasets/sample_dataset.py | 75 +++- pyhealth/datasets/utils.py | 72 +++- pyhealth/interpret/methods/baseline.py | 6 +- pyhealth/interpret/methods/lime.py | 7 +- pyhealth/interpret/methods/shap.py | 7 +- .../kg_emb/datasets/splitter.py | 12 +- pyhealth/metrics/generative/utils.py | 13 +- pyhealth/models/gnn.py | 3 - pyhealth/trainer.py | 44 ++- pyhealth/utils.py | 89 +++++ tests/core/test_reproducibility.py | 329 ++++++++++++++++++ 17 files changed, 746 insertions(+), 32 deletions(-) create mode 100644 examples/reproducible_training.py create mode 100644 tests/core/test_reproducibility.py diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index 30c54a207..aa7c418e4 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -230,6 +230,18 @@ 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. Loaders do not change the dataset, so an unshuffled evaluation +loader on a training dataset leaves the training loader shuffling. See the +Trainer docs for the full checklist. + Things to Watch Out For ------------------------ diff --git a/docs/api/models/pyhealth.models.GNN.rst b/docs/api/models/pyhealth.models.GNN.rst index d7e08a9fd..6325e1936 100644 --- a/docs/api/models/pyhealth.models.GNN.rst +++ b/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: diff --git a/docs/api/trainer.rst b/docs/api/trainer.rst index 0a52af0f7..7f5b7dff6 100644 --- a/docs/api/trainer.rst +++ b/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 ------------- diff --git a/examples/reproducible_training.py b/examples/reproducible_training.py new file mode 100644 index 000000000..05450a54d --- /dev/null +++ b/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() diff --git a/pyhealth/calib/predictionset/scrib/__init__.py b/pyhealth/calib/predictionset/scrib/__init__.py index c0e703bb2..d5a5e697c 100644 --- a/pyhealth/calib/predictionset/scrib/__init__.py +++ b/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: diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py index fa71fabce..98ff1c12f 100644 --- a/pyhealth/datasets/base_dataset.py +++ b/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, diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index e2f8a1986..0ace58ff3 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/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,39 @@ def __getitem__(self, index: Any) -> Any: """ return _decode_record(super().__getitem__(index)) + def set_shuffle(self, shuffle: bool, seed: int | None = None) -> None: + """Turns shuffling on or off when iterating the dataset directly. + + Loaders from :func:`~pyhealth.datasets.get_dataloader` keep their own + shuffle setting and do not use this. + + Args: + shuffle: Whether ``iter(dataset)`` shuffles. + seed: Seed of the shuffle order (litdata reshuffles each epoch + from it). Default: drawn from torch's random generator, so + ``torch.manual_seed`` controls it. + """ + super().set_shuffle(shuffle) + if shuffle: + self.seed = _draw_seed() if seed is None else int(seed) + self.shuffler = None + + def _iter_for_loader(self, shuffle: bool, seed: int, epoch: int) -> Iterable[Any]: + """Iterates one epoch in the order a loader asked for. + + Works on a shallow copy, so loaders on the same dataset do not change + each other's (or the dataset's) shuffle setting. litdata shuffles + chunk by chunk from ``(seed, epoch)`` with its own generator, so the + global random state is not used. + """ + view = copy.copy(self) + view.shuffle = shuffle + view.seed = seed + view.shuffler = None + view._state_dict = None + view.current_epoch = epoch + 1 # litdata counts epochs from 1 + return iter(view) + def _remove_ignored_processors(self): """Remove any processors that are IgnoreProcessor instances.""" for key in [ @@ -718,9 +760,24 @@ def __init__( self._data = [builder.transform({"sample": pickle.dumps(s)}) for s in samples] self._shuffle = False + self._shuffle_seed = 0 + self._epoch = 0 + + def set_shuffle(self, shuffle: bool, seed: int | None = None) -> None: + """Turns shuffling on or off when iterating the dataset directly. + + Loaders from :func:`~pyhealth.datasets.get_dataloader` keep their own + shuffle setting and do not use this. - def set_shuffle(self, shuffle: bool) -> None: + Args: + shuffle: Whether ``iter(dataset)`` shuffles. + seed: Seed of the shuffle order; each new iteration (epoch) gets + a different order derived from it. Default: drawn from + torch's random generator, so ``torch.manual_seed`` controls it. + """ self._shuffle = shuffle + self._shuffle_seed = _draw_seed() if seed is None else int(seed) + self._epoch = 0 def __len__(self) -> int: """Returns the number of samples in the dataset. @@ -747,12 +804,15 @@ def __iter__(self) -> Iterable[Dict[str, Any]]: # type: ignore Returns: An iterator yielding processed sample dictionaries. """ - if self._shuffle: - shuffled_data = self._data[:] - random.shuffle(shuffled_data) - return iter(shuffled_data) - else: + epoch = self._epoch + self._epoch += 1 + return self._iter_for_loader(self._shuffle, self._shuffle_seed, epoch) + + def _iter_for_loader(self, shuffle: bool, seed: int, epoch: int) -> Iterable[Any]: + if not shuffle: return iter(self._data) + order = np.random.default_rng([seed, epoch]).permutation(len(self._data)) + return (self._data[i] for i in order) def subset(self, indices: Union[Sequence[int], slice]) -> SampleDataset: if isinstance(indices, slice): @@ -800,6 +860,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], diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index 24c87a1d5..91c6c7f70 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -8,7 +8,7 @@ import litdata from dateutil.parser import parse as dateutil_parse from torch.nn.utils.rnn import pad_sequence -from torch.utils.data import DataLoader +from torch.utils.data import DataLoader, IterableDataset from pyhealth import BASE_CACHE_PATH from pyhealth.utils import create_directory @@ -328,23 +328,89 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: return collated +class _LoaderView(IterableDataset): + """One loader's view of a dataset: its own shuffle flag, seed and epoch. + + Each pass over the loader is a new epoch with a new order derived from + ``(seed, epoch)``. The dataset itself is not modified, so a training and + an evaluation loader on the same dataset do not interfere. + """ + + def __init__(self, dataset, shuffle: bool, seed: int): + self.dataset = dataset + self.shuffle = shuffle + self.seed = seed + self.epoch = 0 + + def __len__(self) -> int: + return len(self.dataset) + + def __iter__(self): + epoch = self.epoch + self.epoch += 1 + return self.dataset._iter_for_loader(self.shuffle, self.seed, epoch) + + def __getattr__(self, name): + # Expose the dataset's attributes (processors, schemas, ...). + if name == "dataset": + raise AttributeError(name) + return getattr(self.dataset, name) + + def get_dataloader( - dataset: litdata.StreamingDataset, batch_size: int, shuffle: bool = False + dataset: litdata.StreamingDataset, + batch_size: int, + shuffle: bool = False, + seed: int | None = None, ) -> DataLoader: """Creates a DataLoader for a given dataset. + Shuffling is reproducible and private to the loader: each pass (epoch) + gets a new order determined by ``seed`` and the epoch number, without + touching Python's, NumPy's or PyTorch's global random state, and without + changing the dataset (other loaders on it are unaffected). This holds for + in-memory and disk-backed datasets alike; the two backends give different + orders for the same seed. + Args: dataset: The dataset to load data from. batch_size: The number of samples per batch. shuffle: Whether to shuffle the data at every epoch. + seed: Seed of the shuffle order. Default None draws one from torch's + random generator when the loader is created, so + ``torch.manual_seed(...)`` before this call makes runs + reproducible. Unshuffled loaders draw nothing. Returns: A DataLoader instance for the dataset. + + Examples: + >>> from pyhealth.datasets import create_sample_dataset, get_dataloader + >>> samples = [{"patient_id": f"p{i}", "x": [float(i)], "y": i % 2} for i in range(8)] + >>> ds = create_sample_dataset(samples, {"x": "tensor"}, {"y": "binary"}) + >>> order = lambda loader: [pid for b in loader for pid in b["patient_id"]] + >>> order(get_dataloader(ds, 4, shuffle=True, seed=1)) == order( + ... get_dataloader(ds, 4, shuffle=True, seed=1)) + True """ - dataset.set_shuffle(shuffle) + if shuffle and seed is None: + seed = int(torch.randint(0, 2**31 - 1, ()).item()) + seed = 0 if seed is None else int(seed) + # DataLoader draws a base seed from this generator on every pass; a + # private one keeps that draw off torch's global random state. + generator = torch.Generator().manual_seed(seed) + map_style_shuffle = False + if hasattr(dataset, "_iter_for_loader"): + dataset = _LoaderView(dataset, shuffle, seed) + elif hasattr(dataset, "set_shuffle"): + dataset.set_shuffle(shuffle) + else: + map_style_shuffle = shuffle # e.g. a torch map-style Dataset dataloader = DataLoader( dataset, batch_size=batch_size, + shuffle=map_style_shuffle or None, + generator=generator, collate_fn=collate_fn_dict_with_padding, ) diff --git a/pyhealth/interpret/methods/baseline.py b/pyhealth/interpret/methods/baseline.py index 8104023d1..67629c516 100644 --- a/pyhealth/interpret/methods/baseline.py +++ b/pyhealth/interpret/methods/baseline.py @@ -13,6 +13,8 @@ import torch +from pyhealth.utils import _seeded_by + from pyhealth.models import BaseModel from .base_interpreter import BaseInterpreter @@ -88,6 +90,7 @@ def __init__( # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ + @_seeded_by("random_seed") def attribute( self, **kwargs: torch.Tensor | tuple[torch.Tensor, ...], @@ -107,9 +110,6 @@ def attribute( Dictionary mapping each feature key to a random attribution tensor whose shape matches the raw input values. """ - if self.random_seed is not None: - torch.manual_seed(self.random_seed) - device = next(self.model.parameters()).device # Filter kwargs to only include model feature keys diff --git a/pyhealth/interpret/methods/lime.py b/pyhealth/interpret/methods/lime.py index 4e407fccf..3f14e5c89 100644 --- a/pyhealth/interpret/methods/lime.py +++ b/pyhealth/interpret/methods/lime.py @@ -4,6 +4,8 @@ from typing import Dict, Optional, Tuple, Callable, Union, cast import torch + +from pyhealth.utils import _seeded_by import torch.nn.functional as F from torch.nn import CosineSimilarity @@ -162,6 +164,7 @@ def __init__( # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ + @_seeded_by("random_seed") def attribute( self, baseline: Optional[Dict[str, torch.Tensor]] = None, @@ -204,10 +207,6 @@ def attribute( ... ) >>> print(lime_values['conditions']) # Shape: (1, 3) """ - # Set random seed for reproducibility if specified - if self.random_seed is not None: - torch.manual_seed(self.random_seed) - device = next(self.model.parameters()).device # Filter kwargs to only include model feature keys and ensure they are tuples diff --git a/pyhealth/interpret/methods/shap.py b/pyhealth/interpret/methods/shap.py index df52ea732..39163bdb3 100644 --- a/pyhealth/interpret/methods/shap.py +++ b/pyhealth/interpret/methods/shap.py @@ -5,6 +5,8 @@ import torch +from pyhealth.utils import _seeded_by + from pyhealth.models import BaseModel from pyhealth.interpret.api import Interpretable from .base_interpreter import BaseInterpreter @@ -135,6 +137,7 @@ def __init__( # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ + @_seeded_by("random_seed") def attribute( self, baseline: Optional[Dict[str, torch.Tensor]] = None, @@ -176,10 +179,6 @@ def attribute( ... ) >>> print(shap_values['x_continuous']) # Shape: (1, 3) """ - # Set random seed for reproducibility if specified - if self.random_seed is not None: - torch.manual_seed(self.random_seed) - device = next(self.model.parameters()).device # Filter kwargs to only include model feature keys and ensure they are tuples diff --git a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py index 559ecb7c4..71681b16e 100644 --- a/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py +++ b/pyhealth/medcode/pretrained_embeddings/kg_emb/datasets/splitter.py @@ -26,13 +26,19 @@ def split( Note: The original dataset can be accessed by `train_dataset.dataset`, `val_dataset.dataset`, and `test_dataset.dataset`. + The seed shuffles with a private generator; NumPy's global random + state is not changed. + + Examples: + >>> train, val, test = split(dataset, [0.8, 0.1, 0.1], seed=0) # doctest: +SKIP """ - if seed is not None: - np.random.seed(seed) assert sum(ratios) == 1.0, "ratios must sum to 1.0" index = np.arange(len(dataset)) - np.random.shuffle(index) + # A private generator: same order as seeding the global one, without + # changing the caller's global NumPy state. + rng = np.random.RandomState(seed) if seed is not None else np.random + rng.shuffle(index) train_index = index[: int(len(dataset) * ratios[0])] val_index = index[ int(len(dataset) * ratios[0]) : int(len(dataset) * (ratios[0] + ratios[1])) diff --git a/pyhealth/metrics/generative/utils.py b/pyhealth/metrics/generative/utils.py index d5f89d216..09bd5aedb 100644 --- a/pyhealth/metrics/generative/utils.py +++ b/pyhealth/metrics/generative/utils.py @@ -14,6 +14,8 @@ import torch import torch.nn as nn +from pyhealth.utils import preserve_rng_state + __all__ = [ "summarize_metric_runs", "convert_visits_to_sets", @@ -286,6 +288,7 @@ def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: return probs.squeeze(-1) +@preserve_rng_state() def train_lstm_model( train_ehr: pd.DataFrame, test_ehr: pd.DataFrame, @@ -319,8 +322,16 @@ def train_lstm_model( Returns: A tuple ``(model, y_true, y_pred)`` where ``y_true`` and ``y_pred`` are numpy arrays of test labels and binary predictions. + + The seed applies to this call only: the caller's global random state is + restored on return. + + Examples: + >>> model, y_true, y_pred = train_lstm_model( # doctest: +SKIP + ... train_ehr, test_ehr, seed=4 + ... ) """ - torch.manual_seed(seed) + torch.manual_seed(seed) # restored on return by preserve_rng_state all_codes = set() all_codes.update(train_ehr[code_col].unique().tolist()) all_codes.update(test_ehr[code_col].unique().tolist()) diff --git a/pyhealth/models/gnn.py b/pyhealth/models/gnn.py index ccd83e230..f3eec1eb0 100644 --- a/pyhealth/models/gnn.py +++ b/pyhealth/models/gnn.py @@ -17,9 +17,6 @@ from pyhealth.processors import SequenceProcessor from pyhealth.models.embedding import EmbeddingModel -torch.manual_seed(3) -np.random.seed(1) - """Graph Neural Network models for PyHealth. This module provides implementations of Graph Convolutional Network (GCN) and diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index e6939be93..91e7e8d04 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -1,3 +1,4 @@ +import copy import logging import os from datetime import datetime @@ -11,13 +12,25 @@ from tqdm import tqdm from tqdm.autonotebook import trange -from pyhealth.metrics import (binary_metrics_fn, multiclass_metrics_fn, - multilabel_metrics_fn, regression_metrics_fn) +from pyhealth.metrics import ( + binary_metrics_fn, + multiclass_metrics_fn, + multilabel_metrics_fn, + regression_metrics_fn, +) from pyhealth.utils import create_directory logger = logging.getLogger(__name__) +def _cpu_state_dict(model: nn.Module) -> dict: + """A detached CPU copy of the model's state, safe from later updates.""" + return { + key: value.detach().to("cpu", copy=True) if torch.is_tensor(value) else copy.deepcopy(value) + for key, value in model.state_dict().items() + } + + def is_best(best_score: float, score: float, monitor_criterion: str) -> bool: if monitor_criterion == "max": return score > best_score @@ -133,6 +146,7 @@ def train( monitor_criterion: str = "max", load_best_model_at_last: bool = True, patience=None, + keep_best_in_memory: bool = True, ): """Trains the model. @@ -152,6 +166,14 @@ def train( Default is True. patience: Number of epochs to wait for improvement before early stopping. Default is None, which means no early stopping. + keep_best_in_memory: When ``monitor`` is set and the trainer has no + output directory (``enable_logging=False``), keep a CPU copy of + the best epoch's weights in memory so ``load_best_model_at_last`` + can restore them. This costs one extra copy of the model's + weights in CPU memory; set False to save it, in which case the + last epoch's weights are kept (with a warning). With an output + directory the best weights are saved to ``best.ckpt`` instead. + Default is True. Raises: TypeError: If the model is fit outside the trainer (it sets @@ -199,6 +221,19 @@ def train( # initialize data_iterator = iter(train_dataloader) best_score = -1 * float("inf") if monitor_criterion == "max" else float("inf") + best_state = None + restore_from_memory = ( + load_best_model_at_last + and monitor is not None + and val_dataloader is not None + and self.exp_path is None + ) + if restore_from_memory and not keep_best_in_memory: + logger.warning( + "No output directory and keep_best_in_memory=False: the model " + "will keep the last epoch's weights, not the best epoch's." + ) + restore_from_memory = False if steps_per_epoch is None: steps_per_epoch = len(train_dataloader) global_step = 0 @@ -259,6 +294,8 @@ def train( patience_counter = 0 if self.exp_path is not None: self.save_ckpt(os.path.join(self.exp_path, "best.ckpt")) + elif restore_from_memory: + best_state = _cpu_state_dict(self.model) else: patience_counter += 1 # early stopping @@ -273,6 +310,9 @@ def train( os.path.join(self.exp_path, "best.ckpt")): logger.info("Loaded best model") self.load_ckpt(os.path.join(self.exp_path, "best.ckpt")) + elif restore_from_memory and best_state is not None: + logger.info("Loaded best model (kept in memory)") + self.model.load_state_dict(best_state) # test if test_dataloader is not None: diff --git a/pyhealth/utils.py b/pyhealth/utils.py index cc0275d23..8af089b1a 100644 --- a/pyhealth/utils.py +++ b/pyhealth/utils.py @@ -3,6 +3,7 @@ import pickle import random import contextlib +import functools import shutil import urllib.request @@ -56,6 +57,34 @@ def download_file(url: str, path: str, timeout: float = 60.0) -> str: def set_seed(seed): + """Seeds Python's ``random``, NumPy and PyTorch (CPU and CUDA). + + This is the only place PyHealth seeds the global random generators, and + only when you call it. ``torch.manual_seed(seed)`` alone is also enough + for reproducible PyHealth training on CPU: shuffling loaders draw their + seed from torch's generator (see :func:`~pyhealth.datasets.get_dataloader`). + + On CUDA this also sets ``torch.backends.cudnn.deterministic = True`` and + ``torch.backends.cudnn.benchmark = False``, which apply to the whole + process. For bit-identical GPU runs you may additionally need + ``torch.use_deterministic_algorithms(True)`` and the environment variable + ``CUBLAS_WORKSPACE_CONFIG=:4096:8``. + + ``PYTHONHASHSEED`` is set in ``os.environ``; it does not change string + hashing in the running interpreter, only in subprocesses started + afterwards (such as ``set_task`` workers). + + Args: + seed: The seed. + + Examples: + >>> import random + >>> from pyhealth.utils import set_seed + >>> set_seed(0); a = random.random() + >>> set_seed(0); b = random.random() + >>> a == b + True + """ random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) @@ -67,6 +96,66 @@ def set_seed(seed): os.environ["PYTHONHASHSEED"] = str(seed) +@contextlib.contextmanager +def preserve_rng_state(): + """Restores the global random states after the block, whatever it did. + + Saves and restores Python's ``random``, NumPy's legacy global generator, + PyTorch's CPU generator and, if CUDA is initialized, every CUDA + generator. PyHealth wraps cache writes in it because litdata reseeds + the global generators while writing, which would otherwise silently + replace the caller's seed. + + Examples: + >>> import random + >>> from pyhealth.utils import preserve_rng_state + >>> random.seed(1); state = random.getstate() + >>> with preserve_rng_state(): + ... random.seed(42) + >>> random.getstate() == state + True + """ + py_state = random.getstate() + np_state = np.random.get_state() + torch_state = torch.get_rng_state() + cuda_states = ( + torch.cuda.get_rng_state_all() + if torch.cuda.is_available() and torch.cuda.is_initialized() + else None + ) + try: + yield + finally: + random.setstate(py_state) + np.random.set_state(np_state) + torch.set_rng_state(torch_state) + if cuda_states is not None: + torch.cuda.set_rng_state_all(cuda_states) + + +def _seeded_by(attribute: str): + """Method decorator: seed torch from ``self.`` for the call only. + + If the attribute is None the method runs unchanged; otherwise torch is + seeded inside :func:`preserve_rng_state`, so the call is reproducible and + the caller's global random state is left as it was. + """ + + def decorator(method): + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + seed = getattr(self, attribute, None) + if seed is None: + return method(self, *args, **kwargs) + with preserve_rng_state(): + torch.manual_seed(seed) + return method(self, *args, **kwargs) + + return wrapper + + return decorator + + def create_directory(directory): if not os.path.exists(directory): os.makedirs(directory) diff --git a/tests/core/test_reproducibility.py b/tests/core/test_reproducibility.py new file mode 100644 index 000000000..07719b2c8 --- /dev/null +++ b/tests/core/test_reproducibility.py @@ -0,0 +1,329 @@ +"""Reproducible training: seeded per-loader shuffles, no global RNG side +effects from cache writes or imports, and best-epoch restore without logging. +""" + +import os +import random +import subprocess +import sys +import tempfile +import textwrap +import unittest +from pathlib import Path + +import numpy as np +import torch + +from pyhealth.datasets import BaseDataset, create_sample_dataset, get_dataloader +from pyhealth.models import MLP +from pyhealth.tasks import BaseTask +from pyhealth.trainer import Trainer +from pyhealth.utils import preserve_rng_state + +N = 40 + + +def _samples(n=N): + return [ + {"patient_id": f"p{i}", "x": [float(i % 7), float(i % 3)], "y": i % 2} + for i in range(n) + ] + + +def _dataset(in_memory): + return create_sample_dataset(_samples(), {"x": "tensor"}, {"y": "binary"}, in_memory=in_memory) + + +def _order(loader): + return [pid for batch in loader for pid in batch["patient_id"]] + + +def _rng_states(): + return ( + random.getstate(), + tuple(np.asarray(s).tobytes() if isinstance(s, np.ndarray) else s for s in np.random.get_state()), + torch.get_rng_state().numpy().tobytes(), + ) + + +class EventTask(BaseTask): + task_name = "repro_task" + input_schema = {"codes": "sequence", "x": "tensor"} + output_schema = {"y": "binary"} + + def __call__(self, patient): + k = int(patient.patient_id[1:]) + codes = [e.code for e in patient.get_events("events")] + return [{"patient_id": patient.patient_id, "codes": codes, + "x": [float(k % 5), float(len(codes))], "y": k % 2}] + + +def _write_events(root: Path, n_patients=30): + root.mkdir(parents=True, exist_ok=True) + rows = ["patient_id,time,code"] + [ + f"p{k},2020-01-0{1 + v} 00:00:00,c{(k + v) % 6}" for k in range(n_patients) for v in range(1 + k % 3) + ] + (root / "events.csv").write_text("\n".join(rows) + "\n") + (root / "config.yaml").write_text( + 'version: "1.0"\ntables:\n events:\n file_path: "events.csv"\n' + ' patient_id: "patient_id"\n timestamp: "time"\n attributes:\n - "code"\n' + ) + + +class TestSeededShuffle(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.datasets = {"memory": _dataset(True), "disk": _dataset(False)} + + def test_same_seed_same_orders_over_epochs(self): + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + a = get_dataloader(ds, 8, shuffle=True, seed=3) + b = get_dataloader(ds, 8, shuffle=True, seed=3) + epochs_a = [_order(a), _order(a)] + epochs_b = [_order(b), _order(b)] + self.assertEqual(epochs_a, epochs_b) + self.assertNotEqual(epochs_a[0], epochs_a[1]) # new order each epoch + self.assertEqual(sorted(epochs_a[0]), sorted(f"p{i}" for i in range(N))) + # a new loader with the same seed restarts the same sequence + c = get_dataloader(ds, 8, shuffle=True, seed=3) + self.assertEqual([_order(c), _order(c)], epochs_a) + + def test_different_seeds_differ(self): + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + self.assertNotEqual( + _order(get_dataloader(ds, 8, shuffle=True, seed=1)), + _order(get_dataloader(ds, 8, shuffle=True, seed=2)), + ) + + def test_iterating_leaves_global_rng_untouched(self): + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + loader = get_dataloader(ds, 8, shuffle=True, seed=5) + before = _rng_states() + _order(loader), _order(loader) + self.assertEqual(_rng_states(), before) + + def test_torch_seed_alone_controls_default_seed(self): + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + torch.manual_seed(11) + first = _order(get_dataloader(ds, 8, shuffle=True)) + torch.manual_seed(11) + second = _order(get_dataloader(ds, 8, shuffle=True)) + torch.manual_seed(12) + other = _order(get_dataloader(ds, 8, shuffle=True)) + self.assertEqual(first, second) + self.assertNotEqual(first, other) + + def test_unshuffled_loader_keeps_dataset_order(self): + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + self.assertEqual(_order(get_dataloader(ds, 8)), [f"p{i}" for i in range(N)]) + + def test_loaders_on_one_dataset_are_independent(self): + natural = [f"p{i}" for i in range(N)] + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + train = get_dataloader(ds, 8, shuffle=True, seed=4) + expected = _order(get_dataloader(ds, 8, shuffle=True, seed=4)) + evaluation = get_dataloader(ds, 8, shuffle=False) + self.assertEqual(_order(evaluation), natural) + self.assertEqual(_order(train), expected) + self.assertNotEqual(expected, natural) + + def test_set_shuffle_on_dataset_is_seeded_too(self): + ds = _dataset(True) + ds.set_shuffle(True, seed=9) + first = [s["patient_id"] for s in ds] + before = _rng_states() + second = [s["patient_id"] for s in ds] + self.assertEqual(_rng_states(), before) + self.assertNotEqual(first, second) + ds.set_shuffle(True, seed=9) + self.assertEqual([s["patient_id"] for s in ds], first) + + +class TestNoGlobalRngSideEffects(unittest.TestCase): + def test_disk_dataset_creation(self): + before = _rng_states() + create_sample_dataset(_samples(), {"x": "tensor"}, {"y": "binary"}, in_memory=False) + self.assertEqual(_rng_states(), before) + + def test_set_task_cache_write_and_reuse(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) / "data" + _write_events(root) + make = lambda: BaseDataset( + root=str(root), tables=["events"], dataset_name="Repro", + config_path=str(root / "config.yaml"), cache_dir=str(Path(tmp) / "cache"), + ) + orders = [] + for _ in range(2): # first call writes the cache, second reuses it + torch.manual_seed(21) + random.seed(21) + np.random.seed(21) + before = _rng_states() + samples = make().set_task(EventTask()) + self.assertEqual(_rng_states(), before) + orders.append(_order(get_dataloader(samples, 4, shuffle=True))) + self.assertEqual(orders[0], orders[1]) + + def test_preserve_rng_state(self): + before = _rng_states() + with preserve_rng_state(): + random.seed(42) + np.random.seed(42) + torch.manual_seed(42) + self.assertEqual(_rng_states(), before) + + def test_seeded_interpreter_does_not_reseed_caller(self): + from pyhealth.interpret.methods import RandomBaseline + + ds = _dataset(True) + model = MLP(dataset=ds) + batch = next(iter(get_dataloader(ds, 8))) + explainer = RandomBaseline(model, random_seed=0) + torch.manual_seed(1) + before = _rng_states() + first = explainer.attribute(**batch) + self.assertEqual(_rng_states(), before) + second = explainer.attribute(**batch) + self.assertTrue(all(torch.equal(first[k], second[k]) for k in first)) + + def test_importing_models_does_not_reseed(self): + script = textwrap.dedent(""" + import numpy as np, torch + np.random.seed(123); torch.manual_seed(123) + a = (np.random.get_state()[1].tobytes(), torch.get_rng_state().numpy().tobytes()) + import pyhealth.models + b = (np.random.get_state()[1].tobytes(), torch.get_rng_state().numpy().tobytes()) + print("unchanged" if a == b else "changed") + """) + out = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, check=True) + self.assertEqual(out.stdout.strip().splitlines()[-1], "unchanged") + + +class TestBestEpochRestore(unittest.TestCase): + """Validation labels are the inverse of training labels, so validation + PR-AUC is best early and falls as the model fits the training data.""" + + @classmethod + def setUpClass(cls): + rng = np.random.default_rng(0) + x = rng.normal(size=(200, 4)) + y = (x[:, 0] > 0).astype(int) + train = [{"patient_id": f"t{i}", "x": x[i].tolist(), "y": int(y[i])} for i in range(200)] + val = [{"patient_id": f"v{i}", "x": x[i].tolist(), "y": int(1 - y[i])} for i in range(200)] + cls.train_ds = create_sample_dataset(train, {"x": "tensor"}, {"y": "binary"}) + cls.val_ds = create_sample_dataset( + val, {"x": "tensor"}, {"y": "binary"}, + input_processors=cls.train_ds.input_processors, + output_processors=cls.train_ds.output_processors, + ) + + def _train(self, enable_logging, keep_best_in_memory=True, output_path=None): + torch.manual_seed(0) + model = MLP(dataset=self.train_ds, hidden_dim=16) + trainer = Trainer(model=model, metrics=["pr_auc"], enable_logging=enable_logging, + output_path=output_path) + scores = [] + evaluate = trainer.evaluate + trainer.evaluate = lambda loader: scores.append(evaluate(loader)) or scores[-1] + trainer.train( + get_dataloader(self.train_ds, 32, shuffle=True, seed=0), + get_dataloader(self.val_ds, 64), + epochs=12, monitor="pr_auc", patience=3, + optimizer_params={"lr": 0.05}, keep_best_in_memory=keep_best_in_memory, + ) + trainer.evaluate = evaluate + epoch_scores = [s["pr_auc"] for s in scores] + final = trainer.evaluate(get_dataloader(self.val_ds, 64))["pr_auc"] + return model, epoch_scores, final + + def test_best_epoch_restored_without_logging(self): + model, epoch_scores, final = self._train(enable_logging=False) + best = max(epoch_scores) + self.assertLess(epoch_scores.index(best), len(epoch_scores) - 1, "best epoch should not be the last") + self.assertEqual(final, best) + + with tempfile.TemporaryDirectory() as tmp: + logged, logged_scores, logged_final = self._train(enable_logging=True, output_path=tmp) + self.assertEqual(logged_scores, epoch_scores) + self.assertEqual(logged_final, best) + on_disk = logged.state_dict() + for key, value in model.state_dict().items(): + self.assertTrue(torch.equal(value, on_disk[key]), key) + + def test_opt_out_keeps_last_epoch(self): + with self.assertLogs("pyhealth.trainer", level="WARNING"): + _, epoch_scores, final = self._train(enable_logging=False, keep_best_in_memory=False) + self.assertEqual(final, epoch_scores[-1]) + self.assertLess(final, max(epoch_scores)) + + +E2E_SCRIPT = textwrap.dedent(""" + import hashlib, logging, sys, warnings + from pathlib import Path + logging.disable(logging.WARNING); warnings.simplefilter("ignore") + import torch + sys.path.insert(0, sys.argv[5]) + from test_reproducibility import EventTask, _write_events + from pyhealth.datasets import BaseDataset, create_sample_dataset, get_dataloader + from pyhealth.models import MLP, RNN + from pyhealth.trainer import Trainer + + def main(): + seed, backend, cache = int(sys.argv[1]), sys.argv[2], Path(sys.argv[4]) + torch.manual_seed(seed) + if backend == "memory": + samples = [{"patient_id": f"p{i}", "codes": [f"c{(i + j) % 6}" for j in range(1 + i % 3)], + "x": [float(i % 5), float(1 + i % 3)], "y": i % 2} for i in range(30)] + ds = create_sample_dataset(samples, EventTask.input_schema, EventTask.output_schema) + else: + _write_events(cache / "data") + ds = BaseDataset(root=str(cache / "data"), tables=["events"], dataset_name="E2E", + config_path=str(cache / "data" / "config.yaml"), + cache_dir=str(cache / "cache")).set_task(EventTask()) + hashes = [] + for model_cls in (MLP, RNN): + trainer = Trainer(model=model_cls(dataset=ds), enable_logging=False) + trainer.train(get_dataloader(ds, 8, shuffle=True), epochs=3) + _, y_prob, _ = trainer.inference(get_dataloader(ds, 8)) + hashes.append(hashlib.sha256(y_prob.tobytes()).hexdigest()[:16]) + print(" ".join(hashes)) + + if __name__ == "__main__": + main() +""") + + +class TestEndToEndAcrossProcesses(unittest.TestCase): + def _run(self, script, seed, backend, cache): + """Trains MLP and RNN in a fresh process; returns their prediction hashes.""" + env = dict(os.environ, PYTHONPATH=os.pathsep.join(sys.path)) + out = subprocess.run( + [sys.executable, script, str(seed), backend, "-", str(cache), str(Path(__file__).parent)], + capture_output=True, text=True, env=env, timeout=600, + ) + self.assertEqual(out.returncode, 0, out.stderr[-2000:]) + return out.stdout.strip().splitlines()[-1] + + def test_bit_identical_predictions(self): + with tempfile.TemporaryDirectory() as tmp: + script = Path(tmp) / "e2e.py" + script.write_text(E2E_SCRIPT) + for backend in ("memory", "disk"): + with self.subTest(backend=backend): + cache = Path(tmp) / backend + first = self._run(script, 7, backend, cache) # disk: writes the cache + second = self._run(script, 7, backend, cache) # disk: reuses it + other = self._run(script, 8, backend, cache) + self.assertEqual(first, second) # MLP and RNN hashes both equal + for a, b in zip(first.split(), other.split()): + self.assertNotEqual(a, b) + + +if __name__ == "__main__": + unittest.main() From 81ad324a9170742359d9b4ad13921d636fdcdd4c Mon Sep 17 00:00:00 2001 From: solarsys Date: Mon, 12 Oct 2026 14:09:48 -0700 Subject: [PATCH 2/2] Drop _LoaderView: per-loader shuffle settings applied at each pass Review feedback on #1280: the _LoaderView wrapper hid the dataset behind loader.dataset, relied on litdata private fields, and its epoch counter did not advance with num_workers > 0 (iteration runs on worker copies, so every epoch repeated the same order). get_dataloader now returns a small DataLoader subclass that keeps its own (shuffle, seed, epoch) and applies them to the dataset with set_shuffle and set_epoch at the start of each pass, in the main process before workers copy the dataset. loader.dataset is the dataset again, only public litdata calls are used (set_shuffle, set_epoch, seed), and each epoch gets a new reproducible order with or without workers. The in-memory dataset also splits its samples across workers; it used to yield every sample once per worker. Co-Authored-By: Claude Opus 5.5 --- docs/api/datasets.rst | 6 +- pyhealth/datasets/sample_dataset.py | 92 +++++++++++++++++------------ pyhealth/datasets/utils.py | 72 ++++++++++++---------- tests/core/test_reproducibility.py | 22 +++++++ 4 files changed, 119 insertions(+), 73 deletions(-) diff --git a/docs/api/datasets.rst b/docs/api/datasets.rst index aa7c418e4..f629bef8a 100644 --- a/docs/api/datasets.rst +++ b/docs/api/datasets.rst @@ -238,8 +238,10 @@ 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. Loaders do not change the dataset, so an unshuffled evaluation -loader on a training dataset leaves the training loader shuffling. See the +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 diff --git a/pyhealth/datasets/sample_dataset.py b/pyhealth/datasets/sample_dataset.py index 0ace58ff3..2c6e0f19b 100644 --- a/pyhealth/datasets/sample_dataset.py +++ b/pyhealth/datasets/sample_dataset.py @@ -538,37 +538,26 @@ def __getitem__(self, index: Any) -> Any: return _decode_record(super().__getitem__(index)) def set_shuffle(self, shuffle: bool, seed: int | None = None) -> None: - """Turns shuffling on or off when iterating the dataset directly. + """Sets whether iterating the dataset shuffles, and the shuffle seed. - Loaders from :func:`~pyhealth.datasets.get_dataloader` keep their own - shuffle setting and do not use this. + 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 (litdata reshuffles each epoch - from it). Default: drawn from torch's random generator, so - ``torch.manual_seed`` controls it. + 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) - self.shuffler = None - - def _iter_for_loader(self, shuffle: bool, seed: int, epoch: int) -> Iterable[Any]: - """Iterates one epoch in the order a loader asked for. - - Works on a shallow copy, so loaders on the same dataset do not change - each other's (or the dataset's) shuffle setting. litdata shuffles - chunk by chunk from ``(seed, epoch)`` with its own generator, so the - global random state is not used. - """ - view = copy.copy(self) - view.shuffle = shuffle - view.seed = seed - view.shuffler = None - view._state_dict = None - view.current_epoch = epoch + 1 # litdata counts epochs from 1 - return iter(view) def _remove_ignored_processors(self): """Remove any processors that are IgnoreProcessor instances.""" @@ -761,23 +750,43 @@ def __init__( self._shuffle = False self._shuffle_seed = 0 - self._epoch = 0 + self._epoch = 1 # counted from 1, as litdata does def set_shuffle(self, shuffle: bool, seed: int | None = None) -> None: - """Turns shuffling on or off when iterating the dataset directly. + """Sets whether iterating the dataset shuffles, and the shuffle seed. - Loaders from :func:`~pyhealth.datasets.get_dataloader` keep their own - shuffle setting and do not use this. + 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; each new iteration (epoch) gets - a different order derived from it. Default: drawn from - torch's random generator, so ``torch.manual_seed`` controls it. + 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 - self._shuffle_seed = _draw_seed() if seed is None else int(seed) - self._epoch = 0 + 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. @@ -804,14 +813,19 @@ def __iter__(self) -> Iterable[Dict[str, Any]]: # type: ignore Returns: An iterator yielding processed sample dictionaries. """ - epoch = self._epoch + if self._shuffle: + rng = np.random.default_rng([self._shuffle_seed, self._epoch]) + order = rng.permutation(len(self._data)) + else: + order = np.arange(len(self._data)) self._epoch += 1 - return self._iter_for_loader(self._shuffle, self._shuffle_seed, epoch) - - def _iter_for_loader(self, shuffle: bool, seed: int, epoch: int) -> Iterable[Any]: - if not shuffle: - return iter(self._data) - order = np.random.default_rng([seed, epoch]).permutation(len(self._data)) + # 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: diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py index 91c6c7f70..fa0acf5f7 100644 --- a/pyhealth/datasets/utils.py +++ b/pyhealth/datasets/utils.py @@ -8,7 +8,7 @@ import litdata from dateutil.parser import parse as dateutil_parse from torch.nn.utils.rnn import pad_sequence -from torch.utils.data import DataLoader, IterableDataset +from torch.utils.data import DataLoader from pyhealth import BASE_CACHE_PATH from pyhealth.utils import create_directory @@ -328,33 +328,33 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict: return collated -class _LoaderView(IterableDataset): - """One loader's view of a dataset: its own shuffle flag, seed and epoch. +class _SeededDataLoader(DataLoader): + """A DataLoader that applies its own shuffle settings at every pass. - Each pass over the loader is a new epoch with a new order derived from - ``(seed, epoch)``. The dataset itself is not modified, so a training and - an evaluation loader on the same dataset do not interfere. + PyHealth datasets are iterable, so DataLoader cannot shuffle them; the + dataset shuffles itself from a seed and an epoch number. This loader keeps + its own ``(shuffle, seed, epoch)`` and applies them to the dataset at the + start of each pass, in the main process. That way: + + - each epoch gets a new, reproducible order, also with ``num_workers > 0`` + (workers receive a copy of the dataset after the settings are applied); + - a training and an evaluation loader on the same dataset do not change + each other's order. + + ``loader.dataset`` is the dataset itself. """ - def __init__(self, dataset, shuffle: bool, seed: int): - self.dataset = dataset - self.shuffle = shuffle + def __init__(self, dataset, *, shuffle_samples: bool, seed: int, **kwargs): + super().__init__(dataset, **kwargs) + self.shuffle_samples = shuffle_samples self.seed = seed - self.epoch = 0 - - def __len__(self) -> int: - return len(self.dataset) + self.epoch = 1 # counted from 1, as litdata does def __iter__(self): - epoch = self.epoch + self.dataset.set_shuffle(self.shuffle_samples, self.seed) + self.dataset.set_epoch(self.epoch) self.epoch += 1 - return self.dataset._iter_for_loader(self.shuffle, self.seed, epoch) - - def __getattr__(self, name): - # Expose the dataset's attributes (processors, schemas, ...). - if name == "dataset": - raise AttributeError(name) - return getattr(self.dataset, name) + return super().__iter__() def get_dataloader( @@ -365,10 +365,11 @@ def get_dataloader( ) -> DataLoader: """Creates a DataLoader for a given dataset. - Shuffling is reproducible and private to the loader: each pass (epoch) - gets a new order determined by ``seed`` and the epoch number, without - touching Python's, NumPy's or PyTorch's global random state, and without - changing the dataset (other loaders on it are unaffected). This holds for + Shuffling is reproducible: each pass (epoch) gets a new order determined + by ``seed`` and the epoch number, without touching Python's, NumPy's or + PyTorch's global random state. The loader applies its settings to the + dataset at the start of each pass, so several loaders can share one + dataset as long as they are not iterated at the same time. This holds for in-memory and disk-backed datasets alike; the two backends give different orders for the same seed. @@ -399,17 +400,24 @@ def get_dataloader( # DataLoader draws a base seed from this generator on every pass; a # private one keeps that draw off torch's global random state. generator = torch.Generator().manual_seed(seed) - map_style_shuffle = False - if hasattr(dataset, "_iter_for_loader"): - dataset = _LoaderView(dataset, shuffle, seed) - elif hasattr(dataset, "set_shuffle"): + from .sample_dataset import SampleDataset + + if isinstance(dataset, SampleDataset): + return _SeededDataLoader( + dataset, + shuffle_samples=shuffle, + seed=seed, + batch_size=batch_size, + generator=generator, + collate_fn=collate_fn_dict_with_padding, + ) + if hasattr(dataset, "set_shuffle"): # another iterable dataset, e.g. litdata's dataset.set_shuffle(shuffle) - else: - map_style_shuffle = shuffle # e.g. a torch map-style Dataset + shuffle = False # DataLoader cannot shuffle an iterable dataset dataloader = DataLoader( dataset, batch_size=batch_size, - shuffle=map_style_shuffle or None, + shuffle=shuffle or None, generator=generator, collate_fn=collate_fn_dict_with_padding, ) diff --git a/tests/core/test_reproducibility.py b/tests/core/test_reproducibility.py index 07719b2c8..6e02281aa 100644 --- a/tests/core/test_reproducibility.py +++ b/tests/core/test_reproducibility.py @@ -133,6 +133,28 @@ def test_loaders_on_one_dataset_are_independent(self): self.assertEqual(_order(train), expected) self.assertNotEqual(expected, natural) + def test_loader_dataset_is_the_dataset(self): + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + self.assertIs(get_dataloader(ds, 8, shuffle=True, seed=1).dataset, ds) + + def test_worker_processes_get_new_order_each_epoch(self): + # With workers, iteration happens on copies of the dataset, so the + # epoch must advance in the main process. + natural = sorted(f"p{i}" for i in range(N)) + for name, ds in self.datasets.items(): + with self.subTest(backend=name): + runs = [] + for _ in range(2): + loader = get_dataloader(ds, 8, shuffle=True, seed=6) + loader.num_workers, loader.prefetch_factor = 2, 2 + runs.append([_order(loader) for _ in range(3)]) + epochs = runs[0] + self.assertEqual(runs[0], runs[1]) + self.assertEqual(len(set(map(tuple, epochs))), 3) + for epoch in epochs: # every sample exactly once + self.assertEqual(sorted(epoch), natural) + def test_set_shuffle_on_dataset_is_seeded_too(self): ds = _dataset(True) ds.set_shuffle(True, seed=9)