diff --git a/MODEL_ARCHITECTURES.md b/MODEL_ARCHITECTURES.md new file mode 100644 index 0000000..5f7dd15 --- /dev/null +++ b/MODEL_ARCHITECTURES.md @@ -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/ + / + traj_mixer_v5/ + / +``` + +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. diff --git a/backbones.py b/backbones.py index 58d76cd..86e4789 100644 --- a/backbones.py +++ b/backbones.py @@ -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, diff --git a/evaluate_auc.py b/evaluate_auc.py index ce5fd4f..ee0ee4d 100644 --- a/evaluate_auc.py +++ b/evaluate_auc.py @@ -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() diff --git a/evaluate_auc_v2.py b/evaluate_auc_v2.py index 111518e..4c32a3c 100644 --- a/evaluate_auc_v2.py +++ b/evaluate_auc_v2.py @@ -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" diff --git a/evaluate_event_free_survival.py b/evaluate_event_free_survival.py index ef1e7ea..58d3fab 100644 --- a/evaluate_event_free_survival.py +++ b/evaluate_event_free_survival.py @@ -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() diff --git a/evaluate_extra_info_attribution.py b/evaluate_extra_info_attribution.py index cdd8fa9..b521287 100644 --- a/evaluate_extra_info_attribution.py +++ b/evaluate_extra_info_attribution.py @@ -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() diff --git a/evaluate_single_disease_mortality_attribution.py b/evaluate_single_disease_mortality_attribution.py index 1a111b0..b211edb 100644 --- a/evaluate_single_disease_mortality_attribution.py +++ b/evaluate_single_disease_mortality_attribution.py @@ -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() diff --git a/export_tquery_logits_hidden.py b/export_tquery_logits_hidden.py index 692e805..b81bc01 100644 --- a/export_tquery_logits_hidden.py +++ b/export_tquery_logits_hidden.py @@ -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() diff --git a/export_weibull_death_parameter_stats.py b/export_weibull_death_parameter_stats.py index dd1eef2..bfdb1bd 100644 --- a/export_weibull_death_parameter_stats.py +++ b/export_weibull_death_parameter_stats.py @@ -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() diff --git a/model_architectures.py b/model_architectures.py new file mode 100644 index 0000000..330cb43 --- /dev/null +++ b/model_architectures.py @@ -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.`` 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..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 diff --git a/models.py b/models.py index d8b6e4e..d6104dc 100644 --- a/models.py +++ b/models.py @@ -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) diff --git a/run_missing_evaluations.sh b/run_missing_evaluations.sh index b6faf43..e711956 100644 --- a/run_missing_evaluations.sh +++ b/run_missing_evaluations.sh @@ -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 diff --git a/run_missing_training_runs.sh b/run_missing_training_runs.sh index 5aabb91..f14a87b 100755 --- a/run_missing_training_runs.sh +++ b/run_missing_training_runs.sh @@ -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 diff --git a/test_model_architectures.py b/test_model_architectures.py new file mode 100644 index 0000000..729c528 --- /dev/null +++ b/test_model_architectures.py @@ -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() diff --git a/test_traj_mixer.py b/test_traj_mixer.py new file mode 100644 index 0000000..4f886e3 --- /dev/null +++ b/test_traj_mixer.py @@ -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() diff --git a/train_all_future.py b/train_all_future.py index 9e1068f..9a861f0 100644 --- a/train_all_future.py +++ b/train_all_future.py @@ -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") diff --git a/train_next_step.py b/train_next_step.py index 35f86ff..bf242a6 100644 --- a/train_next_step.py +++ b/train_next_step.py @@ -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") diff --git a/train_util.py b/train_util.py index 1fe578c..5c7a806 100644 --- a/train_util.py +++ b/train_util.py @@ -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