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

Refactor TrajMixer to single residual

Browse Source
This commit is contained in:
Jiarui Li
2026-07-24 15:25:02 +08:00
parent 7b48cb8425
commit 8d0d71292e
4 changed files with 304 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
models.py
View File

@@ -15,7 +15,7 @@ from backbones import (
from targets import PAD_IDX
TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v3"
TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v5"
def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
@@ -29,14 +29,12 @@ def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
def validate_traj_mixer_state_dict(state_dict: Mapping[str, object]) -> None:
required_keys = {
"blocks.0.mlp.intra_norm.weight",
"blocks.0.mlp.intra_norm.bias",
"blocks.0.mlp.norm.weight",
"blocks.0.mlp.norm.bias",
"blocks.0.mlp.intra_gate_proj",
"blocks.0.mlp.intra_value_proj",
"blocks.0.mlp.intra_output_proj",
"blocks.0.mlp.cross_norm.weight",
"blocks.0.mlp.cross_norm.bias",
"blocks.0.mlp.group_align",
"blocks.0.mlp.intra_gate_logits",
"blocks.0.mlp.gate_proj",
"blocks.0.mlp.value_proj",
"blocks.0.mlp.output_proj",