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

Enhance training script to support all-future model target mode and update dataset handling

Browse Source
This commit is contained in:
Jiarui Li
2026-06-12 11:49:44 +08:00
parent 3125b6119f
commit 82f70945d9
2 changed files with 345 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
README.md
View File

@@ -221,11 +221,24 @@ all-future / query-conditioned 监督：
## 训练
当前 `train.py` 是 next-step 训练入口，使用：
当前 `train.py` 支持 next-token 和 all-future 两类训练入口：
```python
HealthDataset = NextStepHealthDataset
```
- `--model_target_mode next_token`
- 使用 `NextStepHealthDataset`
- `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
- `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
- `--model_target_mode all_future`
- 使用 `AllFutureHealthDataset`
- 不使用 readout，直接对 query hidden 计算风险
- `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
模型结构组合由 `model_target_mode × time_mode × dist_mode` 决定：
| 维度 | 可选项 |
| --- | --- |
| `model_target_mode` | `next_token`, `all_future` |
| `time_mode` | `relative`, `absolute` |
| `dist_mode` | `exponential`, `weibull`, `mixed` |
示例：
@@ -233,6 +246,7 @@ HealthDataset = NextStepHealthDataset
python train.py \
--data_prefix ukb \
--labels_file labels.csv \
--model_target_mode next_token \
--target_mode uts \
--n_embd 120 \
--n_head 10 \
@@ -240,6 +254,17 @@ python train.py \
--n_tab_layer 4
```
all-future 示例：
```bash
python train.py \
--data_prefix ukb \
--labels_file labels.csv \
--model_target_mode all_future \
--dist_mode weibull \
--time_mode relative
```
选择额外信息变量：
```bash
@@ -252,6 +277,7 @@ python train.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt
- `extra_info_types_file`：训练时使用的列表文件名
- `extra_info_types`：解析后的实际 type id 列表，用于评估脚本复现变量选择
- `model_target_mode`、`time_mode`、`dist_mode`、`dataset_class`、`collate_fn`、`resolved_loss_name`：用于评估脚本重建模型和输入方式
## 评估 AUC

418
train.py
View File

@@ -30,7 +30,12 @@ from torch.optim import AdamW
from torch.nn.utils import clip_grad_norm_
from tqdm.auto import tqdm
from dataset import HealthDataset, collate_fn
from dataset import (
AllFutureHealthDataset,
HealthDataset,
all_future_collate_fn,
collate_fn,
)
from models import DeepHealth
from readouts import build_readout
from losses import build_loss
@@ -225,6 +230,53 @@ def split_dataset(
)
def split_all_future_datasets(
train_dataset: AllFutureHealthDataset,
val_dataset: AllFutureHealthDataset,
test_dataset: AllFutureHealthDataset,
train_ratio: float,
val_ratio: float,
test_ratio: float,
seed: int,
) -> Tuple[Subset, Subset, Subset]:
"""Split all-future datasets by patient, then select validation/test queries."""
total = train_ratio + val_ratio + test_ratio
if not np.isclose(total, 1.0, atol=1e-6):
raise ValueError(
f"train_ratio + val_ratio + test_ratio must equal 1.0, got {total}"
)
n_patients = len(train_dataset.patients)
rng = np.random.RandomState(seed)
patient_indices = rng.permutation(n_patients)
n_train = int(n_patients * train_ratio)
n_val = int(n_patients * val_ratio)
train_patient_idx = patient_indices[:n_train]
val_patient_set = set(int(x) for x in patient_indices[n_train:n_train + n_val])
test_patient_set = set(int(x) for x in patient_indices[n_train + n_val:])
val_query_idx = [
i for i, (pidx, _t_query) in enumerate(val_dataset.valid_queries)
if int(pidx) in val_patient_set
]
test_query_idx = [
i for i, (pidx, _t_query) in enumerate(test_dataset.valid_queries)
if int(pidx) in test_patient_set
]
if not val_query_idx:
raise ValueError("All-future validation split has no valid query samples.")
if not test_query_idx:
raise ValueError("All-future test split has no valid query samples.")
return (
Subset(train_dataset, train_patient_idx),
Subset(val_dataset, np.asarray(val_query_idx, dtype=np.int64)),
Subset(test_dataset, np.asarray(test_query_idx, dtype=np.int64)),
)
def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
"""
Build DeepHealth model using metadata from dataset.
@@ -242,9 +294,9 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
n_categories=dataset.n_categories,
cont_type_ids=dataset.cont_type_ids,
n_bins=args.n_bins,
target_mode="next_token",
target_mode=args.model_target_mode,
time_mode=args.time_mode,
dist_mode="exponential",
dist_mode=args.dist_mode,
dropout=args.dropout,
)
@@ -302,7 +354,7 @@ def move_batch_to_device(batch: Dict, device: torch.device) -> Dict:
def compute_loss(
args: argparse.Namespace,
model: DeepHealth,
readout: nn.Module,
readout: nn.Module | None,
criterion: nn.Module,
batch: Dict[str, torch.Tensor],
device: torch.device,
@@ -340,56 +392,100 @@ def compute_loss(
padding_mask = batch["padding_mask"] # (B, L)
sex = batch["sex"] # (B,)
hidden = model(
event_seq=event_seq,
time_seq=time_seq,
sex=sex,
padding_mask=padding_mask,
other_type=batch["other_type"],
other_value=batch["other_value"],
other_value_kind=batch["other_value_kind"],
other_time=batch["other_time"],
target_mode="next_token",
)
# Apply readout
readout_mask = (
batch["readout_mask"]
if args.readout_name == "same_time_group_end"
else None
)
readout_out = readout(
hidden=hidden,
time_seq=time_seq,
padding_mask=padding_mask,
readout_mask=readout_mask,
)
# Compute risk logits
logits = model.calc_risk(readout_out.hidden)
# Compute loss based on target_mode
if args.target_mode == "delphi2m":
loss_out = criterion(
logits=logits,
target_events=batch["target_event_seq"],
target_times=batch["target_time_seq"],
current_times=batch["time_seq"],
padding_mask=readout_out.readout_mask,
return_components=True,
)
elif args.target_mode == "uts":
loss_out = criterion(
logits=logits,
target_multi_hot=batch["target_multi_hot"],
target_dt_unique=batch["target_dt_unique"],
readout_mask=readout_out.readout_mask,
return_components=True,
if args.model_target_mode == "all_future":
hidden = model(
event_seq=event_seq,
time_seq=time_seq,
sex=sex,
padding_mask=padding_mask,
t_query=batch["t_query"],
other_type=batch["other_type"],
other_value=batch["other_value"],
other_value_kind=batch["other_value_kind"],
other_time=batch["other_time"],
target_mode="all_future",
)
logits = model.calc_risk(hidden)
if args.dist_mode == "exponential":
loss = criterion(
logits=logits,
targets=batch["future_targets"],
exposure=batch["exposure"],
)
elif args.dist_mode == "weibull":
loss = criterion(
logits=logits,
weibull_rho=model.calc_weibull_rho(hidden),
targets=batch["future_targets"],
dt=batch["future_dt"],
exposure=batch["exposure"],
)
elif args.dist_mode == "mixed":
loss = criterion(
logits=logits,
death_rho=model.calc_death_rho(hidden),
targets=batch["future_targets"],
dt=batch["future_dt"],
exposure=batch["exposure"],
)
else:
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
loss_parts = {"total": loss.detach()}
else:
raise ValueError(f"Unknown target_mode: {args.target_mode}")
if readout is None:
raise ValueError("next_token training requires a readout module")
loss, loss_parts = loss_out
hidden = model(
event_seq=event_seq,
time_seq=time_seq,
sex=sex,
padding_mask=padding_mask,
other_type=batch["other_type"],
other_value=batch["other_value"],
other_value_kind=batch["other_value_kind"],
other_time=batch["other_time"],
target_mode="next_token",
)
# Apply readout
readout_mask = (
batch["readout_mask"]
if args.readout_name == "same_time_group_end"
else None
)
readout_out = readout(
hidden=hidden,
time_seq=time_seq,
padding_mask=padding_mask,
readout_mask=readout_mask,
)
# Compute risk logits
logits = model.calc_risk(readout_out.hidden)
# Compute loss based on target_mode
if args.target_mode == "delphi2m":
loss_out = criterion(
logits=logits,
target_events=batch["target_event_seq"],
target_times=batch["target_time_seq"],
current_times=batch["time_seq"],
padding_mask=readout_out.readout_mask,
return_components=True,
)
elif args.target_mode == "uts":
loss_out = criterion(
logits=logits,
target_multi_hot=batch["target_multi_hot"],
target_dt_unique=batch["target_dt_unique"],
readout_mask=readout_out.readout_mask,
return_components=True,
)
else:
raise ValueError(f"Unknown target_mode: {args.target_mode}")
loss, loss_parts = loss_out
# Check for NaN/Inf
if not torch.isfinite(loss):
@@ -406,7 +502,7 @@ def run_one_epoch(
logger: logging.Logger,
args: argparse.Namespace,
model: DeepHealth,
readout: nn.Module,
readout: nn.Module | None,
criterion: nn.Module,
train_loader: DataLoader,
optimizer: AdamW,
@@ -421,7 +517,7 @@ def run_one_epoch(
logger : logging.Logger
args : argparse.Namespace
model : DeepHealth
readout : nn.Module
readout : nn.Module or None
criterion : nn.Module
train_loader : DataLoader
optimizer : AdamW
@@ -522,7 +618,7 @@ def evaluate(
logger: logging.Logger,
args: argparse.Namespace,
model: DeepHealth,
readout: nn.Module,
readout: nn.Module | None,
criterion: nn.Module,
val_loader: DataLoader,
device: torch.device,
@@ -588,13 +684,25 @@ def build_run_metadata(
run_name: str,
) -> Dict[str, Any]:
"""Collect resolved training facts needed to rebuild the model for evaluation."""
dataset_class = (
"AllFutureHealthDataset"
if args.model_target_mode == "all_future"
else "NextStepHealthDataset"
)
collate_name = (
"all_future_collate_fn"
if args.model_target_mode == "all_future"
else "next_step_collate_fn"
)
return {
"run_name": run_name,
"dataset_class": "NextStepHealthDataset",
"collate_fn": "next_step_collate_fn",
"dataset_class": dataset_class,
"collate_fn": collate_name,
"model_class": "DeepHealth",
"model_target_mode": "next_token",
"dist_mode": "exponential",
"model_target_mode": args.model_target_mode,
"dist_mode": args.dist_mode,
"all_future_min_history_events": int(args.all_future_min_history_events),
"all_future_min_future_events": int(args.all_future_min_future_events),
"extra_info_types_file": (
Path(args.extra_info_types_file).name
if args.extra_info_types_file is not None
@@ -621,18 +729,46 @@ def build_run_metadata(
def normalize_training_config(args: argparse.Namespace) -> None:
"""Fill in and validate training options that depend on other flags."""
if args.target_mode not in {"delphi2m", "uts"}:
raise ValueError(f"Unknown target_mode: {args.target_mode}")
if args.model_target_mode not in {"next_token", "all_future"}:
raise ValueError(f"Unknown model_target_mode: {args.model_target_mode}")
if args.dist_mode not in {"exponential", "weibull", "mixed"}:
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
if args.all_future_min_history_events < 1:
raise ValueError("all_future_min_history_events must be >= 1")
if args.all_future_min_future_events < 1:
raise ValueError("all_future_min_future_events must be >= 1")
if args.model_target_mode == "all_future":
args.target_mode = "all_future"
# gap_5y is always enabled, so preserve NO_EVENT target behavior.
args.ignore_no_event_in_delphi2m = False
if args.target_mode == "uts":
if args.model_target_mode == "next_token" and args.target_mode == "uts":
args.include_no_event_in_uts_target = True
def normalize_loss_and_distribution_config(args: argparse.Namespace) -> None:
"""Validate and resolve loss/distribution options after auto-selection."""
if args.loss_name not in {"delphi2m", "uts"}:
next_token_losses = {"delphi2m", "uts"}
all_future_losses = {"exponential", "weibull", "mixed"}
if args.model_target_mode == "all_future":
if args.loss_name not in all_future_losses:
raise ValueError(
"all_future training requires loss_name to be one of "
"exponential, weibull, or mixed."
)
if args.loss_name != args.dist_mode:
raise ValueError(
"all_future loss_name must match dist_mode so risk scoring and "
f"training distribution stay aligned. Got loss_name={args.loss_name!r}, "
f"dist_mode={args.dist_mode!r}."
)
return
if args.target_mode not in {"delphi2m", "uts"}:
raise ValueError(f"Unknown target_mode: {args.target_mode}")
if args.loss_name not in next_token_losses:
raise ValueError(
"Unknown loss_name. Supported values: delphi2m, uts."
)
@@ -692,6 +828,12 @@ def main():
parser.add_argument("--time_mode", type=str, default="relative",
choices=["relative", "absolute"],
help="Time encoding mode for disease history")
parser.add_argument("--model_target_mode", type=str, default="next_token",
choices=["next_token", "all_future"],
help="Model forward/training mode")
parser.add_argument("--dist_mode", type=str, default="exponential",
choices=["exponential", "weibull", "mixed"],
help="Event-time distribution for model heads and all-future loss")
parser.add_argument("--dropout", type=float, default=0.0,
help="Dropout rate")
parser.add_argument("--extra_info_types_file", type=str, default=None,
@@ -700,16 +842,20 @@ def main():
# ---- Training Protocol ----
parser.add_argument("--target_mode", type=str, default="uts",
choices=["delphi2m", "uts"],
help="Target supervision mode")
help="Next-token supervision mode; ignored for all_future model_target_mode")
parser.add_argument("--readout_name", type=str, default=None,
help="Readout name (auto-selected if None)")
parser.add_argument("--readout_reduce", type=str, default="mean",
choices=["mean", "sum"],
help="Readout reduction for SameTimeGroupEndReadout")
parser.add_argument("--all_future_min_history_events", type=int, default=1,
help="Minimum historical events before an all-future query")
parser.add_argument("--all_future_min_future_events", type=int, default=1,
help="Minimum future events after an all-future query")
# ---- Loss ----
parser.add_argument("--loss_name", type=str, default=None,
help="Loss name (auto-selected if None): delphi2m, uts")
help="Loss name (auto-selected if None): delphi2m, uts, exponential, weibull, mixed")
parser.add_argument("--t_min", type=float, default=0.0027378507871321013,
help="Minimum time for loss (1/365.25)")
parser.add_argument("--max_exp_input", type=float, default=60.0,
@@ -758,11 +904,32 @@ def main():
configure_torch_for_training(device)
normalize_training_config(args)
# Auto-select readout if not specified.
if args.model_target_mode == "all_future":
args.readout_name = "none"
elif args.readout_name is None:
args.readout_name = (
"token" if args.target_mode == "delphi2m"
else "same_time_group_end"
)
# Auto-select loss if not specified.
if args.loss_name is None:
if args.model_target_mode == "all_future":
args.loss_name = args.dist_mode
else:
args.loss_name = (
"delphi2m" if args.target_mode == "delphi2m"
else "uts"
)
normalize_loss_and_distribution_config(args)
runs_root = Path("runs")
while True:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
run_name = (
f"{args.time_mode}_exponential_{args.target_mode}_"
f"{args.time_mode}_{args.dist_mode}_{args.model_target_mode}_{args.loss_name}_"
f"other_tokens_gap_5y_{timestamp}"
)
run_dir = runs_root / run_name
@@ -788,45 +955,72 @@ def main():
f"train_ratio + val_ratio + test_ratio must equal 1.0, got {total_ratio}"
)
# Auto-select readout if not specified
if args.readout_name is None:
args.readout_name = (
"token" if args.target_mode == "delphi2m"
else "same_time_group_end"
)
# Auto-select loss if not specified
if args.loss_name is None:
args.loss_name = (
"delphi2m" if args.target_mode == "delphi2m"
else "uts"
)
normalize_loss_and_distribution_config(args)
logger.info(f"Auto-selected readout: {args.readout_name}")
logger.info(f"Auto-selected loss: {args.loss_name}")
# ---- Load Dataset ----
logger.info("Loading dataset...")
dataset = HealthDataset(
data_prefix=args.data_prefix,
labels_file=args.labels_file,
no_event_interval_years=args.no_event_interval_years,
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
extra_info_types=args.extra_info_types,
)
logger.info(
f"Dataset loaded: {len(dataset)} samples, vocab_size={dataset.vocab_size}")
if args.model_target_mode == "all_future":
dataset = AllFutureHealthDataset(
data_prefix=args.data_prefix,
labels_file=args.labels_file,
split="train",
no_event_interval_years=args.no_event_interval_years,
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
min_history_events=args.all_future_min_history_events,
min_future_events=args.all_future_min_future_events,
extra_info_types=args.extra_info_types,
)
val_dataset = AllFutureHealthDataset(
data_prefix=args.data_prefix,
labels_file=args.labels_file,
split="valid",
no_event_interval_years=args.no_event_interval_years,
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
min_history_events=args.all_future_min_history_events,
min_future_events=args.all_future_min_future_events,
extra_info_types=args.extra_info_types,
)
test_dataset = AllFutureHealthDataset(
data_prefix=args.data_prefix,
labels_file=args.labels_file,
split="test",
no_event_interval_years=args.no_event_interval_years,
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
min_history_events=args.all_future_min_history_events,
min_future_events=args.all_future_min_future_events,
extra_info_types=args.extra_info_types,
)
train_subset, val_subset, test_subset = split_all_future_datasets(
train_dataset=dataset,
val_dataset=val_dataset,
test_dataset=test_dataset,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
active_collate_fn = all_future_collate_fn
else:
dataset = HealthDataset(
data_prefix=args.data_prefix,
labels_file=args.labels_file,
no_event_interval_years=args.no_event_interval_years,
include_no_event_in_uts_target=args.include_no_event_in_uts_target,
extra_info_types=args.extra_info_types,
)
train_subset, val_subset, test_subset = split_dataset(
dataset=dataset,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
active_collate_fn = collate_fn
logger.info(
f"Dataset loaded: {len(dataset)} base samples, vocab_size={dataset.vocab_size}")
# ---- Split Dataset ----
train_subset, val_subset, test_subset = split_dataset(
dataset=dataset,
train_ratio=args.train_ratio,
val_ratio=args.val_ratio,
test_ratio=args.test_ratio,
seed=args.seed,
)
logger.info(
f"Dataset split: train={len(train_subset)}, val={len(val_subset)}, test={len(test_subset)}"
)
@@ -837,7 +1031,7 @@ def main():
batch_size=args.batch_size,
sampler=RandomSampler(
train_subset, generator=torch.Generator().manual_seed(args.seed)),
collate_fn=collate_fn,
collate_fn=active_collate_fn,
num_workers=args.num_workers,
pin_memory=device.type == "cuda",
persistent_workers=args.num_workers > 0,
@@ -847,7 +1041,7 @@ def main():
val_subset,
batch_size=args.batch_size,
shuffle=False,
collate_fn=collate_fn,
collate_fn=active_collate_fn,
num_workers=args.num_workers,
pin_memory=device.type == "cuda",
persistent_workers=args.num_workers > 0,
@@ -857,7 +1051,7 @@ def main():
test_subset,
batch_size=args.batch_size,
shuffle=False,
collate_fn=collate_fn,
collate_fn=active_collate_fn,
num_workers=args.num_workers,
pin_memory=device.type == "cuda",
persistent_workers=args.num_workers > 0,
@@ -882,7 +1076,9 @@ def main():
f"Adaptive LR: {adaptive_lr:.6f} (base_lr * sqrt(batch_size/128))")
# ---- Build Readout ----
if args.readout_name == "token":
if args.model_target_mode == "all_future":
readout = None
elif args.readout_name == "token":
readout = build_readout("token")
elif args.readout_name == "same_time_group_end":
readout = build_readout("same_time_group_end",
@@ -894,7 +1090,23 @@ def main():
logger.info(f"Readout: {args.readout_name}")
# ---- Build Loss ----
if args.loss_name == "delphi2m":
if args.model_target_mode == "all_future":
ignored_idx = {PAD_IDX, CHECKUP_IDX}
if args.loss_name == "exponential":
criterion = build_loss("exponential", ignored_idx=ignored_idx)
elif args.loss_name == "weibull":
criterion = build_loss("weibull", ignored_idx=ignored_idx)
elif args.loss_name == "mixed":
criterion = build_loss(
"mixed",
death_idx=dataset.vocab_size - 1,
ignored_idx=ignored_idx,
)
else:
raise ValueError(f"Unknown all_future loss: {args.loss_name}")
logger.info(
f"Loss: {args.loss_name}, dist_mode={args.dist_mode}, ignored_idx={ignored_idx}")
elif args.loss_name == "delphi2m":
ignored_tokens = {PAD_IDX, CHECKUP_IDX}
if args.ignore_no_event_in_delphi2m:
ignored_tokens.add(NO_EVENT_IDX)