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

Remove extra info token pathway

Browse Source
This commit is contained in:
Jiarui Li
2026-07-07 16:57:49 +08:00
parent 6dfeb5a696
commit a0379daf29
13 changed files with 18 additions and 1390 deletions
Download Patch File Download Diff File Expand all files Collapse all files

192
train_next_step.py
View File

@@ -12,7 +12,6 @@ import json
import logging
import math
import time
from pathlib import Path
from typing import Any, Dict
import numpy as np
@@ -30,8 +29,6 @@ from targets import CHECKUP_IDX, NO_EVENT_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,
@@ -48,10 +45,6 @@ MODEL_INPUT_KEYS = (
"time_seq",
"sex",
"padding_mask",
"other_type",
"other_value",
"other_value_kind",
"other_time",
)
@@ -62,7 +55,6 @@ def parse_args() -> argparse.Namespace:
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("--no_event_interval_years", type=float, default=5.0)
parser.add_argument("--include_no_event_in_uts_target", action="store_true")
@@ -76,10 +68,6 @@ def parse_args() -> argparse.Namespace:
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("--dropout", type=float, default=0.0)
@@ -121,11 +109,6 @@ def parse_args() -> argparse.Namespace:
args.include_no_event_in_uts_target = True
else:
args.readout_name = args.readout_name or "token"
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
@@ -153,13 +136,6 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
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="next_token",
time_mode=args.time_mode,
dist_mode="exponential",
@@ -199,153 +175,20 @@ def build_augmented_next_step_targets(
model_out: DeepHealthOutput,
include_uts_targets: bool,
) -> Dict[str, torch.Tensor]:
hidden_len = model_out.hidden.size(1)
event_len = int(model_out.event_len)
extra_len = hidden_len - event_len
device = model_out.hidden.device
non_blocking = device.type == "cuda"
if extra_len <= 0:
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
bsz = batch_cpu["target_event_seq"].size(0)
vocab_size = (
batch_cpu["target_multi_hot"].size(2)
if include_uts_targets
else None
)
other_valid = batch_cpu["other_type"] > 0
extra_time = batch_cpu["other_time"].new_zeros(bsz, extra_len)
extra_mask = torch.zeros(bsz, extra_len, dtype=torch.bool)
for b in range(bsz):
unique_time = torch.unique(batch_cpu["other_time"][b, other_valid[b]], sorted=True)
n_time = min(int(unique_time.numel()), extra_len)
if n_time > 0:
extra_time[b, :n_time] = unique_time[:n_time]
extra_mask[b, :n_time] = True
target_event_seq = torch.cat(
[
batch_cpu["target_event_seq"],
torch.full(
(bsz, extra_len),
PAD_IDX,
dtype=batch_cpu["target_event_seq"].dtype,
),
],
dim=1,
)
target_time_seq = torch.cat(
[
batch_cpu["target_time_seq"],
torch.zeros(
bsz,
extra_len,
dtype=batch_cpu["target_time_seq"].dtype,
),
],
dim=1,
)
readout_mask = torch.cat([batch_cpu["readout_mask"], extra_mask], dim=1)
target_dt_unique = None
target_multi_hot = None
if include_uts_targets:
target_dt_unique = torch.cat(
[
batch_cpu["target_dt_unique"],
torch.zeros(
bsz,
extra_len,
dtype=batch_cpu["target_dt_unique"].dtype,
),
],
dim=1,
)
target_multi_hot = torch.cat(
[
batch_cpu["target_multi_hot"],
torch.zeros(
bsz,
extra_len,
vocab_size,
dtype=batch_cpu["target_multi_hot"].dtype,
),
],
dim=1,
)
for b in range(bsz):
valid_event = batch_cpu["padding_mask"][b].bool()
if not valid_event.any():
continue
n_event = int(valid_event.sum().item())
events = torch.cat(
[
batch_cpu["event_seq"][b, :n_event],
batch_cpu["target_event_seq"][b, n_event - 1:n_event],
]
)
times = torch.cat(
[
batch_cpu["time_seq"][b, :n_event],
batch_cpu["target_time_seq"][b, n_event - 1:n_event],
]
)
valid_full = events > PAD_IDX
events = events[valid_full]
times = times[valid_full]
if events.numel() == 0:
continue
for j in range(extra_len):
if not bool(extra_mask[b, j]):
continue
pos = event_len + j
t = extra_time[b, j]
future = times > t
if not future.any():
readout_mask[b, pos] = False
continue
first_idx = int(torch.nonzero(future, as_tuple=False)[0].item())
next_time = times[first_idx]
next_event = events[first_idx]
target_event_seq[b, pos] = next_event
target_time_seq[b, pos] = next_time
if not include_uts_targets:
continue
same_next_time = times == next_time
next_events = events[same_next_time]
valid_next_events = next_events[
(next_events > PAD_IDX) & (next_events < vocab_size)
].long()
if valid_next_events.numel() == 0:
readout_mask[b, pos] = False
continue
target_multi_hot[b, pos, valid_next_events] = True
target_dt_unique[b, pos] = next_time - t
targets = {
"target_event_seq": target_event_seq.to(device, non_blocking=non_blocking),
"target_time_seq": target_time_seq.to(device, non_blocking=non_blocking),
"readout_mask": readout_mask.to(device, non_blocking=non_blocking),
"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"] = target_dt_unique.to(device, non_blocking=non_blocking)
targets["target_multi_hot"] = target_multi_hot.to(device, non_blocking=non_blocking)
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
@@ -367,10 +210,6 @@ def compute_next_step_loss(
time_seq=batch["time_seq"],
sex=batch["sex"],
padding_mask=batch["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",
return_output=True,
)
@@ -487,19 +326,8 @@ def build_metadata(
"model_target_mode": "next_token",
"target_mode": args.target_mode,
"dist_mode": "exponential",
"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)),
@@ -527,7 +355,6 @@ def main() -> None:
logger.info(f"Starting next-step 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(f"readout={args.readout_name}, target_mode={args.target_mode}")
dataset = HealthDataset(
@@ -535,7 +362,6 @@ def main() -> None:
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,
)
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(