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

Unify FFN and TrajMixer model architectures

Browse Source
This commit is contained in:
Jiarui Li
2026-07-25 12:59:09 +08:00
parent 4526191fe1
commit b13db5e407
18 changed files with 1019 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
MODEL_ARCHITECTURES.md Normal file
View File

@@ -0,0 +1,64 @@
# Model architectures
DeepHealth uses one codebase for both supported history-block architectures.
Select the architecture explicitly when starting a training run:
| `model_architecture` | History block | Checkpoint fingerprint |
| --- | --- | --- |
| `transformer_ffn_v1` | Temporal attention + SwiGLU FFN | `blocks.*.mlp.w1/w2/w3` and `blocks.*.ln2` |
| `traj_mixer_v5` | Temporal attention + TrajMixer | `blocks.*.mlp.intra_*`, `gate_proj`, and `output_proj` |
`transformer_ffn_v1` is the CLI default; pass `traj_mixer_v5` explicitly for
TrajMixer runs.
## Training
Next-step example:
```powershell
python train_next_step.py --model_architecture traj_mixer_v5 --n_layer 12
```
All-future example:
```powershell
python train_all_future.py --model_architecture transformer_ffn_v1 --n_layer 12
```
New runs are separated by architecture:
```text
runs/
transformer_ffn_v1/
<run_name>/
traj_mixer_v5/
<run_name>/
```
Use `--runs_root` to place this structure under a different root. Existing run
directories are not moved or renamed.
Each generated `train_config.json` records `model_architecture`, total parameter
count, and trainable parameter count.
Both training entry points use the single `--n_layer` option to set the number
of history backbone blocks. The same value is passed to `DeepHealth.n_layer`
and saved as `n_layer` in `train_config.json`; it must be at least 1.
## Architecture validation
Evaluation resolves the architecture before constructing the model and always
loads weights with `strict=True`.
- Every config must include an explicit `model_architecture` marker.
- Checkpoint fingerprints are used to validate that the selected architecture
matches the stored weights.
- A config marker that conflicts with the checkpoint fingerprint raises an
error instead of silently choosing one architecture.
- Unsupported historical TrajMixer markers such as `traj_mixer_v2`,
`traj_mixer_v3`, and `traj_mixer_v4` are rejected.
- Checkpoints and configs created before architecture markers were introduced
are intentionally unsupported.
Project code should use the architecture factory rather than instantiate a
history block directly.

203
backbones.py
View File

@@ -4,6 +4,12 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from model_architectures import (
TRAJ_MIXER_ARCHITECTURE,
TRANSFORMER_FFN_ARCHITECTURE,
resolve_model_architecture,
)
class TimeRoPE(nn.Module):
def __init__(self, dim: int, base: float = 10000.0):
@@ -213,7 +219,133 @@ class SwiGLU(nn.Module):
return self.drop(self.w3(F.silu(self.w1(x)) * self.w2(x)))
class GPTBlock(nn.Module):
class TrajMixer(nn.Module):
"""PreNorm gated mixing within and across latent trajectory groups.
The groups are contiguous partitions of the post-``W_O`` residual
representation. They are deliberately not treated as attention heads.
All operations are position-wise, so the sequence dimension remains fully
parallel and no temporal information can leak between positions here.
"""
def __init__(
self,
n_embd: int,
n_head: int = 10,
dropout: float = 0.0,
):
super().__init__()
if n_embd <= 0:
raise ValueError(f"n_embd must be > 0, got {n_embd}")
if n_head <= 0:
raise ValueError(f"n_head must be > 0, got {n_head}")
if n_embd % n_head != 0:
raise ValueError(
f"n_embd must be divisible by n_head, got {n_embd} and {n_head}"
)
self.n_embd = n_embd
self.n_group = n_head
self.d_group = n_embd // n_head
self.intra_hidden = 4 * self.d_group
self.hidden_group = 4 * n_head
self.norm = nn.LayerNorm(self.n_embd)
self.intra_gate_proj = nn.Parameter(
torch.empty(self.n_group, self.d_group, self.intra_hidden)
)
self.intra_value_proj = nn.Parameter(
torch.empty(self.n_group, self.d_group, self.intra_hidden)
)
self.intra_output_proj = nn.Parameter(
torch.empty(self.n_group, self.intra_hidden, self.d_group)
)
self.intra_gate_logits = nn.Parameter(
torch.empty(self.n_group, self.d_group)
)
self.gate_proj = nn.Parameter(
torch.empty(self.d_group, self.n_group, self.hidden_group)
)
self.value_proj = nn.Parameter(
torch.empty(self.d_group, self.n_group, self.hidden_group)
)
self.output_proj = nn.Parameter(
torch.empty(self.d_group, self.hidden_group, self.n_group)
)
self.drop = nn.Dropout(dropout)
self.reset_parameters()
def reset_parameters(self) -> None:
for group_idx in range(self.n_group):
nn.init.xavier_uniform_(self.intra_gate_proj[group_idx])
nn.init.xavier_uniform_(self.intra_value_proj[group_idx])
nn.init.xavier_uniform_(self.intra_output_proj[group_idx])
nn.init.constant_(
self.intra_gate_logits,
math.log(0.1 / 0.9),
)
for feature_idx in range(self.d_group):
nn.init.xavier_uniform_(self.gate_proj[feature_idx])
nn.init.xavier_uniform_(self.value_proj[feature_idx])
nn.init.normal_(self.output_proj, mean=0.0, std=1e-3)
def _intra_mix(self, grouped: torch.Tensor) -> torch.Tensor:
"""Mix features independently inside each residual-space group."""
intra_gate = torch.einsum(
"blgd,gdh->blgh", grouped, self.intra_gate_proj
)
intra_value = torch.einsum(
"blgd,gdh->blgh", grouped, self.intra_value_proj
)
intra_hidden = F.silu(intra_gate) * intra_value
return torch.einsum(
"blgh,ghd->blgd", intra_hidden, self.intra_output_proj
)
def _cross_mix(self, grouped: torch.Tensor) -> torch.Tensor:
"""Mix groups independently for each within-group coordinate."""
gate = torch.einsum(
"blgr,rgh->blhr", grouped, self.gate_proj
)
value = torch.einsum(
"blgr,rgh->blhr", grouped, self.value_proj
)
hidden = F.silu(gate) * value
return torch.einsum(
"blhr,rhg->blgr", hidden, self.output_proj
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Apply one full-width PreNorm and one outer residual update."""
if x.ndim != 3:
raise ValueError(
f"TrajMixer expects a 3D tensor, got shape {tuple(x.shape)}"
)
if x.size(-1) != self.n_embd:
raise ValueError(
f"Expected hidden size {self.n_embd}, got {x.size(-1)}"
)
batch_size, seq_len, _ = x.shape
grouped = self.norm(x).reshape(
batch_size, seq_len, self.n_group, self.d_group
)
intra_output = self._intra_mix(grouped)
intra_gate = torch.sigmoid(self.intra_gate_logits).view(
1, 1, self.n_group, self.d_group
)
mixed_input = grouped + intra_gate * intra_output
update = self._cross_mix(mixed_input).reshape(
batch_size, seq_len, self.n_embd
)
return x + self.drop(update)
class TransformerFFNBlock(nn.Module):
def __init__(
self,
n_embd: int,
@@ -250,6 +382,75 @@ class GPTBlock(nn.Module):
return x
class TrajMixerBlock(nn.Module):
def __init__(
self,
n_embd: int,
n_head: int,
attn_dropout: float = 0.0,
mlp_dropout: float = 0.0,
use_time_rope: bool = False,
use_rbf_bias: bool = False,
n_rbf_bases: int = 16,
):
super().__init__()
self.attn = TemporalAttention(
n_embd=n_embd,
n_head=n_head,
n_rbf_bases=n_rbf_bases,
dropout=attn_dropout,
use_time_rope=use_time_rope,
use_rbf_bias=use_rbf_bias,
)
self.mlp = TrajMixer(
n_embd=n_embd,
n_head=n_head,
dropout=mlp_dropout,
)
self.ln1 = nn.LayerNorm(n_embd)
def forward(
self,
x: torch.Tensor,
rope_cache: tuple[torch.Tensor, torch.Tensor] | None = None,
rbf_cache: torch.Tensor | None = None,
attn_mask: torch.Tensor | None = None,
) -> torch.Tensor:
x = x + self.attn(self.ln1(x), rope_cache, rbf_cache, attn_mask)
return self.mlp(x)
def build_backbone_block(
model_architecture: str,
*,
n_embd: int,
n_head: int,
attn_dropout: float = 0.0,
mlp_dropout: float = 0.0,
use_time_rope: bool = False,
use_rbf_bias: bool = False,
n_rbf_bases: int = 16,
) -> nn.Module:
"""Build one history block for a supported model architecture."""
architecture = resolve_model_architecture(model_architecture)
block_class: type[nn.Module]
if architecture == TRANSFORMER_FFN_ARCHITECTURE:
block_class = TransformerFFNBlock
elif architecture == TRAJ_MIXER_ARCHITECTURE:
block_class = TrajMixerBlock
else: # pragma: no cover - guarded by resolve_model_architecture.
raise ValueError(f"Unsupported model architecture: {architecture!r}")
return block_class(
n_embd=n_embd,
n_head=n_head,
attn_dropout=attn_dropout,
mlp_dropout=mlp_dropout,
use_time_rope=use_time_rope,
use_rbf_bias=use_rbf_bias,
n_rbf_bases=n_rbf_bases,
)
class TokenAutoDiscretization(nn.Module):
def __init__(
self,

23
evaluate_auc.py
View File

@@ -45,6 +45,7 @@ from delphi2m_auc_report import (
build_delphi2m_auc_report,
)
from eval_data import load_sequence_eval_dataset, sequence_eval_collate_fn
from model_architectures import resolve_model_architecture
from models import DeepHealth
from readouts import build_readout
from targets import PAD_IDX, CHECKUP_IDX, NO_EVENT_IDX
@@ -313,19 +314,24 @@ def split_indices(n: int, train_ratio: float, val_ratio: float, test_ratio: floa
return idx[:n_train], idx[n_train:n_train + n_val], idx[n_train + n_val:]
def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
def build_model_from_dataset(
args: argparse.Namespace,
cfg: Dict[str, Any],
dataset: HealthDataset,
state_dict: Optional[Dict[str, Any]] = None,
) -> DeepHealth:
model_target_mode = str(cfg_get(
args, cfg, "model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
)
model_architecture = resolve_model_architecture(cfg, state_dict)
return DeepHealth(
vocab_size=dataset.vocab_size,
n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
n_head=int(cfg_get(args, cfg, "n_head", 10)),
n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
n_layer=int(cfg["n_layer"]),
n_types=dataset.n_types,
n_cont_types=dataset.n_cont_types,
n_categories=dataset.n_categories,
@@ -336,6 +342,7 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
model_architecture=model_architecture,
)
@@ -383,7 +390,7 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
def load_model_state(
model: torch.nn.Module,
model: DeepHealth,
checkpoint_path: str,
device: torch.device,
state_dict: Optional[Dict[str, Any]] = None,
@@ -391,6 +398,7 @@ def load_model_state(
state = state_dict if state_dict is not None else load_checkpoint_state_dict(
checkpoint_path, map_location=device)
resolve_model_architecture(model.model_architecture, state)
model.load_state_dict(state, strict=True)
@@ -1371,14 +1379,19 @@ def main() -> None:
cfg = dict(cfg)
cfg["dist_mode"] = dist_mode
cfg["model_target_mode"] = model_target_mode
model_architecture = resolve_model_architecture(cfg, state_dict)
cfg["model_architecture"] = model_architecture
print(f"Resolved dist_mode for evaluation: {dist_mode}")
print(f"Resolved model architecture: {model_architecture}")
print(f"Model target mode for AUC: {model_target_mode}")
print(
"AUC score semantics: evaluate_auc.py uses disease-specific eta/logit scores; "
"dist_mode affects model loading but is not converted to horizon-specific risk probability."
)
model = build_model_from_dataset(args, cfg, dataset).to(device)
model = build_model_from_dataset(
args, cfg, dataset, state_dict=state_dict
).to(device)
load_model_state(model, str(model_ckpt_path),
device, state_dict=state_dict)
model.eval()

23
evaluate_auc_v2.py
View File

@@ -36,6 +36,7 @@ from delphi2m_auc_report import (
build_delphi2m_auc_report,
)
from eval_data import load_sequence_eval_dataset
from model_architectures import resolve_model_architecture
from models import DeepHealth
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -184,19 +185,24 @@ def resolve_dist_mode_for_checkpoint(cfg_dist_mode: str, state_dict: Dict[str, A
return mode if mode in {"exponential", "weibull", "mixed"} else "exponential"
def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], dataset: HealthDataset) -> DeepHealth:
def build_model_from_dataset(
args: argparse.Namespace,
cfg: Dict[str, Any],
dataset: HealthDataset,
state_dict: Optional[Dict[str, Any]] = None,
) -> DeepHealth:
model_target_mode = str(cfg_get(
args, cfg, "model_target_mode", "next_token")).lower()
if model_target_mode not in {"next_token", "all_future"}:
raise ValueError(
f"model_target_mode must be next_token or all_future, got {model_target_mode!r}"
)
model_architecture = resolve_model_architecture(cfg, state_dict)
return DeepHealth(
vocab_size=dataset.vocab_size,
n_embd=int(cfg_get(args, cfg, "n_embd", 120)),
n_head=int(cfg_get(args, cfg, "n_head", 10)),
n_hist_layer=int(cfg_get(args, cfg, "n_hist_layer", 12)),
n_tab_layer=int(cfg_get(args, cfg, "n_tab_layer", 4)),
n_layer=int(cfg["n_layer"]),
n_types=dataset.n_types,
n_cont_types=dataset.n_cont_types,
n_categories=dataset.n_categories,
@@ -207,10 +213,12 @@ def build_model_from_dataset(args: argparse.Namespace, cfg: Dict[str, Any], data
time_mode=str(cfg_get(args, cfg, "time_mode", "relative")),
dist_mode=str(cfg_get(args, cfg, "dist_mode", "exponential")),
dropout=float(cfg_get(args, cfg, "dropout", 0.0)),
model_architecture=model_architecture,
)
def load_model_state(model: torch.nn.Module, state_dict: Dict[str, Any]) -> None:
def load_model_state(model: DeepHealth, state_dict: Dict[str, Any]) -> None:
resolve_model_architecture(model.model_architecture, state_dict)
model.load_state_dict(state_dict, strict=True)
@@ -1392,12 +1400,17 @@ def main() -> None:
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
model_architecture = resolve_model_architecture(cfg_model, state_dict)
cfg_model["model_architecture"] = model_architecture
print(f"Resolved model architecture: {model_architecture}")
device = resolve_eval_device(args.device)
if device.type == "cuda":
torch.backends.cudnn.benchmark = True
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
model = build_model_from_dataset(
args, cfg_model, dataset, state_dict=state_dict
).to(device)
if (
model_target_mode == "next_token"

4
evaluate_event_free_survival.py
View File

@@ -649,7 +649,9 @@ def main() -> None:
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device)
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
model = build_model_from_dataset(
args, cfg_model, dataset, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()

4
evaluate_extra_info_attribution.py
View File

@@ -758,7 +758,9 @@ def main() -> None:
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device)
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
model = build_model_from_dataset(
args, cfg_model, dataset, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()

4
evaluate_single_disease_mortality_attribution.py
View File

@@ -553,7 +553,9 @@ def main() -> None:
device = resolve_eval_device(args.device)
selected_token_mask = np.zeros(int(dataset.vocab_size), dtype=bool)
selected_token_mask[np.asarray(scanned_disease_tokens, dtype=np.int64)] = True
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
model = build_model_from_dataset(
args, cfg_model, dataset, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()

4
export_tquery_logits_hidden.py
View File

@@ -180,7 +180,9 @@ def main() -> None:
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device)
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
model = build_model_from_dataset(
args, cfg_model, dataset, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()

4
export_weibull_death_parameter_stats.py
View File

@@ -381,7 +381,9 @@ def main() -> None:
cfg_model = dict(cfg)
cfg_model["dist_mode"] = dist_mode
device = resolve_eval_device(args.device)
model = build_model_from_dataset(args, cfg_model, dataset).to(device)
model = build_model_from_dataset(
args, cfg_model, dataset, state_dict=state_dict
).to(device)
load_model_state(model, state_dict)
model.eval()

131
model_architectures.py Normal file
View File

@@ -0,0 +1,131 @@
"""Model-architecture identifiers and checkpoint validation helpers."""
from __future__ import annotations
import re
from collections.abc import Mapping
TRANSFORMER_FFN_ARCHITECTURE = "transformer_ffn_v1"
TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
DEFAULT_MODEL_ARCHITECTURE = TRANSFORMER_FFN_ARCHITECTURE
SUPPORTED_MODEL_ARCHITECTURES = (
TRANSFORMER_FFN_ARCHITECTURE,
TRAJ_MIXER_ARCHITECTURE,
)
_FFN_STATE_KEY = re.compile(
r"(?:^|\.)blocks\.\d+\.mlp\.w[123]\.(?:weight|bias)$"
)
_TRAJ_MIXER_STATE_KEY = re.compile(
r"(?:^|\.)blocks\.\d+\.mlp\.(?:"
r"norm\.(?:weight|bias)|"
r"intra_gate_proj|"
r"intra_value_proj|"
r"intra_output_proj|"
r"intra_gate_logits|"
r"gate_proj|"
r"value_proj|"
r"output_proj"
r")$"
)
def _validate_model_architecture(model_architecture: object) -> str:
if not isinstance(model_architecture, str):
raise ValueError(
"model_architecture must be one of "
f"{SUPPORTED_MODEL_ARCHITECTURES}, got {model_architecture!r}"
)
if model_architecture not in SUPPORTED_MODEL_ARCHITECTURES:
raise ValueError(
f"Unsupported model_architecture={model_architecture!r}; "
f"expected one of {SUPPORTED_MODEL_ARCHITECTURES}."
)
return model_architecture
def detect_model_architecture_from_state_dict(
state_dict: Mapping[str, object],
) -> str:
"""Infer the architecture from block parameter names.
Detection deliberately accepts any ``blocks.<index>`` prefix rather than
assuming that block zero is present.
"""
if not isinstance(state_dict, Mapping):
raise TypeError(
"state_dict must be a mapping, got "
f"{type(state_dict).__name__}"
)
has_ffn = False
has_traj_mixer = False
for raw_key in state_dict:
key = str(raw_key)
has_ffn = has_ffn or _FFN_STATE_KEY.search(key) is not None
has_traj_mixer = (
has_traj_mixer
or _TRAJ_MIXER_STATE_KEY.search(key) is not None
)
if has_ffn and has_traj_mixer:
raise ValueError(
"Checkpoint contains both Transformer FFN and TrajMixer "
"block parameters; its model architecture is ambiguous."
)
if has_ffn:
return TRANSFORMER_FFN_ARCHITECTURE
if has_traj_mixer:
return TRAJ_MIXER_ARCHITECTURE
raise ValueError(
"Could not detect model architecture from checkpoint parameters. "
"Expected a blocks.<index>.mlp FFN or TrajMixer parameter."
)
def resolve_model_architecture(
config_or_marker: Mapping[str, object] | str | None = None,
state_dict: Mapping[str, object] | None = None,
) -> str:
"""Resolve and cross-check a configured and checkpoint architecture.
Every saved run must provide an explicit architecture marker. Checkpoint
parameter names are used only to verify that the marker describes the
weights being loaded.
"""
if isinstance(config_or_marker, Mapping):
configured = config_or_marker.get("model_architecture")
elif isinstance(config_or_marker, str) or config_or_marker is None:
configured = config_or_marker
else:
raise TypeError(
"config_or_marker must be a config mapping, string, or None, got "
f"{type(config_or_marker).__name__}"
)
resolved_config = (
_validate_model_architecture(configured)
if configured is not None
else None
)
detected = (
detect_model_architecture_from_state_dict(state_dict)
if state_dict is not None
else None
)
if resolved_config is None:
raise ValueError(
"model_architecture is required; expected one of "
f"{SUPPORTED_MODEL_ARCHITECTURES}."
)
if detected is not None and resolved_config != detected:
raise ValueError(
"Configured model architecture conflicts with checkpoint: "
f"config={resolved_config!r}, checkpoint={detected!r}."
)
return resolved_config

22
models.py
View File

@@ -6,11 +6,12 @@ import torch.nn.functional as F
from backbones import (
AgeSinusoidalEncoding,
GPTBlock,
GaussianRBFTimeBasis,
TimeRoPE,
TokenAutoDiscretization,
build_backbone_block,
)
from model_architectures import resolve_model_architecture
from targets import PAD_IDX
@@ -147,8 +148,7 @@ class DeepHealth(nn.Module):
vocab_size: int,
n_embd: int,
n_head: int,
n_hist_layer: int,
n_tab_layer: int,
n_layer: int,
n_types: int,
n_cont_types: int,
n_categories: int,
@@ -160,6 +160,7 @@ class DeepHealth(nn.Module):
dist_mode: str = "exponential", # "exponential", "weibull" or "mixed"
extra_pool_reduce: str = "mean",
dropout: float = 0.0,
model_architecture: str | None = None,
):
super().__init__()
if target_mode not in ["next_token", "all_future"]:
@@ -173,6 +174,9 @@ class DeepHealth(nn.Module):
"dist_mode must be either 'exponential', 'weibull' or 'mixed'")
if extra_pool_reduce not in {"mean", "sum"}:
raise ValueError("extra_pool_reduce must be either 'mean' or 'sum'")
if n_layer < 1:
raise ValueError(f"n_layer must be >= 1, got {n_layer}")
model_architecture = resolve_model_architecture(model_architecture)
self.token_embedding = nn.Embedding(vocab_size, n_embd, padding_idx=0)
self.gender_embedding = nn.Embedding(
2, n_embd) # Assuming binary gender
@@ -189,6 +193,8 @@ class DeepHealth(nn.Module):
self.time_mode = time_mode
self.dist_mode = dist_mode
self.extra_pool_reduce = extra_pool_reduce
self.model_architecture = model_architecture
self.n_layer = n_layer
self.n_embd = n_embd
self.vocab_size = vocab_size
nn.init.normal_(self.token_embedding.weight, mean=0.0, std=0.02)
@@ -208,26 +214,28 @@ class DeepHealth(nn.Module):
if time_mode == "absolute":
self.age_encoding = AgeSinusoidalEncoding(n_embd)
self.blocks = nn.ModuleList([
GPTBlock(
build_backbone_block(
model_architecture,
n_embd=n_embd,
n_head=n_head,
use_time_rope=False,
use_rbf_bias=False,
mlp_dropout=dropout,
) for _ in range(n_hist_layer)
) for _ in range(n_layer)
])
self.rope = None
self.rbf = None
elif time_mode == "relative":
self.age_encoding = None
self.blocks = nn.ModuleList([
GPTBlock(
build_backbone_block(
model_architecture,
n_embd=n_embd,
n_head=n_head,
use_time_rope=True,
use_rbf_bias=True,
mlp_dropout=dropout,
) for _ in range(n_hist_layer)
) for _ in range(n_layer)
])
self.rope = TimeRoPE(n_embd // n_head)
self.rbf = GaussianRBFTimeBasis(n_bases=16, max_time_diff=40.0)

59
run_missing_evaluations.sh
View File

@@ -1,10 +1,11 @@
#!/usr/bin/env bash
set -euo pipefail
# Run all non-wrapper evaluation scripts for every completed experiment under
# runs/. The script is written for Linux servers with bash 4.2.
# Run all non-wrapper evaluation scripts for every completed current-format
# experiment under runs/. The script is written for Linux servers with bash 4.2.
cd "$(dirname "${BASH_SOURCE[0]}")"
shopt -s globstar nullglob
PYTHON_BIN="${PYTHON_BIN:-python}"
DEVICE="${DEVICE:-cuda}"
@@ -81,6 +82,20 @@ run_dir_result_if_missing() {
run_command "$@"
}
run_file_result_if_missing() {
local label="$1"
local result_dir="$2"
local required="$3"
shift 3
if [[ -s "${result_dir}/${required}" ]]; then
echo " skip ${label}: found ${result_dir}/${required}"
return 0
fi
run_command "$@"
}
run_has_extra_info() {
"${PYTHON_BIN}" - "$1" <<'PY'
import json
@@ -115,8 +130,30 @@ raise SystemExit(0 if mode == "all_future" else 1)
PY
}
for run_path in runs/*; do
[[ -d "${run_path}" ]] || continue
run_has_current_model_config() {
"${PYTHON_BIN}" - "$1" <<'PY'
import json
import sys
from pathlib import Path
cfg_path = Path(sys.argv[1]) / "train_config.json"
try:
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
n_layer = int(cfg.get("n_layer", 0))
except Exception:
raise SystemExit(1)
supported = {"transformer_ffn_v1", "traj_mixer_v5"}
raise SystemExit(
0
if cfg.get("model_architecture") in supported and n_layer >= 1
else 1
)
PY
}
for config_path in runs/**/train_config.json; do
run_path="${config_path%/train_config.json}"
echo "==> ${run_path}"
if [[ ! -f "${run_path}/train_config.json" ]]; then
@@ -127,6 +164,10 @@ for run_path in runs/*; do
echo " skip run: missing best_model.pt"
continue
fi
if ! run_has_current_model_config "${run_path}"; then
echo " skip run: config lacks current model_architecture/n_layer fields"
continue
fi
common=()
while IFS= read -r arg; do common+=("${arg}"); done < <(common_args_with_device "${run_path}")
@@ -137,18 +178,16 @@ for run_path in runs/*; do
cpu_reduce_extra=()
while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)
run_dir_result_if_missing \
run_file_result_if_missing \
"evaluate_auc.py" \
"${run_path}" \
"df_both.csv" \
"df_auc_unpooled.csv" \
"df_auc_delphi2m_report.csv" \
"${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"
run_dir_result_if_missing \
run_file_result_if_missing \
"evaluate_auc_v2.py" \
"${run_path}" \
"df_auc_landmark.csv" \
"df_auc_landmark_unpooled.csv" \
"df_auc_landmark_delphi2m_report.csv" \
"${PYTHON_BIN}" evaluate_auc_v2.py "${common[@]}" "${auc_extra[@]}"
if ! run_is_all_future "${run_path}"; then

27
run_missing_training_runs.sh
View File

@@ -10,7 +10,8 @@ set -euo pipefail
# all_future + relative time + mixed death/risk head
#
# This script only launches those missing training jobs. It intentionally does
# not call evaluate_*.py and does not add extra random seeds.
# not call evaluate_*.py and does not add extra random seeds. Set
# MODEL_ARCHITECTURE=traj_mixer_v5 to run the TrajMixer variant.
cd "$(dirname "${BASH_SOURCE[0]}")"
@@ -18,6 +19,8 @@ PYTHON_BIN="${PYTHON_BIN:-python}"
DEVICE="${DEVICE:-cuda}"
NUM_WORKERS="${NUM_WORKERS:-4}"
PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"
MODEL_ARCHITECTURE="${MODEL_ARCHITECTURE:-transformer_ffn_v1}"
N_LAYER="${N_LAYER:-12}"
TIME_MODE="relative"
DIST_MODE="mixed"
@@ -36,8 +39,8 @@ COMMON_ARGS=(
--min_future_events 1
--n_embd 120
--n_head 10
--n_hist_layer 12
--n_tab_layer 4
--n_layer "${N_LAYER}"
--model_architecture "${MODEL_ARCHITECTURE}"
--n_bins 16
--extra_pool_reduce mean
--dropout 0.0
@@ -57,15 +60,23 @@ COMMON_ARGS=(
already_trained() {
local extra_file="$1"
"${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" <<'PY'
"${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" "$MODEL_ARCHITECTURE" "$N_LAYER" <<'PY'
import json
import sys
from pathlib import Path
time_mode, dist_mode, extra_file, seed, validation_query_seed = sys.argv[1:6]
(
time_mode,
dist_mode,
extra_file,
seed,
validation_query_seed,
model_architecture,
n_layer,
) = sys.argv[1:8]
extra_name = Path(extra_file).name
for config_path in Path("runs").glob("*/train_config.json"):
for config_path in Path("runs").rglob("train_config.json"):
try:
cfg = json.loads(config_path.read_text(encoding="utf-8"))
except Exception:
@@ -78,6 +89,8 @@ for config_path in Path("runs").glob("*/train_config.json"):
if (
cfg.get("model_target_mode") == "all_future"
and cfg.get("model_architecture") == model_architecture
and int(cfg.get("n_layer", -1)) == int(n_layer)
and cfg.get("time_mode") == time_mode
and cfg.get("dist_mode") == dist_mode
and Path(str(cfg.get("extra_info_types_file", ""))).name == extra_name
@@ -100,7 +113,7 @@ train_if_missing() {
return 2
fi
echo "==> Checking ${label}: ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
echo "==> Checking ${label}: ${MODEL_ARCHITECTURE} n_layer=${N_LAYER} ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
if existing_run="$(already_trained "$extra_file")"; then
echo " skip: already trained at ${existing_run}"
return 0

247
test_model_architectures.py Normal file
View File

@@ -0,0 +1,247 @@
import unittest
import torch
from backbones import (
SwiGLU,
TrajMixer,
TrajMixerBlock,
TransformerFFNBlock,
build_backbone_block,
)
from model_architectures import (
TRAJ_MIXER_ARCHITECTURE,
TRANSFORMER_FFN_ARCHITECTURE,
detect_model_architecture_from_state_dict,
resolve_model_architecture,
)
from models import DeepHealth
def _build_block(model_architecture: str):
return build_backbone_block(
model_architecture,
n_embd=12,
n_head=3,
use_time_rope=False,
use_rbf_bias=False,
mlp_dropout=0.0,
)
def _as_model_state_dict(block: torch.nn.Module) -> dict[str, torch.Tensor]:
return {
f"blocks.0.{name}": value.detach().clone()
for name, value in block.state_dict().items()
}
def _build_model(
model_architecture: str | None,
*,
n_layer: int = 1,
) -> DeepHealth:
return DeepHealth(
vocab_size=8,
n_embd=12,
n_head=3,
n_layer=n_layer,
n_types=2,
n_cont_types=0,
n_categories=2,
cont_type_ids=[],
time_mode="absolute",
model_architecture=model_architecture,
)
class ModelArchitectureFactoryTest(unittest.TestCase):
def test_factory_builds_both_architectures_with_expected_topology(self) -> None:
ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
self.assertIsInstance(ffn_block, TransformerFFNBlock)
self.assertIsInstance(ffn_block.mlp, SwiGLU)
self.assertTrue(hasattr(ffn_block, "ln1"))
self.assertTrue(hasattr(ffn_block, "ln2"))
traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
self.assertIsInstance(traj_block, TrajMixerBlock)
self.assertIsInstance(traj_block.mlp, TrajMixer)
self.assertTrue(hasattr(traj_block, "ln1"))
self.assertFalse(hasattr(traj_block, "ln2"))
def test_both_architectures_forward_and_backward(self) -> None:
for architecture in (
TRANSFORMER_FFN_ARCHITECTURE,
TRAJ_MIXER_ARCHITECTURE,
):
with self.subTest(architecture=architecture):
torch.manual_seed(0)
block = _build_block(architecture)
x = torch.randn(2, 5, 12, requires_grad=True)
output = block(x)
self.assertEqual(output.shape, x.shape)
output.square().mean().backward()
self.assertIsNotNone(x.grad)
self.assertTrue(torch.isfinite(x.grad).all())
self.assertGreater(x.grad.abs().sum().item(), 0.0)
self.assertIsNotNone(block.attn.qkv.weight.grad)
self.assertGreater(
block.attn.qkv.weight.grad.abs().sum().item(),
0.0,
)
if architecture == TRANSFORMER_FFN_ARCHITECTURE:
branch_parameters = (
block.mlp.w1.weight,
block.mlp.w2.weight,
block.mlp.w3.weight,
)
else:
branch_parameters = (
block.mlp.intra_gate_proj,
block.mlp.intra_value_proj,
block.mlp.output_proj,
)
for parameter in branch_parameters:
self.assertIsNotNone(parameter.grad)
self.assertTrue(torch.isfinite(parameter.grad).all())
self.assertGreater(parameter.grad.abs().sum().item(), 0.0)
def test_unknown_architecture_is_rejected(self) -> None:
with self.assertRaises(ValueError):
_build_block("unknown_architecture")
with self.assertRaisesRegex(ValueError, "model_architecture is required"):
_build_model(None)
def test_deephealth_rejects_fewer_than_one_layer(self) -> None:
for n_layer in (0, -1):
with self.subTest(n_layer=n_layer):
with self.assertRaisesRegex(ValueError, "n_layer must be >= 1"):
_build_model(
TRANSFORMER_FFN_ARCHITECTURE,
n_layer=n_layer,
)
def test_deephealth_uses_factory_and_strictly_reloads_both_models(self) -> None:
for architecture, block_class in (
(TRANSFORMER_FFN_ARCHITECTURE, TransformerFFNBlock),
(TRAJ_MIXER_ARCHITECTURE, TrajMixerBlock),
):
with self.subTest(architecture=architecture):
model = _build_model(architecture)
self.assertEqual(model.model_architecture, architecture)
self.assertIsInstance(model.blocks[0], block_class)
self.assertEqual(
detect_model_architecture_from_state_dict(
model.state_dict()
),
architecture,
)
reloaded = _build_model(architecture)
incompatible = reloaded.load_state_dict(
model.state_dict(),
strict=True,
)
self.assertEqual(incompatible.missing_keys, [])
self.assertEqual(incompatible.unexpected_keys, [])
class ModelArchitectureResolutionTest(unittest.TestCase):
def setUp(self) -> None:
self.ffn_block = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
self.traj_block = _build_block(TRAJ_MIXER_ARCHITECTURE)
self.ffn_state = _as_model_state_dict(self.ffn_block)
self.traj_state = _as_model_state_dict(self.traj_block)
def test_state_dict_detection_recognizes_both_architectures(self) -> None:
self.assertEqual(
detect_model_architecture_from_state_dict(self.ffn_state),
TRANSFORMER_FFN_ARCHITECTURE,
)
self.assertEqual(
detect_model_architecture_from_state_dict(self.traj_state),
TRAJ_MIXER_ARCHITECTURE,
)
def test_explicit_markers_resolve_when_checkpoint_matches(self) -> None:
for architecture, state_dict in (
(TRANSFORMER_FFN_ARCHITECTURE, self.ffn_state),
(TRAJ_MIXER_ARCHITECTURE, self.traj_state),
):
with self.subTest(architecture=architecture):
self.assertEqual(
resolve_model_architecture(
{"model_architecture": architecture},
state_dict,
),
architecture,
)
def test_architecture_marker_is_required_for_checkpoint_loading(self) -> None:
with self.assertRaisesRegex(ValueError, "model_architecture is required"):
resolve_model_architecture({}, self.ffn_state)
with self.assertRaisesRegex(ValueError, "model_architecture is required"):
resolve_model_architecture(None, self.traj_state)
def test_explicit_marker_conflicting_with_state_dict_is_rejected(self) -> None:
conflicts = (
(TRANSFORMER_FFN_ARCHITECTURE, self.traj_state),
(TRAJ_MIXER_ARCHITECTURE, self.ffn_state),
)
for architecture, state_dict in conflicts:
with self.subTest(architecture=architecture):
with self.assertRaises(ValueError):
resolve_model_architecture(
{"model_architecture": architecture},
state_dict,
)
def test_unknown_marker_and_ambiguous_state_dict_are_rejected(self) -> None:
with self.assertRaises(ValueError):
resolve_model_architecture(
{"model_architecture": "traj_mixer_v4"}
)
ambiguous_state = dict(self.ffn_state)
ambiguous_state.update(self.traj_state)
with self.assertRaises(ValueError):
detect_model_architecture_from_state_dict(ambiguous_state)
with self.assertRaises(ValueError):
detect_model_architecture_from_state_dict(
{"token_embedding.weight": torch.empty(2, 2)}
)
def test_ffn_block_schema_is_stable_and_strictly_loadable(self) -> None:
expected_keys = {
"attn.time_bias_scale",
"attn.qkv.weight",
"attn.out_proj.weight",
"attn.rbf_proj.weight",
"mlp.w1.weight",
"mlp.w1.bias",
"mlp.w2.weight",
"mlp.w2.bias",
"mlp.w3.weight",
"mlp.w3.bias",
"ln1.weight",
"ln1.bias",
"ln2.weight",
"ln2.bias",
}
state = self.ffn_block.state_dict()
self.assertSetEqual(set(state), expected_keys)
self.assertEqual(tuple(state["mlp.w1.weight"].shape), (30, 12))
self.assertEqual(tuple(state["mlp.w3.weight"].shape), (12, 30))
reloaded = _build_block(TRANSFORMER_FFN_ARCHITECTURE)
incompatible = reloaded.load_state_dict(state, strict=True)
self.assertEqual(incompatible.missing_keys, [])
self.assertEqual(incompatible.unexpected_keys, [])
if __name__ == "__main__":
unittest.main()

166
test_traj_mixer.py Normal file
View File

@@ -0,0 +1,166 @@
import unittest
import torch
from backbones import TrajMixer
class TrajMixerTest(unittest.TestCase):
def test_default_shape_parameters_and_initialization(self) -> None:
mixer = TrajMixer(
n_embd=120,
n_head=10,
dropout=0.0,
)
x = torch.randn(2, 7, 120)
self.assertEqual(mixer(x).shape, x.shape)
self.assertEqual(sum(p.numel() for p in mixer.parameters()), 32_040)
self.assertEqual(tuple(mixer.norm.normalized_shape), (120,))
self.assertEqual(tuple(mixer.intra_gate_logits.shape), (10, 12))
torch.testing.assert_close(
torch.sigmoid(mixer.intra_gate_logits.detach()),
torch.full((10, 12), 0.1),
)
self.assertEqual(mixer.intra_hidden, 48)
self.assertEqual(
tuple(mixer.intra_gate_proj.shape),
(10, 12, 48),
)
self.assertEqual(
tuple(mixer.intra_value_proj.shape),
(10, 12, 48),
)
self.assertEqual(
tuple(mixer.intra_output_proj.shape),
(10, 48, 12),
)
self.assertEqual(mixer.hidden_group, 40)
self.assertEqual(tuple(mixer.gate_proj.shape), (12, 10, 40))
self.assertEqual(tuple(mixer.value_proj.shape), (12, 10, 40))
self.assertEqual(tuple(mixer.output_proj.shape), (12, 40, 10))
def test_zero_final_output_projection_makes_mixer_identity(self) -> None:
torch.manual_seed(0)
mixer = TrajMixer(12, n_head=3, dropout=0.0)
with torch.no_grad():
mixer.output_proj.zero_()
x = torch.randn(2, 5, 12)
torch.testing.assert_close(mixer(x), x)
def test_forward_matches_single_outer_residual_formula(self) -> None:
torch.manual_seed(0)
mixer = TrajMixer(12, n_head=3, dropout=0.0)
mixer.eval()
x = torch.randn(2, 5, 12)
grouped = mixer.norm(x).reshape(2, 5, 3, 4)
intra_output = mixer._intra_mix(grouped)
static_gate = torch.sigmoid(mixer.intra_gate_logits).view(
1, 1, 3, 4
)
mixed_input = grouped + static_gate * intra_output
update = mixer._cross_mix(mixed_input).reshape(2, 5, 12)
torch.testing.assert_close(mixer(x), x + update)
def test_intra_stage_is_independent_across_groups(self) -> None:
torch.manual_seed(0)
mixer = TrajMixer(12, n_head=3, dropout=0.0)
mixer.eval()
grouped = torch.randn(2, 4, 3, 4)
changed = grouped.clone()
changed[:, :, 1, :] += torch.randn_like(changed[:, :, 1, :])
original_out = mixer._intra_mix(grouped)
changed_out = mixer._intra_mix(changed)
unchanged_groups = torch.tensor([0, 2])
torch.testing.assert_close(
original_out.index_select(2, unchanged_groups),
changed_out.index_select(2, unchanged_groups),
)
def test_cross_stage_mixes_groups_without_mixing_coordinates(self) -> None:
mixer = TrajMixer(6, n_head=3, dropout=0.0)
mixer.eval()
with torch.no_grad():
mixer.gate_proj.zero_()
mixer.value_proj.zero_()
mixer.output_proj.zero_()
# Coordinate 0 reads group 0 through hidden unit 0 and writes it
# into group 1. Coordinate 1 must remain independent.
mixer.gate_proj[0, 0, 0] = 1.0
mixer.value_proj[0, 0, 0] = 1.0
mixer.output_proj[0, 0, 1] = 1.0
grouped = torch.tensor(
[[[
[-1.0, 4.0],
[0.0, 5.0],
[1.0, 6.0],
]]]
)
changed = grouped.clone()
changed[0, 0, 0, 0] = 2.0
original_out = mixer._cross_mix(grouped)
changed_out = mixer._cross_mix(changed)
self.assertNotEqual(
original_out[0, 0, 1, 0].item(),
changed_out[0, 0, 1, 0].item(),
)
torch.testing.assert_close(
original_out[..., 1],
changed_out[..., 1],
)
def test_mixer_does_not_mix_sequence_positions(self) -> None:
torch.manual_seed(0)
mixer = TrajMixer(12, n_head=3, dropout=0.0)
mixer.eval()
x = torch.randn(2, 5, 12)
changed = x.clone()
changed[:, 3, :] += torch.randn_like(changed[:, 3, :])
original_out = mixer(x)
changed_out = mixer(changed)
unchanged_positions = torch.tensor([0, 1, 2, 4])
torch.testing.assert_close(
original_out.index_select(1, unchanged_positions),
changed_out.index_select(1, unchanged_positions),
)
def test_gradients_reach_every_projection_family(self) -> None:
torch.manual_seed(1)
mixer = TrajMixer(12, n_head=3, dropout=0.0)
x = torch.randn(2, 4, 12, requires_grad=True)
mixer(x).square().mean().backward()
self.assertIsNotNone(x.grad)
self.assertTrue(torch.isfinite(x.grad).all())
self.assertGreater(x.grad.abs().sum().item(), 0.0)
projection_names = (
"intra_gate_proj",
"intra_value_proj",
"intra_output_proj",
"gate_proj",
"value_proj",
"output_proj",
)
for name in projection_names:
parameter = getattr(mixer, name)
self.assertIsNotNone(parameter.grad, name)
self.assertTrue(torch.isfinite(parameter.grad).all(), name)
self.assertGreater(parameter.grad.abs().sum().item(), 0.0, name)
def test_invalid_group_partition_is_rejected(self) -> None:
with self.assertRaisesRegex(ValueError, "divisible"):
TrajMixer(n_embd=121, n_head=10)
if __name__ == "__main__":
unittest.main()

36
train_all_future.py
View File

@@ -27,12 +27,17 @@ from tqdm.auto import tqdm
from dataset import AllFutureHealthDataset, all_future_collate_fn
from losses import build_loss
from model_architectures import (
DEFAULT_MODEL_ARCHITECTURE,
SUPPORTED_MODEL_ARCHITECTURES,
)
from models import DeepHealth
from targets import CHECKUP_IDX, PAD_IDX
from train_util import (
configure_torch_for_training,
create_unique_run_dir,
format_extra_info_types,
get_model_parameter_counts,
load_extra_info_types_file,
resolve_device,
save_checkpoint,
@@ -64,6 +69,7 @@ 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("--runs_root", type=str, default="runs")
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--extra_info_types_file", type=str, default=None)
@@ -79,8 +85,7 @@ 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_layer", type=int, default=12)
parser.add_argument("--n_bins", type=int, default=16)
parser.add_argument("--extra_pool_reduce", type=str, default="mean",
choices=["mean", "sum"])
@@ -89,6 +94,12 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--dist_mode", type=str, default="exponential",
choices=["exponential", "weibull", "mixed"])
parser.add_argument("--dropout", type=float, default=0.0)
parser.add_argument(
"--model_architecture",
type=str,
default=DEFAULT_MODEL_ARCHITECTURE,
choices=SUPPORTED_MODEL_ARCHITECTURES,
)
parser.add_argument("--batch_size", type=int, default=128)
parser.add_argument("--base_lr", type=float, default=3e-4)
@@ -147,8 +158,7 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
vocab_size=dataset.vocab_size,
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_layer=args.n_layer,
n_types=dataset.n_types,
n_cont_types=dataset.n_cont_types,
n_categories=dataset.n_categories,
@@ -159,6 +169,7 @@ def build_model(args: argparse.Namespace, dataset: AllFutureHealthDataset) -> De
time_mode=args.time_mode,
dist_mode=args.dist_mode,
dropout=args.dropout,
model_architecture=args.model_architecture,
)
@@ -298,6 +309,7 @@ def build_metadata(
"dataset_class": "AllFutureHealthDataset",
"collate_fn": "all_future_collate_fn",
"model_class": "DeepHealth",
"model_architecture": args.model_architecture,
"model_target_mode": "all_future",
"target_mode": "all_future",
"dist_mode": args.dist_mode,
@@ -335,12 +347,14 @@ def main() -> None:
configure_torch_for_training(device)
run_dir, run_name = create_unique_run_dir(
lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}"
lambda timestamp: f"{args.time_mode}_{args.dist_mode}_all_future_pure_disease_{timestamp}",
runs_root=Path(args.runs_root) / args.model_architecture,
)
logger = setup_logging(run_dir)
logger.info(f"Starting all-future training run: {run_name}")
logger.info(f"Device: {device}")
logger.info(f"Model architecture: {args.model_architecture}")
logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
logger.info("Loading all-future datasets...")
@@ -434,6 +448,12 @@ def main() -> None:
)
model = build_model(args, train_dataset).to(device)
parameter_counts = get_model_parameter_counts(model)
logger.info(
"Model parameters: "
f"total={parameter_counts['model_parameter_count']:,}, "
f"trainable={parameter_counts['trainable_parameter_count']:,}"
)
optimizer = AdamW(
model.parameters(),
lr=args.base_lr,
@@ -443,10 +463,14 @@ def main() -> None:
criterion = build_criterion(args, train_dataset)
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
train_metadata = build_metadata(
args, train_dataset, run_name, train_subset, val_subset, test_subset
)
train_metadata.update(parameter_counts)
save_config(
args,
run_dir / "train_config.json",
extra=build_metadata(args, train_dataset, run_name, train_subset, val_subset, test_subset),
extra=train_metadata,
)
best_val = float("inf")

36
train_next_step.py
View File

@@ -24,6 +24,10 @@ from tqdm.auto import tqdm
from dataset import HealthDataset, collate_fn
from losses import build_loss
from model_architectures import (
DEFAULT_MODEL_ARCHITECTURE,
SUPPORTED_MODEL_ARCHITECTURES,
)
from models import DeepHealth, DeepHealthOutput
from readouts import build_readout
from targets import CHECKUP_IDX, NO_EVENT_IDX, PAD_IDX
@@ -31,6 +35,7 @@ from train_util import (
configure_torch_for_training,
create_unique_run_dir,
format_extra_info_types,
get_model_parameter_counts,
load_extra_info_types_file,
resolve_device,
save_checkpoint,
@@ -61,6 +66,7 @@ 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("--runs_root", type=str, default="runs")
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)
@@ -75,14 +81,19 @@ 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_layer", type=int, default=12)
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)
parser.add_argument(
"--model_architecture",
type=str,
default=DEFAULT_MODEL_ARCHITECTURE,
choices=SUPPORTED_MODEL_ARCHITECTURES,
)
parser.add_argument("--target_mode", type=str, default="uts",
choices=["delphi2m", "uts"])
@@ -152,8 +163,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
vocab_size=dataset.vocab_size,
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_layer=args.n_layer,
n_types=dataset.n_types,
n_cont_types=dataset.n_cont_types,
n_categories=dataset.n_categories,
@@ -164,6 +174,7 @@ def build_model(args: argparse.Namespace, dataset: HealthDataset) -> DeepHealth:
time_mode=args.time_mode,
dist_mode="exponential",
dropout=args.dropout,
model_architecture=args.model_architecture,
)
@@ -484,6 +495,7 @@ def build_metadata(
"dataset_class": "NextStepHealthDataset",
"collate_fn": "next_step_collate_fn",
"model_class": "DeepHealth",
"model_architecture": args.model_architecture,
"model_target_mode": "next_token",
"target_mode": args.target_mode,
"dist_mode": "exponential",
@@ -521,12 +533,14 @@ def main() -> None:
lambda timestamp: (
f"{args.time_mode}_exponential_next_token_{args.target_mode}_"
f"gap_{args.no_event_interval_years:g}y_{timestamp}"
)
),
runs_root=Path(args.runs_root) / args.model_architecture,
)
logger = setup_logging(run_dir)
logger.info(f"Starting next-step training run: {run_name}")
logger.info(f"Device: {device}")
logger.info(f"Model architecture: {args.model_architecture}")
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}")
@@ -596,6 +610,12 @@ def main() -> None:
)
model = build_model(args, dataset).to(device)
parameter_counts = get_model_parameter_counts(model)
logger.info(
"Model parameters: "
f"total={parameter_counts['model_parameter_count']:,}, "
f"trainable={parameter_counts['trainable_parameter_count']:,}"
)
readout = build_next_step_readout(args).to(device)
criterion = build_next_step_loss(args)
optimizer = AdamW(
@@ -606,10 +626,14 @@ def main() -> None:
)
adaptive_lr = args.base_lr * math.sqrt(args.batch_size / 128)
train_metadata = build_metadata(
args, dataset, run_name, train_subset, val_subset, test_subset
)
train_metadata.update(parameter_counts)
save_config(
args,
run_dir / "train_config.json",
extra=build_metadata(args, dataset, run_name, train_subset, val_subset, test_subset),
extra=train_metadata,
)
best_val = float("inf")

14
train_util.py
View File

@@ -300,6 +300,20 @@ def build_optimizer(args: Any, model: DeepHealth) -> AdamW:
)
def get_model_parameter_counts(model: torch.nn.Module) -> Dict[str, int]:
"""Return stable total and trainable parameter counts."""
return {
"model_parameter_count": sum(
parameter.numel() for parameter in model.parameters()
),
"trainable_parameter_count": sum(
parameter.numel()
for parameter in model.parameters()
if parameter.requires_grad
),
}
def set_optimizer_lr(optimizer: AdamW, lr: float) -> None:
for param_group in optimizer.param_groups:
param_group["lr"] = lr