From 2388d816785e2999c7a66ecf182ca200bf653ca5 Mon Sep 17 00:00:00 2001 From: Jiarui Li Date: Wed, 8 Jul 2026 11:10:56 +0800 Subject: [PATCH] Optimize exposure index training IO --- dataset.py | 110 ++++++++++++++++++++++++++------ train_next_step.py | 152 ++++++++++++++++++++++++++++++++++++++------- 2 files changed, 222 insertions(+), 40 deletions(-) diff --git a/dataset.py b/dataset.py index f8a33e1..84cc6ef 100644 --- a/dataset.py +++ b/dataset.py @@ -45,7 +45,7 @@ def _monthly_exposure_columns() -> list[str]: class ExposureCache: """Random-access view over files produced by prepare_exposure_cache.py.""" - def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 16): + def __init__(self, cache_dir: str | Path, row_group_cache_size: int = 4): cache_dir = Path(cache_dir) self.cache_dir = cache_dir manifest_path = cache_dir / "exposure_manifest.json" @@ -129,6 +129,21 @@ class ExposureCache: self._key_to_index: dict[tuple[int, int, int], int] | None = None + def locality_key(self, indices: np.ndarray) -> tuple[int, int]: + """Return a stable parquet locality key for sampler-side batching.""" + indices = np.asarray(indices, dtype=np.int64) + valid = indices[indices >= 0] + if len(valid) == 0: + return (2**31 - 1, 2**31 - 1) + if self.storage != "parquet_index": + return (0, int(valid[0] // 1024)) + file_ids = np.asarray(self.daily_file_ids[valid], dtype=np.int64) + row_groups = np.asarray(self.daily_row_groups[valid], dtype=np.int64) + groups = np.stack([file_ids, row_groups], axis=1) + unique_groups, counts = np.unique(groups, axis=0, return_counts=True) + best = unique_groups[int(np.argmax(counts))] + return (int(best[0]), int(best[1])) + def build_age_index(self, birth_date_by_eid: dict[int, np.datetime64]) -> None: keys: dict[tuple[int, int, int], int] = {} eids = np.asarray(self.eids, dtype=np.int64) @@ -159,18 +174,67 @@ class ExposureCache: return out def daily_window(self, index: int) -> np.ndarray: - if index < 0: - return np.full(DAILY_EXPOSURE_SHAPE, np.nan, dtype=np.float32) - if self.storage == "dense_npy": - return np.asarray(self.daily[index], dtype=np.float32) - return self._parquet_window("daily", index) + return self.daily_windows(np.asarray([index], dtype=np.int64))[0] def monthly_window(self, index: int) -> np.ndarray: - if index < 0: - return np.full(MONTHLY_EXPOSURE_SHAPE, np.nan, dtype=np.float32) + return self.monthly_windows(np.asarray([index], dtype=np.int64))[0] + + def daily_windows(self, indices: np.ndarray) -> np.ndarray: + return self._windows("daily", indices) + + def monthly_windows(self, indices: np.ndarray) -> np.ndarray: + return self._windows("monthly", indices) + + def _windows( + self, + kind: Literal["daily", "monthly"], + indices: np.ndarray, + ) -> np.ndarray: + indices = np.asarray(indices, dtype=np.int64) + shape = DAILY_EXPOSURE_SHAPE if kind == "daily" else MONTHLY_EXPOSURE_SHAPE + out = np.full((len(indices), shape[0], shape[1]), np.nan, dtype=np.float32) + valid_pos = np.nonzero(indices >= 0)[0] + if len(valid_pos) == 0: + return out + + valid_indices = indices[valid_pos] if self.storage == "dense_npy": - return np.asarray(self.monthly[index], dtype=np.float32) - return self._parquet_window("monthly", index) + source = self.daily if kind == "daily" else self.monthly + out[valid_pos] = np.asarray(source[valid_indices], dtype=np.float32) + return out + + if kind == "daily": + file_ids = np.asarray(self.daily_file_ids[valid_indices], dtype=np.int64) + row_groups = np.asarray(self.daily_row_groups[valid_indices], dtype=np.int64) + row_in_groups = np.asarray(self.daily_row_in_groups[valid_indices], dtype=np.int64) + columns = _daily_exposure_columns() + else: + file_ids = np.asarray(self.monthly_file_ids[valid_indices], dtype=np.int64) + row_groups = np.asarray(self.monthly_row_groups[valid_indices], dtype=np.int64) + row_in_groups = np.asarray( + self.monthly_row_in_groups[valid_indices], + dtype=np.int64, + ) + columns = _monthly_exposure_columns() + + group_keys = np.stack([file_ids, row_groups], axis=1) + unique_groups, inverse = np.unique(group_keys, axis=0, return_inverse=True) + for group_idx, (file_id, row_group) in enumerate(unique_groups): + group_pos = np.nonzero(inverse == group_idx)[0] + frame = self._read_parquet_row_group( + kind, + int(file_id), + int(row_group), + columns, + ) + row_values = frame.iloc[row_in_groups[group_pos]].reindex(columns=columns) + values = ( + row_values.to_numpy(dtype=np.float32, copy=True) + .reshape(len(group_pos), shape[1], shape[0]) + .transpose(0, 2, 1) + ) + out[valid_pos[group_pos]] = values + return out def _parquet_window(self, kind: Literal["daily", "monthly"], index: int) -> np.ndarray: if kind == "daily": @@ -308,13 +372,17 @@ class _ExpoBaseDataset(Dataset): include_no_event_in_uts_target: bool = False, exposure_cache_dir: str | Path | None = None, mask_onset_exposure: bool = False, + exposure_row_group_cache_size: int = 4, ) -> None: self.data_prefix = data_prefix self.labels_file = labels_file self.no_event_interval_years = float(no_event_interval_years) self.include_no_event_in_uts_target = bool(include_no_event_in_uts_target) self.exposure_cache = ( - ExposureCache(exposure_cache_dir) + ExposureCache( + exposure_cache_dir, + row_group_cache_size=exposure_row_group_cache_size, + ) if exposure_cache_dir is not None else None ) @@ -450,14 +518,14 @@ class _ExpoBaseDataset(Dataset): if self.exposure_cache is None: raise RuntimeError("Exposure cache is not enabled") - daily = np.stack( - [self.exposure_cache.daily_window(int(idx)) for idx in exposure_index], - axis=0, - ).astype(np.float32, copy=True) - monthly = np.stack( - [self.exposure_cache.monthly_window(int(idx)) for idx in exposure_index], - axis=0, - ).astype(np.float32, copy=True) + daily = self.exposure_cache.daily_windows(exposure_index).astype( + np.float32, + copy=False, + ) + monthly = self.exposure_cache.monthly_windows(exposure_index).astype( + np.float32, + copy=False, + ) if self.mask_onset_exposure: daily[:, 0, :] = np.nan @@ -485,6 +553,7 @@ class NextStepHealthDataset(_ExpoBaseDataset): include_no_event_in_uts_target: bool = False, exposure_cache_dir: str | Path | None = None, mask_onset_exposure: bool = False, + exposure_row_group_cache_size: int = 4, ) -> None: super().__init__( data_prefix=data_prefix, @@ -493,6 +562,7 @@ class NextStepHealthDataset(_ExpoBaseDataset): include_no_event_in_uts_target=include_no_event_in_uts_target, exposure_cache_dir=exposure_cache_dir, mask_onset_exposure=mask_onset_exposure, + exposure_row_group_cache_size=exposure_row_group_cache_size, ) self.samples: List[Dict] = [] @@ -581,6 +651,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): validation_query_seed: int = 42, exposure_cache_dir: str | Path | None = None, mask_onset_exposure: bool = False, + exposure_row_group_cache_size: int = 4, ) -> None: if split not in {"train", "valid", "test"}: raise ValueError(f"split must be train/valid/test, got {split!r}") @@ -592,6 +663,7 @@ class AllFutureHealthDataset(_ExpoBaseDataset): include_no_event_in_uts_target=include_no_event_in_uts_target, exposure_cache_dir=exposure_cache_dir, mask_onset_exposure=mask_onset_exposure, + exposure_row_group_cache_size=exposure_row_group_cache_size, ) self.split = split diff --git a/train_next_step.py b/train_next_step.py index b19570a..b5c9b5a 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -12,13 +12,13 @@ import json import logging import math import time -from typing import Any, Dict +from typing import Any, Dict, Iterator, List import numpy as np import torch from torch.nn.utils import clip_grad_norm_ from torch.optim import AdamW -from torch.utils.data import DataLoader, RandomSampler +from torch.utils.data import DataLoader, RandomSampler, Sampler from tqdm.auto import tqdm from dataset import HealthDataset, collate_fn @@ -53,6 +53,59 @@ EXPOSURE_INPUT_KEYS = ( ) +class ExposureLocalityBatchSampler(Sampler[List[int]]): + """Randomized batches with within-buffer sorting by exposure parquet locality.""" + + def __init__( + self, + data_source, + batch_size: int, + buffer_size: int, + seed: int, + drop_last: bool = False, + ) -> None: + self.data_source = data_source + self.batch_size = int(batch_size) + self.buffer_size = max(int(buffer_size), self.batch_size) + self.seed = int(seed) + self.drop_last = bool(drop_last) + self.epoch = 0 + + def __iter__(self) -> Iterator[List[int]]: + n = len(self.data_source) + generator = torch.Generator().manual_seed(self.seed + self.epoch) + self.epoch += 1 + shuffled = torch.randperm(n, generator=generator).tolist() + for start in range(0, n, self.buffer_size): + buffer = shuffled[start:start + self.buffer_size] + buffer.sort(key=self._locality_key) + for batch_start in range(0, len(buffer), self.batch_size): + batch = buffer[batch_start:batch_start + self.batch_size] + if len(batch) < self.batch_size and self.drop_last: + continue + yield batch + + def __len__(self) -> int: + n = len(self.data_source) + if self.drop_last: + return n // self.batch_size + return (n + self.batch_size - 1) // self.batch_size + + def _locality_key(self, local_idx: int) -> tuple[int, int, int]: + dataset = self.data_source + raw_idx = int(local_idx) + if hasattr(dataset, "dataset") and hasattr(dataset, "indices"): + raw_idx = int(dataset.indices[local_idx]) + dataset = dataset.dataset + sample = getattr(dataset, "samples", [])[raw_idx] + exposure_index = sample.get("exposure_index") + exposure_cache = getattr(dataset, "exposure_cache", None) + if exposure_index is None or exposure_cache is None: + return (2**31 - 1, 2**31 - 1, raw_idx) + file_id, row_group = exposure_cache.locality_key(exposure_index) + return (file_id, row_group, raw_idx) + + def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train DeepHealth with next-token/point supervision") @@ -83,6 +136,15 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--exposure_conv_kernel_size", type=int, default=7) parser.add_argument("--exposure_mlp_ratio", type=float, default=4.0) parser.add_argument("--no_exposure_gate", action="store_true") + parser.add_argument( + "--exposure_row_group_cache_size", + type=int, + default=4, + help=( + "Number of parquet exposure row groups cached per DataLoader worker " + "when using indexed exposure storage." + ), + ) parser.add_argument("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) @@ -106,6 +168,21 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--patience", type=int, default=15) parser.add_argument("--min_lr_ratio", type=float, default=0.1) parser.add_argument("--num_workers", type=int, default=4) + parser.add_argument( + "--prefetch_factor", + type=int, + default=4, + help="DataLoader batches prefetched per worker when num_workers > 0.", + ) + parser.add_argument( + "--exposure_locality_buffer_size", + type=int, + default=4096, + help=( + "Training-only shuffle buffer sorted by exposure parquet locality. " + "Set 0 to use the standard RandomSampler." + ), + ) parser.add_argument("--device", type=str, default="cuda") parser.add_argument("--progress_interval", type=int, default=20) @@ -116,6 +193,12 @@ def parse_args() -> argparse.Namespace: ) if not use_eid_split and not np.isclose(args.train_ratio + args.val_ratio + args.test_ratio, 1.0): raise ValueError("train_ratio + val_ratio + test_ratio must equal 1.0") + if args.num_workers > 0 and args.prefetch_factor <= 0: + raise ValueError("prefetch_factor must be positive when num_workers > 0") + if args.exposure_row_group_cache_size < 0: + raise ValueError("exposure_row_group_cache_size must be non-negative") + if args.exposure_locality_buffer_size < 0: + raise ValueError("exposure_locality_buffer_size must be non-negative") if args.target_mode == "uts": args.readout_name = args.readout_name or "same_time_group_end" args.include_no_event_in_uts_target = True @@ -364,6 +447,10 @@ def build_metadata( "exposure_conv_kernel_size": int(args.exposure_conv_kernel_size), "exposure_mlp_ratio": float(args.exposure_mlp_ratio), "exposure_use_gate": not bool(args.no_exposure_gate), + "exposure_row_group_cache_size": int(args.exposure_row_group_cache_size), + "num_workers": int(args.num_workers), + "prefetch_factor": int(args.prefetch_factor), + "exposure_locality_buffer_size": int(args.exposure_locality_buffer_size), "split_sizes": { "train": int(len(train_subset)), "val": int(len(val_subset)), @@ -394,6 +481,13 @@ def main() -> None: logger.info(f"Device: {device}") logger.info(f"readout={args.readout_name}, target_mode={args.target_mode}") logger.info(f"exposure_cache_dir={args.exposure_cache_dir}") + logger.info( + "DataLoader IO: " + f"num_workers={args.num_workers}, " + f"prefetch_factor={args.prefetch_factor if args.num_workers > 0 else None}, " + f"exposure_row_group_cache_size={args.exposure_row_group_cache_size}, " + f"exposure_locality_buffer_size={args.exposure_locality_buffer_size}" + ) dataset = HealthDataset( data_prefix=args.data_prefix, @@ -402,6 +496,7 @@ def main() -> None: include_no_event_in_uts_target=args.include_no_event_in_uts_target, exposure_cache_dir=args.exposure_cache_dir, mask_onset_exposure=args.mask_onset_exposure, + exposure_row_group_cache_size=args.exposure_row_group_cache_size, ) if args.train_eid_file and args.val_eid_file and args.test_eid_file: train_subset, val_subset, test_subset = split_dataset_by_eid_files( @@ -430,35 +525,50 @@ def main() -> None: f"Samples: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}" ) - train_loader = DataLoader( - train_subset, - batch_size=args.batch_size, - sampler=RandomSampler(train_subset, generator=torch.Generator().manual_seed(args.seed)), - collate_fn=collate_fn, - num_workers=args.num_workers, - pin_memory=device.type == "cuda", - persistent_workers=args.num_workers > 0, - prefetch_factor=2 if args.num_workers > 0 else None, + dataloader_kwargs = { + "collate_fn": collate_fn, + "num_workers": args.num_workers, + "pin_memory": device.type == "cuda", + "persistent_workers": args.num_workers > 0, + "prefetch_factor": args.prefetch_factor if args.num_workers > 0 else None, + } + use_locality_sampler = ( + args.exposure_cache_dir is not None + and args.exposure_locality_buffer_size > 0 ) + if use_locality_sampler: + logger.info("Using exposure-locality batch sampler for training") + train_loader = DataLoader( + train_subset, + batch_sampler=ExposureLocalityBatchSampler( + train_subset, + batch_size=args.batch_size, + buffer_size=args.exposure_locality_buffer_size, + seed=args.seed, + ), + **dataloader_kwargs, + ) + else: + train_loader = DataLoader( + train_subset, + batch_size=args.batch_size, + sampler=RandomSampler( + train_subset, + generator=torch.Generator().manual_seed(args.seed), + ), + **dataloader_kwargs, + ) val_loader = DataLoader( val_subset, batch_size=args.batch_size, shuffle=False, - collate_fn=collate_fn, - num_workers=args.num_workers, - pin_memory=device.type == "cuda", - persistent_workers=args.num_workers > 0, - prefetch_factor=2 if args.num_workers > 0 else None, + **dataloader_kwargs, ) test_loader = DataLoader( test_subset, batch_size=args.batch_size, shuffle=False, - collate_fn=collate_fn, - num_workers=args.num_workers, - pin_memory=device.type == "cuda", - persistent_workers=args.num_workers > 0, - prefetch_factor=2 if args.num_workers > 0 else None, + **dataloader_kwargs, ) model = build_model(args, dataset).to(device)