fix: compute evaluation metrics once on pooled predictions, not per batch (#859)

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
 ------------------------
 

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()

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,15 +317,17 @@ 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:
                     for key in additional_outputs.keys():
                         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()