Jiarui Li
1757bcd25b
- Removed BaselineEncoder and CrossAttention classes from models.py. - Introduced OtherInfoTokenizer for handling additional token types. - Updated DeepHealth class to integrate OtherInfoTokenizer and manage extra pooling logic. - Added support for extra_pool_reduce parameter to control pooling behavior. - Modified forward methods to return structured output using DeepHealthOutput dataclass. - Updated training scripts to accommodate changes in model architecture and output handling. - Enhanced error handling and validation for input shapes and types.
443 lines
16 KiB
Markdown
443 lines
16 KiB
Markdown
# DeepHealthNew
|
||
|
||
这个目录包含 DeepHealth 的数据准备、数据集、模型、readout 和 loss 代码。当前版本的核心设计是：
|
||
|
||
```text
|
||
疾病序列 stream + 统一的额外信息 token stream
|
||
```
|
||
|
||
疾病、死亡、checkup 事件仍然保存在事件序列里；性别单独保存在 `basic_info`；其他体检、暴露、生活方式等信息统一整理成 `(type, value, value_kind, time)` token。
|
||
|
||
## 数据准备
|
||
|
||
运行：
|
||
|
||
```bash
|
||
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)`
|
||
- 每行是：
|
||
```text
|
||
(eid, type, value, value_kind, time)
|
||
```
|
||
- `type=0` 预留给 padding
|
||
- `value_kind=1` 表示连续变量
|
||
- `value_kind=2` 表示分类变量
|
||
- 缺失值不会生成 token
|
||
- 当前 UKB 额外信息只有一个时间点，所以 `time` 暂时都是 `date_of_assessment`
|
||
|
||
- `cate_types.csv`
|
||
- 分类变量元信息
|
||
- 字段：
|
||
```text
|
||
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`
|
||
|
||
为了兼容旧训练入口：
|
||
|
||
```python
|
||
HealthDataset = NextStepHealthDataset
|
||
collate_fn = next_step_collate_fn
|
||
```
|
||
|
||
all-future 训练入口会显式使用 `AllFutureHealthDataset` 和 `all_future_collate_fn`。
|
||
|
||
dataset 会输出：
|
||
|
||
```python
|
||
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` 选择纳入哪些额外信息变量：
|
||
|
||
```python
|
||
dataset = NextStepHealthDataset(extra_info_types=[1, 3, 7])
|
||
```
|
||
|
||
如果不传，则使用全部可用 other-info type。
|
||
|
||
dataset 会暴露模型初始化需要的元信息：
|
||
|
||
```python
|
||
dataset.n_types
|
||
dataset.n_cont_types
|
||
dataset.n_categories
|
||
dataset.cont_type_ids
|
||
dataset.vocab_size
|
||
```
|
||
|
||
## 模型
|
||
|
||
模型主体定义在 `models.py`，通用网络模块定义在 `backbones.py`。
|
||
|
||
### OtherInfoTokenizer
|
||
|
||
`DeepHealth` 内部的 `OtherInfoTokenizer` 编码统一的 other-info token：
|
||
|
||
```python
|
||
other_type # (B, K)
|
||
other_value # (B, K)
|
||
other_value_kind # (B, K)
|
||
other_time # (B, K)
|
||
```
|
||
|
||
连续值使用 `TokenAutoDiscretization`：
|
||
|
||
```text
|
||
type_id -> continuous type index -> soft bin embedding
|
||
```
|
||
|
||
分类值使用 dataset 动态计算后的 global category id：
|
||
|
||
```text
|
||
selected type offsets + local category id
|
||
```
|
||
|
||
同一患者、同一 `other_time` 的 extra-info token 会先被池化为一个 token，再并入主序列。默认池化方式是平均池化：
|
||
|
||
```bash
|
||
--extra_pool_reduce mean
|
||
```
|
||
|
||
也可以设为 `sum`。池化后，每个 extra-info 时间点最多产生一个主序列 token。
|
||
|
||
### DeepHealth
|
||
|
||
`DeepHealth` 的统一路径是：
|
||
|
||
```text
|
||
disease tokens + pooled extra-info tokens
|
||
-> temporal backbone
|
||
-> risk head
|
||
```
|
||
|
||
extra-info 不再通过独立的 `BaselineEncoder` 或 `CrossAttention` 注入；它们作为主序列 token 直接参与同一个 temporal transformer。相对时间模式下，主序列内所有 token 共用 `TimeRoPE` 和 `GaussianRBFTimeBasis`。
|
||
|
||
两种目标模式的读出语义不同：
|
||
|
||
- `next_token`
|
||
- 模型内部序列包含 disease tokens 和 pooled extra-info tokens
|
||
- 默认 Tensor 返回值仍只返回 disease token hidden，形状为 `(B, L, D)`，兼容评估和旧调用
|
||
- 训练时使用 `return_output=True` 取完整主序列输出；pooled extra-info tokens 也会产生 logits，并在有未来 disease target 时参与 prediction/loss 监督
|
||
- next-token 模式下 `risk_head.weight` 与 `token_embedding.weight` 使用 weight tying
|
||
|
||
- `all_future`
|
||
- 在合并后的主序列末尾拼接一个 query token
|
||
- query token 的时间是 `t_query`
|
||
- 只读出最后一个 query hidden，形状为 `(B, D)`
|
||
- pooled extra-info tokens 只作为 query 的上下文输入，不单独读出、不参与 loss 监督
|
||
- all-future 模式不使用 weight tying
|
||
|
||
如果需要拿到完整 next-token 输出，可使用结构化返回：
|
||
|
||
```python
|
||
out = model(..., target_mode="next_token", return_output=True)
|
||
out.hidden # disease tokens + pooled extra-info tokens
|
||
out.time_seq # 与 hidden 对齐的时间
|
||
out.padding_mask # 与 hidden 对齐的有效位置
|
||
out.event_len # 原 disease token 长度
|
||
```
|
||
|
||
模型初始化示例：
|
||
|
||
```python
|
||
model = DeepHealth(
|
||
vocab_size=dataset.vocab_size,
|
||
n_embd=120,
|
||
n_head=10,
|
||
n_hist_layer=12,
|
||
n_tab_layer=4, # 兼容旧配置；当前不再创建独立 tabular transformer
|
||
n_types=dataset.n_types,
|
||
n_cont_types=dataset.n_cont_types,
|
||
n_categories=dataset.n_categories,
|
||
cont_type_ids=dataset.cont_type_ids,
|
||
extra_pool_reduce="mean",
|
||
)
|
||
```
|
||
|
||
## Loss
|
||
|
||
`losses.py` 中保留：
|
||
|
||
next-token 监督：
|
||
|
||
- `Delphi2MLoss`
|
||
- `UniqueTimeSetExponentialLoss`
|
||
|
||
next-token 训练中，模型会请求 `return_output=True`，因此 loss 的预测位置包括：
|
||
|
||
- 原 disease token readout 位置
|
||
- 同一时间点 extra-info 池化后的 pooled extra-info token
|
||
|
||
pooled extra-info token 的监督目标在训练时动态构造：对 pooled extra-info token 的时间 `t`，寻找该患者 `t` 之后的下一个 disease 事件时间；`delphi2m` 使用第一个未来事件作为 next-token target，`uts` 使用下一唯一时间点上的事件集合做 multi-hot target。若该 extra-info 时间点之后没有未来 disease target，则该位置不参与 loss。
|
||
|
||
all-future / query-conditioned 监督：
|
||
|
||
- `ExponentialLoss`
|
||
- `WeibullLoss`
|
||
- `MixedLoss`
|
||
|
||
all-future 训练只读出 `t_query` 对应的 query hidden。pooled extra-info tokens 作为主序列上下文输入，但不会被单独读出，也不会被纳入 loss 监督。
|
||
|
||
`UniqueTimeSetExponentialLoss` 的 observed term 固定使用 sum reduction，不再暴露旧的 `observed_reduction` 参数。
|
||
|
||
## 训练
|
||
|
||
当前提供两类训练入口：
|
||
|
||
- `train_next_step.py`
|
||
- 使用 `NextStepHealthDataset`
|
||
- `--target_mode delphi2m` 默认搭配 `Delphi2MLoss` + `token` readout
|
||
- `--target_mode uts` 默认搭配 `UniqueTimeSetExponentialLoss` + `same_time_group_end` readout
|
||
- 当前 next-token 训练只支持 exponential time loss
|
||
- pooled extra-info tokens 会加入 prediction/loss 监督
|
||
- `train_all_future.py`
|
||
- 使用 `AllFutureHealthDataset`
|
||
- 不使用 readout，直接对 query hidden 计算风险
|
||
- `--dist_mode exponential/weibull/mixed` 分别搭配 `ExponentialLoss`、`WeibullLoss`、`MixedLoss`
|
||
- pooled extra-info tokens 只作为 query 上下文，不单独监督
|
||
|
||
当前 `train_next_step.py` / `train_all_future.py` 支持所有已有训练目标定义的组合：
|
||
|
||
| 训练模式 | 时间模式 | 分布/监督 | 默认 loss/readout |
|
||
| --- | --- | --- | --- |
|
||
| `next_token` | `relative`, `absolute` | `target_mode=delphi2m`, `dist_mode=exponential` | `Delphi2MLoss` + `token` |
|
||
| `next_token` | `relative`, `absolute` | `target_mode=uts`, `dist_mode=exponential` | `UniqueTimeSetExponentialLoss` + `same_time_group_end` |
|
||
| `all_future` | `relative`, `absolute` | `dist_mode=exponential` | `ExponentialLoss`，无 readout |
|
||
| `all_future` | `relative`, `absolute` | `dist_mode=weibull` | `WeibullLoss`，无 readout |
|
||
| `all_future` | `relative`, `absolute` | `dist_mode=mixed` | `MixedLoss`，无 readout |
|
||
|
||
示例：
|
||
|
||
```bash
|
||
python train_next_step.py \
|
||
--data_prefix ukb \
|
||
--labels_file labels.csv \
|
||
--target_mode uts \
|
||
--n_embd 120 \
|
||
--n_head 10 \
|
||
--n_hist_layer 12 \
|
||
--n_tab_layer 4 \
|
||
--extra_pool_reduce mean
|
||
```
|
||
|
||
all-future 示例：
|
||
|
||
```bash
|
||
python train_all_future.py \
|
||
--data_prefix ukb \
|
||
--labels_file labels.csv \
|
||
--dist_mode weibull \
|
||
--time_mode relative \
|
||
--extra_pool_reduce mean
|
||
```
|
||
|
||
选择额外信息变量：
|
||
|
||
```bash
|
||
python train_next_step.py --extra_info_types_file extra_info_types_smoking_alcohol_bmi.txt
|
||
```
|
||
|
||
`train_next_step.py` / `train_all_future.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 列表，用于评估脚本复现变量选择
|
||
- `extra_pool_reduce`：同一 `other_time` 的 extra-info tokens 池化方式，默认为 `mean`
|
||
- `model_target_mode`、`time_mode`、`dist_mode`、`dataset_class`、`collate_fn`、`resolved_loss_name`：用于评估脚本重建模型和输入方式
|
||
|
||
## 评估 AUC
|
||
|
||
当前提供两个 AUC 评估入口，二者都已适配新的 `DeepHealth` 模型和统一的 other-info token 输入；AUC 的 DeLong 计算、病例/对照筛选和分层聚合逻辑保持原评估脚本口径。
|
||
|
||
评估脚本默认使用常规 Tensor 返回值：next-token checkpoint 只缓存 disease token/readout 位置的 hidden；all-future checkpoint 只缓存 `t_query` query hidden。训练中额外加入监督的 pooled extra-info tokens 不作为 AUC 评估位置单独输出。
|
||
|
||
### `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。
|
||
- 输出未池化分层结果和按疾病汇总后的结果。
|
||
|
||
典型用法：
|
||
|
||
```bash
|
||
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。
|
||
|
||
查询方式由 `train_config.json` 中的 `model_target_mode` 决定：
|
||
|
||
- `model_target_mode="next_token"`：在 landmark age 人工插入一个 `<NO_EVENT>` token，取该 token 的 hidden 做风险分数。
|
||
- `model_target_mode="all_future"`：不插入 `<NO_EVENT>`，直接把 landmark age 作为 `t_query` 传入模型，取 query hidden 做风险分数。
|
||
|
||
核心流程：
|
||
|
||
- 为每个患者和 landmark age 构造 landmark query 样本。
|
||
- 根据模型模式插入 `<NO_EVENT>` token 或直接传 `t_query`，取 landmark/query hidden。
|
||
- 对疾病 token 分块投影到 `risk_head`；`score_mode="risk"` 时会根据 `dist_mode` 把线性输出转换为固定 horizon 风险概率。
|
||
- 分布转换规则与 all-future 训练损失一致：`exponential` 使用 `1 - exp(-rate * horizon)`；`weibull` 使用 `1 - exp(-rate * horizon ** rho)`；`mixed` 中普通疾病使用 exponential，死亡 endpoint 使用 Weibull death rho。
|
||
- `score_mode="eta"` 是诊断用排序分数，不使用 `rho`，因此不区分不同分布的风险曲线。
|
||
- 按疾病、性别、landmark age、horizon 计算 incident disease AUC。
|
||
- 可选择排除 horizon 内先于目标疾病发生的死亡竞争风险。
|
||
|
||
典型用法：
|
||
|
||
```bash
|
||
python evaluate_auc_v2.py \
|
||
--run_path runs/your_run_dir \
|
||
--eval_split test \
|
||
--landmark_start 40 \
|
||
--landmark_stop 80 \
|
||
--landmark_step 5 \
|
||
--horizons 1,5,10
|
||
```
|
||
|
||
主要输出：
|
||
|
||
- `df_auc_landmark_unpooled.csv`
|
||
- 疾病 token 在 sex、landmark age、horizon 分层下的 AUC。
|
||
- `df_auc_landmark.csv`
|
||
- 按疾病 token 和 horizon 聚合后的 landmark AUC。
|
||
|
||
适合回答的问题：
|
||
|
||
- “一个人在 40/45/50/... 岁这个固定年龄点，如果此前未患某病，未来 1/5/10 年内发生该病的风险区分能力如何？”
|
||
- “模型能否作为 landmark risk prediction 模型使用？”
|
||
|
||
### 两者区别
|
||
|
||
| 项目 | `evaluate_auc.py` | `evaluate_auc_v2.py` |
|
||
| --- | --- | --- |
|
||
| 评估口径 | next-step/token-level 预测点 | landmark fixed-horizon incident risk |
|
||
| 查询位置 | next-token 用满足 offset 条件的最新 readout token；all-future 直接用该预测点年龄作为 `t_query` | next-token 用人工插入的 `<NO_EVENT>` landmark token；all-future 直接用 `t_query` |
|
||
| 时间参数 | `offsets`：预测点至少早于目标事件多少年 | `landmark_*` 和 `horizons`：固定年龄点与未来窗口 |
|
||
| score 与分布 | 使用 eta/logit 排序分数；不按 `dist_mode` 转换风险概率 | `score_mode="risk"` 按 `dist_mode` 区分 exponential / Weibull / mixed；`score_mode="eta"` 不区分分布 |
|
||
| 病例定义 | target table 中出现目标疾病的患者/事件 | landmark 后 horizon 内首次发生目标疾病 |
|
||
| 对照定义 | 从未出现该疾病的患者的 eligible target occurrence | landmark 时未患病，且 horizon 内未发病并有足够随访 |
|
||
| 分层 | sex + age bracket + offset | sex + landmark age + horizon |
|
||
| 输出文件 | `df_auc_unpooled.csv`, `df_both.csv` | `df_auc_landmark_unpooled.csv`, `df_auc_landmark.csv` |
|
||
| 适用问题 | 提前若干年预测未来目标事件的 token-level AUC | 固定年龄点未来固定年限 incident disease risk AUC |
|
||
|
||
简单选择：
|
||
|
||
- 想复现/延续旧的 next-token Delphi 风格 AUC：用 `evaluate_auc.py`。
|
||
- 想做临床上更像 “某年龄点未来 N 年发病风险” 的 landmark AUC：用 `evaluate_auc_v2.py`。
|
||
|
||
## 主要文件
|
||
|
||
- `prepare_data.py`
|
||
- UKB 原始数据到模型输入文件的 ETL
|
||
|
||
- `dataset.py`
|
||
- next-step 和 all-future dataset
|
||
- 动态选择 other-info type
|
||
- 动态计算 categorical global id
|
||
|
||
- `models.py`
|
||
- `DeepHealth`
|
||
- `DeepHealthOutput`
|
||
- `OtherInfoTokenizer`
|
||
|
||
- `backbones.py`
|
||
- `TimeRoPE`
|
||
- `GaussianRBFTimeBasis`
|
||
- `TemporalAttention`
|
||
- `GPTBlock`
|
||
- `TokenAutoDiscretization`
|
||
- `AgeSinusoidalEncoding`
|
||
|
||
- `losses.py`
|
||
- next-token 和 all-future losses
|
||
|
||
- `readouts.py`
|
||
- token readout
|
||
- same-time group readout
|
||
- last-valid readout
|
||
|
||
- `evaluate_auc.py`
|
||
- next-step/token-level 疾病 AUC 评估
|
||
- 使用 prediction offset、sex、age bracket 分层
|
||
- next-token 模型使用 readout hidden；all-future 模型使用 `t_query` hidden
|
||
- score 是疾病 eta/logit，不按分布转换为固定 horizon 风险概率
|
||
|
||
- `evaluate_auc_v2.py`
|
||
- landmark fixed-horizon incident disease AUC 评估
|
||
- next-token 模型通过插入 `<NO_EVENT>` landmark token 查询固定年龄点风险
|
||
- all-future 模型直接通过 `t_query` 查询固定年龄点风险
|
||
- `score_mode="risk"` 按 exponential / Weibull / mixed 分布计算固定 horizon 风险
|