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

README.md

@@ -221,11 +221,24 @@ all-future / query-conditioned 监督：
 
 ## 训练
 
-当前 `train.py` 是 next-step 训练入口，使用：
+当前 `train.py` 支持 next-token 和 all-future 两类训练入口：
 
-```python
+- `--model_target_mode next_token`
-HealthDataset = NextStepHealthDataset
+  - 使用 `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
 

train.py

@@ -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(
+    if args.model_target_mode == "all_future":
-        event_seq=event_seq,
+        hidden = model(
-        time_seq=time_seq,
+            event_seq=event_seq,
-        sex=sex,
+            time_seq=time_seq,
-        padding_mask=padding_mask,
+            sex=sex,
-        other_type=batch["other_type"],
+            padding_mask=padding_mask,
-        other_value=batch["other_value"],
+            t_query=batch["t_query"],
-        other_value_kind=batch["other_value_kind"],
+            other_type=batch["other_type"],
-        other_time=batch["other_time"],
+            other_value=batch["other_value"],
-        target_mode="next_token",
+            other_value_kind=batch["other_value_kind"],
-    )
+            other_time=batch["other_time"],
-
+            target_mode="all_future",
-    # 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,
         )
+        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",
+        "dataset_class": dataset_class,
-        "collate_fn": "next_step_collate_fn",
+        "collate_fn": collate_name,
         "model_class": "DeepHealth",
-        "model_target_mode": "next_token",
+        "model_target_mode": args.model_target_mode,
-        "dist_mode": "exponential",
+        "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"}:
+    if args.model_target_mode not in {"next_token", "all_future"}:
-        raise ValueError(f"Unknown target_mode: {args.target_mode}")
+        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(
+    if args.model_target_mode == "all_future":
-        data_prefix=args.data_prefix,
+        dataset = AllFutureHealthDataset(
-        labels_file=args.labels_file,
+            data_prefix=args.data_prefix,
-        no_event_interval_years=args.no_event_interval_years,
+            labels_file=args.labels_file,
-        include_no_event_in_uts_target=args.include_no_event_in_uts_target,
+            split="train",
-        extra_info_types=args.extra_info_types,
+            no_event_interval_years=args.no_event_interval_years,
-    )
+            include_no_event_in_uts_target=args.include_no_event_in_uts_target,
-    logger.info(
+            min_history_events=args.all_future_min_history_events,
-        f"Dataset loaded: {len(dataset)} samples, vocab_size={dataset.vocab_size}")
+            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)