Files

Jiarui Li 82f70945d9 Enhance training script to support all-future model target mode and update dataset handling

2026-06-12 11:49:44 +08:00

13 KiB

Raw Blame History

DeepHealthNew

这个目录包含 DeepHealth 的数据准备、数据集、模型、readout 和 loss 代码。当前版本的核心设计是：

疾病序列 stream + 统一的额外信息 token stream

疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 basic_info；其他体检、暴露、生活方式等信息统一整理成 (type, value, value_kind, time) token。

数据准备

运行：

python prepare_data.py

输入文件：

ukb_data.csv
field_ids_enriched.csv
icd10_codes_mod.tsv
labels.csv

输出文件：

ukb_event_data.npy
- 形状为 (N, 3)
- 每行是 (eid, days, label)
- 包含疾病、死亡、checkup 事件
ukb_basic_info.csv
- index 为 eid
- 当前只保留 sex
ukb_other_info.npy
- 形状为 (M, 5)
- 每行是：
```
(eid, type, value, value_kind, time)
```
- type=0 预留给 padding
- value_kind=1 表示连续变量
- value_kind=2 表示分类变量
- 缺失值不会生成 token
- 当前 UKB 额外信息只有一个时间点，所以 time 暂时都是 date_of_assessment
cate_types.csv
- 分类变量元信息
- 字段：
```
type,name,n_categories
```
- ukb_other_info.npy 里的分类 value 是变量内部的局部 id；global category id 在 dataset 中根据实验选择的变量动态计算。

Dataset

dataset.py 提供两个 dataset：

NextStepHealthDataset
- 用于 next-token / next-time-point 监督
- 对应 Delphi2MLoss 和 UniqueTimeSetExponentialLoss
AllFutureHealthDataset
- 用于 query-conditioned all-future 监督
- 对应 ExponentialLoss、WeibullLoss、MixedLoss

为了兼容旧训练入口：

HealthDataset = NextStepHealthDataset
collate_fn = next_step_collate_fn

dataset 会输出：

event_seq
time_seq
padding_mask
sex
other_type
other_value
other_value_kind
other_time

其中 other_type=0 表示 padding，不额外传 other-token mask。

可以通过 extra_info_types 选择纳入哪些额外信息变量：

dataset = NextStepHealthDataset(extra_info_types=[1, 3, 7])

如果不传，则使用全部可用 other-info type。

dataset 会暴露模型初始化需要的元信息：

dataset.n_types
dataset.n_cont_types
dataset.n_categories
dataset.cont_type_ids
dataset.vocab_size

模型

模型主体定义在 models.py，通用网络模块定义在 backbones.py。

BaselineEncoder

BaselineEncoder 编码统一的 other-info token：

other_type        # (B, K)
other_value       # (B, K)
other_value_kind  # (B, K)

它暂时不直接使用 other_time。时间信息保留给后续 CrossAttention，用于建模疾病/query 与 other-info token 的相对时间关系。

连续值使用 TokenAutoDiscretization：

type_id -> continuous type index -> soft bin embedding

分类值使用 dataset 动态计算后的 global category id：

selected type offsets + local category id

CrossAttention

CrossAttention 让 disease-side hidden state 注意到 other-info token：

h_disease  # (B, L, D)
t_disease  # (B, L)
h_token    # (B, K, D)
t_token    # (B, K)

时间信息通过两种方式进入注意力：

TimeRoPE
- 使用 query time 和 key time 旋转 q/k
- 让 q-k 相似度带有时间位置信息
GaussianRBFTimeBasis
- 对 t_disease - t_token 做 RBF 编码
- 投影成每个 attention head 的时间 bias

注意力是时间因果的：

other_info_time <= disease_or_query_time

如果某个 disease/query 位置没有任何可见 other-info token，则该位置保持原 hidden 不变。

DeepHealth

DeepHealth 的统一路径是：

disease-side sequence
-> disease temporal backbone
-> CrossAttention 到 other-info tokens
-> risk head

两种目标模式共用同一套语义：

next_token
- h_disease 长度为 L
- 输出 (B, L, D)
all_future
- 在 disease-side 序列末尾拼接一个 query token
- query token 的时间是 t_query
- h_disease 长度为 L + 1
- 输出最后一个 query hidden，形状为 (B, D)

模型初始化示例：

model = DeepHealth(
    vocab_size=dataset.vocab_size,
    n_embd=120,
    n_head=10,
    n_hist_layer=12,
    n_tab_layer=4,
    n_types=dataset.n_types,
    n_cont_types=dataset.n_cont_types,
    n_categories=dataset.n_categories,
    cont_type_ids=dataset.cont_type_ids,
)

Loss

losses.py 中保留：

next-token 监督：

Delphi2MLoss
UniqueTimeSetExponentialLoss

all-future / query-conditioned 监督：

ExponentialLoss
WeibullLoss
MixedLoss

UniqueTimeSetExponentialLoss 的 observed term 固定使用 sum reduction，不再暴露旧的 observed_reduction 参数。

训练

当前 train.py 支持 next-token 和 all-future 两类训练入口：

--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`

示例：

python train.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --model_target_mode next_token \
  --target_mode uts \
  --n_embd 120 \
  --n_head 10 \
  --n_hist_layer 12 \
  --n_tab_layer 4

all-future 示例：

python train.py \
  --data_prefix ukb \
  --labels_file labels.csv \
  --model_target_mode all_future \
  --dist_mode weibull \
  --time_mode relative

选择额外信息变量：

python train.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt

train.py 只接受 --extra_info_types_file 指定变量列表，不接受在 CLI 里直接输入 type id。文件可以每行一个 type id，也可以带 # 注释；如果不传 --extra_info_types_file，默认使用全部 other-info type。

训练输出的 train_config.json 会记录：

extra_info_types_file：训练时使用的列表文件名
extra_info_types：解析后的实际 type id 列表，用于评估脚本复现变量选择
model_target_mode、time_mode、dist_mode、dataset_class、collate_fn、resolved_loss_name：用于评估脚本重建模型和输入方式

评估 AUC

当前提供两个 AUC 评估入口，二者都已适配新的 DeepHealth 模型和统一的 other-info token 输入；AUC 的 DeLong 计算、病例/对照筛选和分层聚合逻辑保持原评估脚本口径。

`evaluate_auc.py`

evaluate_auc.py 评估的是 next-step / token-level 预测位置上的疾病 AUC。

查询方式由 train_config.json 中的 model_target_mode 决定：

model_target_mode="next_token"：使用训练 readout 对应的历史 token hidden 作为预测点。
model_target_mode="all_future"：不使用 readout token，直接把每个预测点年龄作为 t_query 传入模型，取 query hidden 作为预测点。

核心流程：

按训练配置重新构建 HealthDataset 和 DeepHealth。
对评估 split 中的患者做模型推理，缓存每个预测点的 hidden。
对疾病 token 分块投影到 risk_head，避免一次性保存全词表 logits。
AUC score 使用疾病对应的 eta/logit 排序分数；dist_mode 只用于正确构建模型，不会把分数转换成 horizon-specific risk probability。
对每个疾病、性别、年龄段、prediction offset 分别计算 AUC。
输出未池化分层结果和按疾病汇总后的结果。

典型用法：

python evaluate_auc.py \
  --run_path runs/your_run_dir \
  --eval_split test \
  --offsets 0.1,1,5,10

主要输出：

df_auc_unpooled.csv
- 疾病 token 在 sex、age bracket、offset 分层下的 AUC。
df_both.csv
- 按疾病 token 和 offset 聚合后的 AUC。

适合回答的问题：

“模型在历史序列中的某个预测 token 上，提前 offset 年预测未来疾病的区分能力如何？”
“不同年龄段、性别、提前量下，next-step 训练模型的疾病预测 AUC 如何？”

`evaluate_auc_v2.py`

evaluate_auc_v2.py 评估的是 landmark fixed-horizon incident disease AUC。

它不是使用已有序列中的普通 readout 位置，而是在指定 landmark age 构造一个 landmark query，然后评估该 landmark 后固定 horizon 内是否发生 incident disease。