From b2f941b751bcaf2723f42150a70cad8e8d14151d Mon Sep 17 00:00:00 2001 From: Azra Bano Date: Mon, 24 Aug 2026 02:50:15 -0400 Subject: [PATCH] fix: compute evaluation metrics once on pooled predictions (#859) trainer.evaluate now documents and guarantees pooled-metric semantics, and the evaluation loss is pooled as an example-weighted mean instead of an unweighted mean of per-batch means. Previously a partial final batch was over-weighted, making the reported loss (and anything monitoring it) depend on the evaluation batch size and shuffle order, so repeated benchmark runs disagreed even with a fixed seed. - inference(): accumulate loss as sum(batch_mean * batch_size) / n_samples - evaluate() mode=None path: same example-weighted pooling, with batch size inferred from model outputs or the input batch - add regression tests: determinism under a fixed seed, batch-order invariance, batch-size invariance (16 vs 64), pooled AUROC/AUPRC/loss match sklearn/torch computed directly on the full dataset - document pooled evaluation semantics in docs/api/trainer.rst and add examples/trainer_pooled_evaluation.py --- docs/api/trainer.rst | 17 +++ examples/trainer_pooled_evaluation.py | 94 +++++++++++++ pyhealth/trainer.py | 68 ++++++++-- tests/core/test_trainer_pooled_eval.py | 179 +++++++++++++++++++++++++ 4 files changed, 349 insertions(+), 9 deletions(-) create mode 100644 examples/trainer_pooled_evaluation.py create mode 100644 tests/core/test_trainer_pooled_eval.py diff --git a/docs/api/trainer.rst b/docs/api/trainer.rst index 0a52af0f7..bae394038 100644 --- a/docs/api/trainer.rst +++ b/docs/api/trainer.rst @@ -77,6 +77,23 @@ Controlling the Training Loop restores the best checkpoint at the end of training rather than keeping the weights from the final epoch. +Evaluation Semantics +-------------------- + +``trainer.evaluate()`` accumulates predictions and labels across all batches +and computes every metric **exactly once on the pooled arrays** — metrics are +never computed per batch and then averaged. The reported ``loss`` is the +example-weighted mean over all samples (each batch's mean loss is weighted by +its batch size), so a smaller final batch does not skew it. + +As a result, evaluation scores are invariant to the evaluation batch size and +to the order in which batches arrive: running the same model on the same data +at ``batch_size=16`` or ``batch_size=64``, shuffled or not, yields the same +numbers. This matters for reproducible benchmarking — batch-averaged metrics +such as AUROC or AUPRC are not decomposable over batches, and averaging them +per batch silently produces batch-size-dependent (and thus non-reproducible) +results. + Getting the Test Scores ------------------------ diff --git a/examples/trainer_pooled_evaluation.py b/examples/trainer_pooled_evaluation.py new file mode 100644 index 000000000..5fea2624d --- /dev/null +++ b/examples/trainer_pooled_evaluation.py @@ -0,0 +1,94 @@ +"""Demonstrates that trainer.evaluate() is batch-invariant and reproducible. + +Context (https://github.com/sunlabuiuc/PyHealth/issues/859): metrics such as +AUROC and AUPRC are not decomposable over batches — computing them per batch +and averaging gives batch-size-dependent, non-reproducible numbers. PyHealth's +``Trainer.evaluate`` therefore pools predictions and labels across all batches +and computes each metric exactly once on the pooled arrays, and pools the loss +as an example-weighted mean. + +This script evaluates the same fixed model on the same synthetic data with +different batch sizes and shuffle orders and shows the scores are identical. + +Run with: python examples/trainer_pooled_evaluation.py +""" + +import numpy as np +import torch +from torch import nn +from torch.utils.data import DataLoader, Dataset + +from pyhealth.trainer import Trainer + + +class SyntheticBinaryDataset(Dataset): + def __init__(self, n_samples=100, n_features=4, seed=7): + rng = np.random.default_rng(seed) + self.x = rng.normal(size=(n_samples, n_features)).astype("float32") + self.weights = rng.normal(size=(n_features,)).astype("float32") + noise = rng.normal(scale=2.0, size=n_samples).astype("float32") + self.y = (self.x @ self.weights + noise > 0).astype("float32") + + def __len__(self): + return len(self.y) + + def __getitem__(self, index): + return {"x": self.x[index], "y": self.y[index]} + + +class DeterministicBinaryModel(nn.Module): + def __init__(self, weights): + super().__init__() + self.mode = "binary" + self.linear = nn.Linear(len(weights), 1) + with torch.no_grad(): + self.linear.weight.copy_(torch.from_numpy(weights).reshape(1, -1)) + self.linear.bias.zero_() + + def forward(self, x, y, **kwargs): + logits = self.linear(x).squeeze(-1) + y_true = y.float() + loss = nn.functional.binary_cross_entropy_with_logits(logits, y_true) + return {"loss": loss, "y_true": y_true, "y_prob": torch.sigmoid(logits)} + + +def main(): + dataset = SyntheticBinaryDataset() + trainer = Trainer( + model=DeterministicBinaryModel(dataset.weights), + metrics=["roc_auc", "pr_auc", "f1"], + device="cpu", + enable_logging=False, + ) + + configs = { + "batch_size=16, unshuffled": DataLoader(dataset, batch_size=16), + "batch_size=64, unshuffled": DataLoader(dataset, batch_size=64), + "batch_size=16, shuffled(seed=0)": DataLoader( + dataset, + batch_size=16, + shuffle=True, + generator=torch.Generator().manual_seed(0), + ), + "batch_size=16, shuffled(seed=1)": DataLoader( + dataset, + batch_size=16, + shuffle=True, + generator=torch.Generator().manual_seed(1), + ), + } + + results = {name: trainer.evaluate(loader) for name, loader in configs.items()} + for name, scores in results.items(): + printable = {k: round(v, 6) for k, v in scores.items()} + print(f"{name}: {printable}") + + reference = next(iter(results.values())) + for scores in results.values(): + for key in reference: + np.testing.assert_allclose(scores[key], reference[key], rtol=1e-6) + print("All configurations produced identical pooled metrics.") + + +if __name__ == "__main__": + main() diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py index bc6a28677..ba02d3d30 100644 --- a/pyhealth/trainer.py +++ b/pyhealth/trainer.py @@ -64,6 +64,22 @@ class Trainer: enable_logging: Whether to enable logging. Default is True. output_path: Path to save the output. Default is "./output". exp_name: Name of the experiment. Default is current datetime. + + Examples: + >>> from pyhealth.trainer import Trainer + >>> trainer = Trainer(model=model, metrics=["roc_auc", "pr_auc"]) + >>> trainer.train( + ... train_dataloader=train_loader, + ... val_dataloader=val_loader, + ... epochs=5, + ... monitor="roc_auc", + ... ) + >>> scores = trainer.evaluate(test_loader) + + ``evaluate`` pools predictions and labels across all batches and + computes each metric exactly once on the pooled arrays, and the + reported loss is the example-weighted mean, so the scores do not + depend on the batch size or on the order in which batches arrive. """ def __init__( @@ -268,6 +284,12 @@ def inference(self, dataloader, additional_outputs=None, return_patient_ids=False) -> Dict[str, float]: """Model inference. + Predictions and labels are accumulated across all batches and + returned as pooled arrays. The loss is pooled as an + example-weighted mean (each batch's mean loss is weighted by its + batch size), so the returned value does not depend on the batch + size or on the order in which batches arrive. + Args: dataloader: Dataloader for evaluation. additional_outputs: List of additional output to collect. @@ -276,11 +298,13 @@ def inference(self, dataloader, additional_outputs=None, Returns: y_true_all: List of true labels. y_prob_all: List of predicted probabilities. - loss_mean: Mean loss over batches. + loss_mean: Example-weighted mean loss over all samples. additional_outputs (only if requested): Dict of additional results. - patient_ids (only if requested): List of patient ids in the same order as y_true_all/y_prob_all. + patient_ids (only if requested): List of patient ids in the same + order as y_true_all/y_prob_all. """ - loss_all = [] + loss_sum = 0.0 + sample_count = 0 y_true_all = [] y_prob_all = [] patient_ids = [] @@ -293,7 +317,9 @@ def inference(self, dataloader, additional_outputs=None, loss = output["loss"] y_true = output["y_true"].cpu().numpy() y_prob = output["y_prob"].cpu().numpy() - loss_all.append(loss.item()) + batch_size = y_true.shape[0] if y_true.ndim > 0 else 1 + loss_sum += loss.item() * batch_size + sample_count += batch_size y_true_all.append(y_true) y_prob_all.append(y_prob) if additional_outputs is not None: @@ -301,7 +327,7 @@ def inference(self, dataloader, additional_outputs=None, additional_outputs[key].append(output[key].cpu().numpy()) if return_patient_ids: patient_ids.extend(data["patient_id"]) - loss_mean = sum(loss_all) / len(loss_all) + loss_mean = loss_sum / sample_count y_true_all = np.concatenate(y_true_all, axis=0) y_prob_all = np.concatenate(y_prob_all, axis=0) outputs = [y_true_all, y_prob_all, loss_mean] @@ -316,6 +342,13 @@ def inference(self, dataloader, additional_outputs=None, def evaluate(self, dataloader) -> Dict[str, float]: """Evaluates the model. + Predictions and labels are first pooled across all batches via + :meth:`inference`, and each metric is computed exactly once on the + pooled arrays (never computed per batch and averaged). The reported + loss is the example-weighted mean over all samples. As a result, the + scores are invariant to the evaluation batch size and to the order + in which batches arrive. + Args: dataloader: Dataloader for evaluation. @@ -329,17 +362,34 @@ def evaluate(self, dataloader) -> Dict[str, float]: scores = metrics_fn(y_true_all, y_prob_all, metrics=self.metrics) scores["loss"] = loss_mean else: - loss_all = [] + loss_sum = 0.0 + sample_count = 0 + self.model.eval() for data in tqdm(dataloader, desc="Evaluation"): - self.model.eval() with torch.no_grad(): output = self.model(**data) loss = output["loss"] - loss_all.append(loss.item()) - loss_mean = sum(loss_all) / len(loss_all) + batch_size = self._get_batch_size(output, data) + loss_sum += loss.item() * batch_size + sample_count += batch_size + loss_mean = loss_sum / sample_count scores = {"loss": loss_mean} return scores + @staticmethod + def _get_batch_size(output: dict, data: dict) -> int: + """Infers the number of samples in a batch for loss pooling.""" + for key in ("y_true", "y_prob"): + value = output.get(key) + if value is not None and hasattr(value, "shape") and len(value.shape): + return value.shape[0] + for value in data.values(): + if isinstance(value, torch.Tensor) and value.ndim > 0: + return value.shape[0] + if isinstance(value, (list, tuple)): + return len(value) + return 1 + def save_ckpt(self, ckpt_path: str) -> None: """Saves the model checkpoint.""" state_dict = self.model.state_dict() diff --git a/tests/core/test_trainer_pooled_eval.py b/tests/core/test_trainer_pooled_eval.py new file mode 100644 index 000000000..b9fb8013f --- /dev/null +++ b/tests/core/test_trainer_pooled_eval.py @@ -0,0 +1,179 @@ +"""Tests for pooled (batch-invariant) evaluation metrics in Trainer. + +Regression tests for https://github.com/sunlabuiuc/PyHealth/issues/859: +``trainer.evaluate`` must accumulate predictions and labels across batches +and compute every metric exactly once on the pooled arrays, and must pool +the loss as an example-weighted mean. Otherwise the reported scores depend +on the evaluation batch size and on the order in which batches arrive, so +repeated benchmark runs disagree even with a fixed seed. +""" + +import unittest + +import numpy as np +import torch +from sklearn.metrics import average_precision_score, roc_auc_score +from torch import nn +from torch.utils.data import DataLoader, Dataset + +from pyhealth.trainer import Trainer + + +class SyntheticBinaryDataset(Dataset): + """Small deterministic binary-classification dataset.""" + + def __init__(self, n_samples: int = 100, n_features: int = 4, seed: int = 7): + rng = np.random.default_rng(seed) + self.x = rng.normal(size=(n_samples, n_features)).astype("float32") + self.weights = rng.normal(size=(n_features,)).astype("float32") + noise = rng.normal(scale=2.0, size=n_samples).astype("float32") + self.y = (self.x @ self.weights + noise > 0).astype("float32") + + def __len__(self): + return len(self.y) + + def __getitem__(self, index): + return {"x": self.x[index], "y": self.y[index]} + + +class DeterministicBinaryModel(nn.Module): + """Fixed-weight logistic model exposing the PyHealth output contract.""" + + def __init__(self, weights: np.ndarray): + super().__init__() + self.mode = "binary" + self.linear = nn.Linear(len(weights), 1) + with torch.no_grad(): + self.linear.weight.copy_(torch.from_numpy(weights).reshape(1, -1)) + self.linear.bias.zero_() + + def forward(self, x, y, **kwargs): + logits = self.linear(x).squeeze(-1) + y_true = y.float() + loss = nn.functional.binary_cross_entropy_with_logits(logits, y_true) + return {"loss": loss, "y_true": y_true, "y_prob": torch.sigmoid(logits)} + + +class LossOnlyModel(nn.Module): + """Model without prediction outputs (exercises the mode=None path).""" + + def __init__(self, n_features: int = 4): + super().__init__() + self.mode = None + self.linear = nn.Linear(n_features, 1) + with torch.no_grad(): + self.linear.weight.fill_(0.1) + self.linear.bias.zero_() + + def forward(self, x, y, **kwargs): + logits = self.linear(x).squeeze(-1) + loss = nn.functional.binary_cross_entropy_with_logits(logits, y.float()) + return {"loss": loss} + + +METRICS = ["roc_auc", "pr_auc", "f1", "accuracy"] + + +class TestPooledEvaluation(unittest.TestCase): + """trainer.evaluate must be deterministic and batch-invariant.""" + + @classmethod + def setUpClass(cls): + cls.dataset = SyntheticBinaryDataset(n_samples=100) + cls.model = DeterministicBinaryModel(cls.dataset.weights) + cls.trainer = Trainer( + model=cls.model, + metrics=METRICS, + device="cpu", + enable_logging=False, + ) + + def _evaluate(self, batch_size, shuffle_seed=None): + if shuffle_seed is None: + loader = DataLoader(self.dataset, batch_size=batch_size) + else: + generator = torch.Generator().manual_seed(shuffle_seed) + loader = DataLoader( + self.dataset, + batch_size=batch_size, + shuffle=True, + generator=generator, + ) + return self.trainer.evaluate(loader) + + def test_determinism_same_seed(self): + """Two evaluate() calls with the same seed give identical metrics.""" + scores_1 = self._evaluate(batch_size=16, shuffle_seed=0) + scores_2 = self._evaluate(batch_size=16, shuffle_seed=0) + self.assertEqual(set(scores_1), set(scores_2)) + for key in scores_1: + self.assertEqual(scores_1[key], scores_2[key], msg=key) + + def test_batch_order_invariance(self): + """Metrics do not depend on the order in which batches arrive.""" + scores_1 = self._evaluate(batch_size=16, shuffle_seed=0) + scores_2 = self._evaluate(batch_size=16, shuffle_seed=1) + self.assertEqual(set(scores_1), set(scores_2)) + for key in scores_1: + np.testing.assert_allclose( + scores_1[key], scores_2[key], rtol=1e-6, err_msg=key + ) + + def test_batch_size_invariance(self): + """Same data at batch_size 16 vs 64 gives identical pooled metrics. + + 100 samples do not divide evenly into either batch size, so any + per-batch averaging would over-weight the partial final batch and + make the two runs disagree. + """ + scores_16 = self._evaluate(batch_size=16) + scores_64 = self._evaluate(batch_size=64) + self.assertEqual(set(scores_16), set(scores_64)) + for key in scores_16: + np.testing.assert_allclose( + scores_16[key], scores_64[key], rtol=1e-6, err_msg=key + ) + + def test_pooled_metrics_match_sklearn(self): + """Pooled AUROC/AUPRC match sklearn computed directly on all data.""" + with torch.no_grad(): + logits = self.model.linear(torch.from_numpy(self.dataset.x)) + y_prob = torch.sigmoid(logits.squeeze(-1)) + y_true = self.dataset.y + expected_roc_auc = roc_auc_score(y_true, y_prob.numpy()) + expected_pr_auc = average_precision_score(y_true, y_prob.numpy()) + expected_loss = nn.functional.binary_cross_entropy_with_logits( + logits.squeeze(-1), torch.from_numpy(y_true) + ).item() + + scores = self._evaluate(batch_size=16) + np.testing.assert_allclose(scores["roc_auc"], expected_roc_auc, rtol=1e-6) + np.testing.assert_allclose(scores["pr_auc"], expected_pr_auc, rtol=1e-6) + np.testing.assert_allclose(scores["loss"], expected_loss, rtol=1e-5) + + def test_loss_only_model_batch_size_invariance(self): + """The mode=None loss path is also example-weighted, not batch-averaged.""" + trainer = Trainer( + model=LossOnlyModel(), + device="cpu", + enable_logging=False, + ) + loss_16 = trainer.evaluate( + DataLoader(self.dataset, batch_size=16) + )["loss"] + loss_64 = trainer.evaluate( + DataLoader(self.dataset, batch_size=64) + )["loss"] + np.testing.assert_allclose(loss_16, loss_64, rtol=1e-6) + + def test_inference_returns_all_samples_in_order(self): + """inference() pools every sample exactly once, in dataloader order.""" + loader = DataLoader(self.dataset, batch_size=16) + y_true_all, y_prob_all, _ = self.trainer.inference(loader) + self.assertEqual(y_true_all.shape[0], len(self.dataset)) + self.assertEqual(y_prob_all.shape[0], len(self.dataset)) + np.testing.assert_array_equal(y_true_all, self.dataset.y) + + +if __name__ == "__main__": + unittest.main()