Jiarui Li
5e979e061b
- Implemented target construction in `targets.py` for next-token and unique-time set supervision. - Added validation functions and utility methods for target building. - Created a comprehensive training script in `train.py` that includes data loading, model building, optimizer setup, and training loop with early stopping and logging. - Integrated loss functions and readout mechanisms based on target modes. - Established dataset splitting and DataLoader configurations for training, validation, and testing.
281 lines
5.8 KiB
Markdown
281 lines
5.8 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
|
||
```
|
||
|
||
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`。
|
||
|
||
### BaselineEncoder
|
||
|
||
`BaselineEncoder` 编码统一的 other-info token：
|
||
|
||
```python
|
||
other_type # (B, K)
|
||
other_value # (B, K)
|
||
other_value_kind # (B, K)
|
||
```
|
||
|
||
它暂时不直接使用 `other_time`。时间信息保留给后续 `CrossAttention`，用于建模疾病/query 与 other-info token 的相对时间关系。
|
||
|
||
连续值使用 `TokenAutoDiscretization`：
|
||
|
||
```text
|
||
type_id -> continuous type index -> soft bin embedding
|
||
```
|
||
|
||
分类值使用 dataset 动态计算后的 global category id：
|
||
|
||
```text
|
||
selected type offsets + local category id
|
||
```
|
||
|
||
### CrossAttention
|
||
|
||
`CrossAttention` 让 disease-side hidden state 注意到 other-info token：
|
||
|
||
```python
|
||
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
|
||
|
||
注意力是时间因果的：
|
||
|
||
```text
|
||
other_info_time <= disease_or_query_time
|
||
```
|
||
|
||
如果某个 disease/query 位置没有任何可见 other-info token，则该位置保持原 hidden 不变。
|
||
|
||
### DeepHealth
|
||
|
||
`DeepHealth` 的统一路径是：
|
||
|
||
```text
|
||
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)`
|
||
|
||
模型初始化示例：
|
||
|
||
```python
|
||
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-step 训练入口，使用：
|
||
|
||
```python
|
||
HealthDataset = NextStepHealthDataset
|
||
```
|
||
|
||
示例：
|
||
|
||
```bash
|
||
python train.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
|
||
```
|
||
|
||
选择额外信息变量：
|
||
|
||
```bash
|
||
python train.py --extra_info_types 1 3 7
|
||
```
|
||
|
||
如果不传 `--extra_info_types`，默认使用全部 other-info type。
|
||
|
||
## 主要文件
|
||
|
||
- `prepare_data.py`
|
||
- UKB 原始数据到模型输入文件的 ETL
|
||
|
||
- `dataset.py`
|
||
- next-step 和 all-future dataset
|
||
- 动态选择 other-info type
|
||
- 动态计算 categorical global id
|
||
|
||
- `models.py`
|
||
- `DeepHealth`
|
||
|
||
- `backbones.py`
|
||
- `TimeRoPE`
|
||
- `GaussianRBFTimeBasis`
|
||
- `TemporalAttention`
|
||
- `GPTBlock`
|
||
- `TokenAutoDiscretization`
|
||
- `BaselineEncoder`
|
||
- `CrossAttention`
|
||
- `AgeSinusoidalEncoding`
|
||
|
||
- `losses.py`
|
||
- next-token 和 all-future losses
|
||
|
||
- `readouts.py`
|
||
- token readout
|
||
- same-time group readout
|
||
- last-valid readout
|