505 lines
18 KiB
Python
505 lines
18 KiB
Python
"""
|
|
Train DeepHealth with query-conditioned all-future supervision.
|
|
|
|
Training samples are patient-level. For each patient and each __getitem__ call,
|
|
AllFutureHealthDataset randomly samples a query time t_query, uses events at or
|
|
before t_query as history, and uses events after t_query as the future target set.
|
|
|
|
Validation/test samples are deterministic query points built from future event
|
|
times, then split by patient.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import logging
|
|
import math
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict
|
|
|
|
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 tqdm.auto import tqdm
|
|
|
|
from dataset import AllFutureHealthDataset, all_future_collate_fn
|
|
from losses import build_loss
|
|
from models import DeepHealth
|
|
from targets import CHECKUP_IDX, PAD_IDX
|
|
from train_util import (
|
|
configure_torch_for_training,
|
|
create_unique_run_dir,
|
|
format_extra_info_types,
|
|
load_extra_info_types_file,
|
|
resolve_device,
|
|
save_checkpoint,
|
|
save_config,
|
|
set_optimizer_lr,
|
|
set_seed,
|
|
setup_logging,
|
|
split_all_future_datasets,
|
|
split_all_future_datasets_by_eid_files,
|
|
)
|
|
|
|
|
|
MODEL_INPUT_KEYS = (
|
|
"event_seq",
|
|
"time_seq",
|
|
"sex",
|
|
"padding_mask",
|
|
"t_query",
|
|
"other_type",
|
|
"other_value",
|
|
"other_value_kind",
|
|
"other_time",
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Train DeepHealth with all-future supervision")
|
|
|
|
parser.add_argument("--data_prefix", type=str, default="ukb")
|
|
parser.add_argument("--labels_file", type=str, default="labels.csv")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--extra_info_types_file", type=str, default=None)
|
|
|
|
parser.add_argument("--train_ratio", type=float, default=0.7)
|
|
parser.add_argument("--val_ratio", type=float, default=0.15)
|
|
parser.add_argument("--test_ratio", type=float, default=0.15)
|
|
parser.add_argument("--train_eid_file", type=str, default="ukb_train_eid.csv")
|
|
parser.add_argument("--val_eid_file", type=str, default="ukb_val_eid.csv")
|
|
parser.add_argument("--test_eid_file", type=str, default="ukb_test_eid.csv")
|
|
parser.add_argument("--min_history_events", type=int, default=1)
|
|
parser.add_argument("--min_future_events", type=int, default=1)
|
|
parser.add_argument("--validation_query_seed", type=int, default=None)
|
|
|
|
parser.add_argument("--n_embd", type=int, default=120)
|
|
parser.add_argument("--n_head", type=int, default=10)
|
|
parser.add_argument("--n_hist_layer", type=int, default=12)
|
|
parser.add_argument("--n_tab_layer", type=int, default=4)
|
|
parser.add_argument("--n_bins", type=int, default=16)
|
|
parser.add_argument("--extra_pool_reduce", type=str, default="mean",
|
|
choices=["mean", "sum"])
|
|
parser.add_argument("--time_mode", type=str, default="relative",
|
|
choices=["relative", "absolute"])
|
|
parser.add_argument("--dist_mode", type=str, default="exponential",
|
|
choices=["exponential", "weibull", "mixed"])
|
|
parser.add_argument("--dropout", type=float, default=0.0)
|
|
|
|
parser.add_argument("--batch_size", type=int, default=128)
|
|
parser.add_argument("--base_lr", type=float, default=3e-4)
|
|
parser.add_argument("--weight_decay", type=float, default=0.1)
|
|
parser.add_argument("--betas", type=float, nargs=2, default=(0.9, 0.99))
|
|
parser.add_argument("--grad_clip", type=float, default=1.0)
|
|
parser.add_argument("--max_epochs", type=int, default=200)
|
|
parser.add_argument("--warmup_epochs", type=int, default=10)
|
|
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("--device", type=str, default="cuda")
|
|
parser.add_argument("--progress_interval", type=int, default=20)
|
|
|
|
args = parser.parse_args()
|
|
if args.min_history_events < 1:
|
|
raise ValueError("min_history_events must be >= 1")
|
|
if args.min_future_events < 1:
|
|
raise ValueError("min_future_events must be >= 1")
|
|
use_eid_split = all(
|
|
getattr(args, name)
|
|
for name in ("train_eid_file", "val_eid_file", "test_eid_file")
|
|
)
|
|
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.validation_query_seed is None:
|
|
args.validation_query_seed = int(args.seed)
|
|
args.extra_info_types = (
|
|
load_extra_info_types_file(args.extra_info_types_file)
|
|
if args.extra_info_types_file is not None
|
|
else None
|
|
)
|
|
return args
|
|
|
|
|
|
def get_lr(epoch: int, args: argparse.Namespace, adaptive_lr: float) -> float:
|
|
if epoch < args.warmup_epochs:
|
|
return adaptive_lr * (epoch + 1) / args.warmup_epochs
|
|
progress = (epoch - args.warmup_epochs) / max(1, args.max_epochs - args.warmup_epochs)
|
|
cosine = 0.5 * (1 + math.cos(math.pi * progress))
|
|
return adaptive_lr * (args.min_lr_ratio + cosine * (1 - args.min_lr_ratio))
|
|
|
|
|
|
def move_batch_to_device(batch: Dict[str, torch.Tensor], device: torch.device) -> Dict[str, torch.Tensor]:
|
|
non_blocking = device.type == "cuda"
|
|
return {
|
|
key: value.to(device, non_blocking=non_blocking)
|
|
if isinstance(value, torch.Tensor)
|
|
else value
|
|
for key, value in batch.items()
|
|
}
|
|
|
|
|
|
def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> DeepHealth:
|
|
return DeepHealth(
|
|
vocab_size=dataset.vocab_size,
|
|
n_embd=args.n_embd,
|
|
n_head=args.n_head,
|
|
n_hist_layer=args.n_hist_layer,
|
|
n_tab_layer=args.n_tab_layer,
|
|
n_types=dataset.n_types,
|
|
n_cont_types=dataset.n_cont_types,
|
|
n_categories=dataset.n_categories,
|
|
cont_type_ids=dataset.cont_type_ids,
|
|
n_bins=args.n_bins,
|
|
extra_pool_reduce=args.extra_pool_reduce,
|
|
target_mode="all_future",
|
|
time_mode=args.time_mode,
|
|
dist_mode=args.dist_mode,
|
|
dropout=args.dropout,
|
|
)
|
|
|
|
|
|
def build_criterion(args: argparse.Namespace, dataset: AllFutureHealthDataset):
|
|
ignored_idx = {PAD_IDX, CHECKUP_IDX}
|
|
if args.dist_mode == "exponential":
|
|
return build_loss("exponential", ignored_idx=ignored_idx)
|
|
if args.dist_mode == "weibull":
|
|
return build_loss("weibull", ignored_idx=ignored_idx)
|
|
if args.dist_mode == "mixed":
|
|
return build_loss(
|
|
"mixed",
|
|
death_idx=dataset.vocab_size - 1,
|
|
ignored_idx=ignored_idx,
|
|
)
|
|
raise ValueError(f"Unknown dist_mode: {args.dist_mode}")
|
|
|
|
|
|
def compute_all_future_loss(
|
|
args: argparse.Namespace,
|
|
model: DeepHealth,
|
|
criterion,
|
|
batch: Dict[str, torch.Tensor],
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
required_keys = set(MODEL_INPUT_KEYS)
|
|
required_keys.update(("future_targets", "exposure"))
|
|
if args.dist_mode in {"weibull", "mixed"}:
|
|
required_keys.add("future_dt")
|
|
batch = move_batch_to_device(
|
|
{key: batch[key] for key in required_keys},
|
|
device,
|
|
)
|
|
|
|
hidden = model(
|
|
event_seq=batch["event_seq"],
|
|
time_seq=batch["time_seq"],
|
|
sex=batch["sex"],
|
|
padding_mask=batch["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"],
|
|
)
|
|
else:
|
|
loss = criterion(
|
|
logits=logits,
|
|
death_rho=model.calc_death_rho(hidden),
|
|
targets=batch["future_targets"],
|
|
dt=batch["future_dt"],
|
|
exposure=batch["exposure"],
|
|
)
|
|
|
|
if not torch.isfinite(loss):
|
|
raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}")
|
|
return loss
|
|
|
|
|
|
def run_epoch(
|
|
logger: logging.Logger,
|
|
args: argparse.Namespace,
|
|
model: DeepHealth,
|
|
criterion,
|
|
loader: DataLoader,
|
|
optimizer: AdamW | None,
|
|
device: torch.device,
|
|
is_train: bool,
|
|
) -> float:
|
|
model.train(is_train)
|
|
total = torch.zeros((), device=device)
|
|
n_batches = 0
|
|
skipped = 0
|
|
desc = "train" if is_train else "val"
|
|
progress_interval = max(1, int(args.progress_interval))
|
|
|
|
progress = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)
|
|
for batch_idx, batch in enumerate(progress):
|
|
try:
|
|
loss = compute_all_future_loss(args, model, criterion, batch, device)
|
|
if is_train:
|
|
if optimizer is None:
|
|
raise ValueError("optimizer is required for training")
|
|
optimizer.zero_grad(set_to_none=True)
|
|
loss.backward()
|
|
if args.grad_clip > 0:
|
|
clip_grad_norm_(model.parameters(), args.grad_clip)
|
|
optimizer.step()
|
|
|
|
total = total + loss.detach()
|
|
n_batches += 1
|
|
if (batch_idx + 1) % progress_interval == 0:
|
|
avg = total / max(1, n_batches)
|
|
progress.set_postfix(
|
|
loss=f"{float(loss.detach().cpu()):.4f}",
|
|
avg=f"{float(avg.detach().cpu()):.4f}",
|
|
skipped=skipped,
|
|
)
|
|
except RuntimeError as exc:
|
|
if "Loss is not finite" not in str(exc):
|
|
raise
|
|
skipped += 1
|
|
logger.warning(f"Batch {batch_idx} skipped: {str(exc)[:120]}")
|
|
|
|
if skipped:
|
|
logger.info(f"Skipped {skipped} batches due to non-finite loss")
|
|
return float((total / max(1, n_batches)).detach().cpu()) if n_batches else float("inf")
|
|
|
|
|
|
def build_metadata(
|
|
args: argparse.Namespace,
|
|
dataset: AllFutureHealthDataset,
|
|
run_name: str,
|
|
train_subset,
|
|
val_subset,
|
|
test_subset,
|
|
) -> Dict[str, Any]:
|
|
return {
|
|
"run_name": run_name,
|
|
"dataset_class": "AllFutureHealthDataset",
|
|
"collate_fn": "all_future_collate_fn",
|
|
"model_class": "DeepHealth",
|
|
"model_target_mode": "all_future",
|
|
"target_mode": "all_future",
|
|
"dist_mode": args.dist_mode,
|
|
"all_future_min_history_events": int(args.min_history_events),
|
|
"all_future_min_future_events": int(args.min_future_events),
|
|
"all_future_validation_query_seed": int(args.validation_query_seed),
|
|
"extra_info_types_file": (
|
|
Path(args.extra_info_types_file).name
|
|
if args.extra_info_types_file is not None
|
|
else None
|
|
),
|
|
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
|
"dataset_metadata": {
|
|
"vocab_size": int(dataset.vocab_size),
|
|
"n_types": int(dataset.n_types),
|
|
"n_cont_types": int(dataset.n_cont_types),
|
|
"n_categories": int(dataset.n_categories),
|
|
"cont_type_ids": [int(x) for x in dataset.cont_type_ids],
|
|
"extra_info_types": [int(x) for x in dataset.extra_info_types],
|
|
},
|
|
"split_sizes": {
|
|
"train": int(len(train_subset)),
|
|
"val": int(len(val_subset)),
|
|
"test": int(len(test_subset)),
|
|
},
|
|
"resolved_readout_name": "none",
|
|
"resolved_loss_name": args.dist_mode,
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
set_seed(args.seed)
|
|
device = resolve_device(args.device)
|
|
configure_torch_for_training(device)
|
|
|
|
run_dir, run_name = create_unique_run_dir(
|
|
lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}"
|
|
)
|
|
logger = setup_logging(run_dir)
|
|
|
|
logger.info(f"Starting all-future training run: {run_name}")
|
|
logger.info(f"Device: {device}")
|
|
logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
|
|
|
|
logger.info("Loading all-future datasets...")
|
|
train_dataset = AllFutureHealthDataset(
|
|
data_prefix=args.data_prefix,
|
|
labels_file=args.labels_file,
|
|
split="train",
|
|
min_history_events=args.min_history_events,
|
|
min_future_events=args.min_future_events,
|
|
validation_query_seed=args.validation_query_seed,
|
|
extra_info_types=args.extra_info_types,
|
|
)
|
|
val_dataset = AllFutureHealthDataset(
|
|
data_prefix=args.data_prefix,
|
|
labels_file=args.labels_file,
|
|
split="valid",
|
|
min_history_events=args.min_history_events,
|
|
min_future_events=args.min_future_events,
|
|
validation_query_seed=args.validation_query_seed,
|
|
extra_info_types=args.extra_info_types,
|
|
)
|
|
test_dataset = AllFutureHealthDataset(
|
|
data_prefix=args.data_prefix,
|
|
labels_file=args.labels_file,
|
|
split="test",
|
|
min_history_events=args.min_history_events,
|
|
min_future_events=args.min_future_events,
|
|
validation_query_seed=args.validation_query_seed,
|
|
extra_info_types=args.extra_info_types,
|
|
)
|
|
if args.train_eid_file and args.val_eid_file and args.test_eid_file:
|
|
train_subset, val_subset, test_subset = split_all_future_datasets_by_eid_files(
|
|
train_dataset=train_dataset,
|
|
val_dataset=val_dataset,
|
|
test_dataset=test_dataset,
|
|
train_eid_file=args.train_eid_file,
|
|
val_eid_file=args.val_eid_file,
|
|
test_eid_file=args.test_eid_file,
|
|
)
|
|
logger.info(
|
|
"Using eid split files: "
|
|
f"train={args.train_eid_file}, val={args.val_eid_file}, test={args.test_eid_file}"
|
|
)
|
|
else:
|
|
train_subset, val_subset, test_subset = split_all_future_datasets(
|
|
train_dataset=train_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,
|
|
)
|
|
logger.info(
|
|
f"Using random ratio split: train={args.train_ratio}, "
|
|
f"val={args.val_ratio}, test={args.test_ratio}, seed={args.seed}"
|
|
)
|
|
logger.info(
|
|
f"Patients/queries: 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=all_future_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,
|
|
)
|
|
val_loader = DataLoader(
|
|
val_subset,
|
|
batch_size=args.batch_size,
|
|
shuffle=False,
|
|
collate_fn=all_future_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,
|
|
)
|
|
test_loader = DataLoader(
|
|
test_subset,
|
|
batch_size=args.batch_size,
|
|
shuffle=False,
|
|
collate_fn=all_future_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,
|
|
)
|
|
|
|
model = build_model(args, train_dataset).to(device)
|
|
optimizer = AdamW(
|
|
model.parameters(),
|
|
lr=args.base_lr,
|
|
betas=tuple(args.betas),
|
|
weight_decay=args.weight_decay,
|
|
)
|
|
criterion = build_criterion(args, train_dataset)
|
|
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
|
|
|
|
save_config(
|
|
args,
|
|
run_dir / "train_config.json",
|
|
extra=build_metadata(args, train_dataset, run_name, train_subset, val_subset, test_subset),
|
|
)
|
|
|
|
best_val = float("inf")
|
|
patience = 0
|
|
history = []
|
|
best_model_path = run_dir / "best_model.pt"
|
|
start = time.time()
|
|
|
|
for epoch in range(args.max_epochs):
|
|
lr = get_lr(epoch, args, adaptive_lr)
|
|
set_optimizer_lr(optimizer, lr)
|
|
|
|
train_loss = run_epoch(logger, args, model, criterion, train_loader, optimizer, device, True)
|
|
with torch.no_grad():
|
|
val_loss = run_epoch(logger, args, model, criterion, val_loader, None, device, False)
|
|
|
|
is_best = val_loss < best_val
|
|
if is_best:
|
|
best_val = val_loss
|
|
patience = 0
|
|
save_checkpoint(model, best_model_path)
|
|
else:
|
|
patience += 1
|
|
|
|
logger.info(
|
|
f"Epoch {epoch + 1}/{args.max_epochs} | lr={lr:.6f} | "
|
|
f"train_loss={train_loss:.6f} | val_loss={val_loss:.6f} | "
|
|
f"best_val_loss={best_val:.6f} | patience={patience}/{args.patience} | "
|
|
f"elapsed={time.time() - start:.1f}s"
|
|
)
|
|
history.append({
|
|
"epoch": epoch + 1,
|
|
"lr": lr,
|
|
"train_loss": train_loss,
|
|
"val_loss": val_loss,
|
|
"best_val_loss": best_val,
|
|
"is_best": int(is_best),
|
|
})
|
|
if patience >= args.patience:
|
|
logger.info(f"Early stopping triggered at epoch {epoch + 1}")
|
|
break
|
|
|
|
with (run_dir / "history.json").open("w", encoding="utf-8") as f:
|
|
json.dump(history, f, indent=2)
|
|
|
|
logger.info("Evaluating best model on all-future test queries...")
|
|
model.load_state_dict(torch.load(best_model_path, map_location=device))
|
|
with torch.no_grad():
|
|
test_loss = run_epoch(logger, args, model, criterion, test_loader, None, device, False)
|
|
logger.info(f"Test loss: {test_loss:.6f}")
|
|
logger.info(f"Best checkpoint: {best_model_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|