JiaruiLi/DeepHealthExpo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Optimize exposure index training IO

Browse Source
This commit is contained in:
Jiarui Li
2026-07-08 11:10:56 +08:00
parent 1288087959
commit 2388d81678
2 changed files with 222 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
train_next_step.py
View File

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