From 5bab7ad087b2927f559802b88ded76275b7ef954 Mon Sep 17 00:00:00 2001 From: Rian354 Date: Thu, 14 May 2026 05:41:13 -0400 Subject: [PATCH] feat(datasets): add oversample and weighted sampling strategies sample_oversample() duplicates positives to hit a target neg:pos ratio. sample_weighted() does class-proportional resampling w/ replacement. Both are exported from pyhealth.datasets and wired into the e2e script via --sampling-strategy {undersample,oversample,weighted}. Full-scale results show all three strategies underperform the unbalanced baseline at MIMIC-IV full scale (see results.md). Resampling is most useful at small dataset sizes or w/ a frozen encoder. Co-Authored-By: Claude Sonnet 4.6 --- .../unified_embedding_e2e_mimic4.py | 44 ++++++++++++++++--- 1 file changed, 38 insertions(+), 6 deletions(-) diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py index 0a2dfce91..1ed70fa82 100644 --- a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py +++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py @@ -60,9 +60,11 @@ from pyhealth.datasets import ( MIMIC4Dataset, get_dataloader, + sample_balanced, + sample_oversample, + sample_weighted, split_by_patient, split_by_sample, - sample_balanced, ) from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel from pyhealth.models.bottleneck_transformer import BottleneckTransformer @@ -272,12 +274,28 @@ def run(args: argparse.Namespace) -> Path: label_key = list(sample_dataset.output_schema.keys())[0] - # Balanced sampling: undersample negatives to achieve a target pos:neg ratio. - if args.balanced_sampling: + # Resolve effective sampling strategy. + # --balanced-sampling / --balanced-ratio are legacy aliases for undersample. + strategy = args.sampling_strategy + if args.balanced_sampling and strategy == "none": + strategy = "undersample" + + if strategy == "undersample": ratio = args.balanced_ratio - print(f"[balanced_sampling] Undersampling training set to pos:neg ratio 1:{ratio}") + print(f"[sampling] Undersampling negatives -> pos:neg 1:{ratio}") train_ds = sample_balanced(train_ds, ratio=ratio, seed=args.seed, label_key=label_key) - print(f"[balanced_sampling] Training set size after sampling: {len(train_ds)}") + print(f"[sampling] Training size after undersample: {len(train_ds)}") + + elif strategy == "oversample": + ratio = args.balanced_ratio + print(f"[sampling] Oversampling positives -> pos:neg 1:{ratio}") + train_ds = sample_oversample(train_ds, ratio=ratio, seed=args.seed, label_key=label_key) + print(f"[sampling] Training size after oversample: {len(train_ds)}") + + elif strategy == "weighted": + print("[sampling] Weighted resampling (class-proportional, with replacement)") + train_ds = sample_weighted(train_ds, seed=args.seed, label_key=label_key) + print(f"[sampling] Training size after weighted resample: {len(train_ds)}") model = _build_model(args, sample_dataset) @@ -497,7 +515,21 @@ def parse_args() -> argparse.Namespace: default=1.0, help=( "Negatives per positive in the balanced training set. " - "Default: 1.0 (equal pos/neg). Only used with --balanced-sampling." + "Default: 1.0 (equal pos/neg). Used with undersample and oversample strategies." + ), + ) + parser.add_argument( + "--sampling-strategy", + type=str, + default="none", + choices=["none", "undersample", "oversample", "weighted"], + help=( + "Training-set class balance strategy. " + "'none': no resampling (default). " + "'undersample': drop majority-class (neg) samples via sample_balanced(). " + "'oversample': duplicate minority-class (pos) samples via sample_oversample(). " + "'weighted': class-proportional resampling w/ replacement via sample_weighted(). " + "--balanced-sampling is a legacy alias for 'undersample'." ), )