""" Train DeepHealth with next-token / next-time-point supervision. The next-step dataset uses observed event histories, including CHECKUP state tokens, plus optional gap imputation. UTS training reads out only same-time group ends. """ from __future__ import annotations import argparse import json import logging import math import time 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 HealthDataset, collate_fn from losses import build_loss from models import DeepHealth, DeepHealthOutput from readouts import build_readout from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX from train_util import ( configure_torch_for_training, create_unique_run_dir, resolve_device, save_checkpoint, save_config, set_optimizer_lr, set_seed, setup_logging, split_dataset, split_dataset_by_eid_files, ) MODEL_INPUT_KEYS = ( "event_seq", "time_seq", "sex", "padding_mask", ) EXPOSURE_INPUT_KEYS = ( "exposure_daily", "exposure_monthly", ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train DeepHealth with next-token/point 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("--no_event_interval_years", type=float, default=5.0) parser.add_argument("--include_no_event_in_uts_target", action="store_true") 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("--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("--dropout", type=float, default=0.0) parser.add_argument("--exposure_cache_dir", type=str, default=None) parser.add_argument("--mask_onset_exposure", action="store_true") parser.add_argument("--exposure_d_model", type=int, default=None) parser.add_argument("--exposure_n_layers", type=int, default=2) parser.add_argument("--exposure_top_k", type=int, default=3) parser.add_argument("--exposure_n_convnext_blocks", type=int, default=2) 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("--target_mode", type=str, default="uts", choices=["delphi2m", "uts"]) parser.add_argument("--readout_name", type=str, default=None, choices=["token", "same_time_group_end", "last_valid"]) parser.add_argument("--readout_reduce", type=str, default="mean", choices=["mean", "sum"]) parser.add_argument("--t_min", type=float, default=0.0027378507871321013) parser.add_argument("--max_exp_input", type=float, default=60.0) parser.add_argument("--ce_weight", type=float, default=1.0) parser.add_argument("--time_weight", type=float, default=1.0) parser.add_argument("--ignore_no_event_in_delphi2m", action="store_true") 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() 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.target_mode == "uts": args.readout_name = args.readout_name or "same_time_group_end" args.include_no_event_in_uts_target = True else: args.readout_name = args.readout_name or "token" 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: HealthDataset) -> 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, target_mode="next_token", dist_mode="exponential", dropout=args.dropout, use_exposure_encoder=args.exposure_cache_dir is not None, exposure_d_model=args.exposure_d_model, exposure_n_layers=args.exposure_n_layers, exposure_top_k=args.exposure_top_k, exposure_n_convnext_blocks=args.exposure_n_convnext_blocks, exposure_conv_kernel_size=args.exposure_conv_kernel_size, exposure_mlp_ratio=args.exposure_mlp_ratio, exposure_use_gate=not args.no_exposure_gate, ) def build_next_step_readout(args: argparse.Namespace): if args.readout_name == "same_time_group_end": return build_readout("same_time_group_end", reduce=args.readout_reduce) return build_readout(args.readout_name) def build_next_step_loss(args: argparse.Namespace): if args.target_mode == "delphi2m": ignored_tokens = {PAD_IDX, CHECKUP_IDX} if args.ignore_no_event_in_delphi2m: ignored_tokens.add(NO_EVENT_IDX) return build_loss( "delphi2m", ignored_tokens=ignored_tokens, t_min=args.t_min, max_exp_input=args.max_exp_input, ce_weight=args.ce_weight, time_weight=args.time_weight, ) return build_loss( "uts", ignored_idx={PAD_IDX, CHECKUP_IDX}, t_min=args.t_min, max_exp_input=args.max_exp_input, ) def build_augmented_next_step_targets( batch_cpu: Dict[str, torch.Tensor], model_out: DeepHealthOutput, include_uts_targets: bool, ) -> Dict[str, torch.Tensor]: device = model_out.hidden.device non_blocking = device.type == "cuda" targets = { "target_event_seq": batch_cpu["target_event_seq"].to(device, non_blocking=non_blocking), "target_time_seq": batch_cpu["target_time_seq"].to(device, non_blocking=non_blocking), "readout_mask": batch_cpu["readout_mask"].to(device, non_blocking=non_blocking), } if include_uts_targets: targets["target_dt_unique"] = batch_cpu["target_dt_unique"].to( device, non_blocking=non_blocking ) targets["target_multi_hot"] = batch_cpu["target_multi_hot"].to( device, non_blocking=non_blocking ) return targets def compute_next_step_loss( args: argparse.Namespace, model: DeepHealth, readout, criterion, batch: Dict[str, torch.Tensor], device: torch.device, ) -> tuple[torch.Tensor, Dict[str, torch.Tensor]]: batch_cpu = batch input_keys = list(MODEL_INPUT_KEYS) input_keys.extend(key for key in EXPOSURE_INPUT_KEYS if key in batch_cpu) batch = move_batch_to_device( {key: batch_cpu[key] for key in input_keys}, device, ) model_kwargs = { "event_seq": batch["event_seq"], "time_seq": batch["time_seq"], "sex": batch["sex"], "padding_mask": batch["padding_mask"], "target_mode": "next_token", "return_output": True, } if "exposure_daily" in batch: model_kwargs["exposure_daily"] = batch["exposure_daily"] model_kwargs["exposure_monthly"] = batch["exposure_monthly"] model_out = model(**model_kwargs) if not isinstance(model_out, DeepHealthOutput): raise TypeError("DeepHealth return_output=True must return DeepHealthOutput") targets = build_augmented_next_step_targets( batch_cpu=batch_cpu, model_out=model_out, include_uts_targets=args.target_mode == "uts", ) readout_out = readout( hidden=model_out.hidden, time_seq=model_out.time_seq, padding_mask=model_out.padding_mask, readout_mask=targets["readout_mask"] if args.readout_name == "same_time_group_end" else None, ) logits = model.calc_risk(readout_out.hidden) if args.target_mode == "delphi2m": loss, parts = criterion( logits=logits, target_events=targets["target_event_seq"], target_times=targets["target_time_seq"], current_times=model_out.time_seq, padding_mask=readout_out.readout_mask, return_components=True, ) else: loss, parts = criterion( logits=logits, target_multi_hot=targets["target_multi_hot"], target_dt_unique=targets["target_dt_unique"], readout_mask=readout_out.readout_mask, return_components=True, ) if not torch.isfinite(loss): raise RuntimeError(f"Loss is not finite: {float(loss.detach().cpu())}") return loss, parts def run_epoch( logger: logging.Logger, args: argparse.Namespace, model: DeepHealth, readout, criterion, loader: DataLoader, optimizer: AdamW | None, device: torch.device, is_train: bool, ) -> float: model.train(is_train) readout.train(is_train) total = torch.zeros((), device=device) n_batches = 0 skipped = 0 parts_sum: Dict[str, torch.Tensor] = {} 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, parts = compute_next_step_loss(args, model, readout, 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 for name, value in parts.items(): parts_sum[name] = parts_sum.get(name, torch.zeros((), device=device)) + value.detach() if (batch_idx + 1) % progress_interval == 0: avg = total / max(1, n_batches) postfix = { "loss": f"{float(loss.detach().cpu()):.4f}", "avg": f"{float(avg.detach().cpu()):.4f}", "skipped": skipped, } for name, value in parts_sum.items(): postfix[name] = f"{float((value / max(1, n_batches)).detach().cpu()):.4f}" progress.set_postfix(postfix) 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: HealthDataset, run_name: str, train_subset, val_subset, test_subset, ) -> Dict[str, Any]: return { "run_name": run_name, "dataset_class": "NextStepHealthDataset", "collate_fn": "next_step_collate_fn", "model_class": "DeepHealth", "model_target_mode": "next_token", "target_mode": args.target_mode, "dist_mode": "exponential", "dataset_metadata": { "vocab_size": int(dataset.vocab_size), }, "use_exposure_encoder": args.exposure_cache_dir is not None, "exposure_cache_dir": args.exposure_cache_dir, "mask_onset_exposure": bool(args.mask_onset_exposure), "exposure_d_model": args.exposure_d_model, "exposure_n_layers": int(args.exposure_n_layers), "exposure_top_k": int(args.exposure_top_k), "exposure_n_convnext_blocks": int(args.exposure_n_convnext_blocks), "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), "split_sizes": { "train": int(len(train_subset)), "val": int(len(val_subset)), "test": int(len(test_subset)), }, "resolved_readout_name": args.readout_name, "resolved_loss_name": args.target_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"absolute_exponential_next_token_{args.target_mode}_" f"gap_{args.no_event_interval_years:g}y_" f"{'exposure' if args.exposure_cache_dir else 'noexposure'}_" f"{timestamp}" ) ) logger = setup_logging(run_dir) logger.info(f"Starting next-step training run: {run_name}") 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}") 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, exposure_cache_dir=args.exposure_cache_dir, mask_onset_exposure=args.mask_onset_exposure, ) 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( dataset=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_dataset( dataset=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"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, ) 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, ) 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, ) model = build_model(args, dataset).to(device) readout = build_next_step_readout(args).to(device) criterion = build_next_step_loss(args) optimizer = AdamW( model.parameters(), lr=args.base_lr, betas=tuple(args.betas), weight_decay=args.weight_decay, ) adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128) save_config( args, run_dir / "train_config.json", extra=build_metadata(args, 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, readout, criterion, train_loader, optimizer, device, True) with torch.no_grad(): val_loss = run_epoch(logger, args, model, readout, 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 next-step test split...") model.load_state_dict(torch.load(best_model_path, map_location=device)) with torch.no_grad(): test_loss = run_epoch(logger, args, model, readout, 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()