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

Add two-stage TrajMixer mixing

Browse Source
This commit is contained in:
Jiarui Li
2026-07-24 10:51:41 +08:00
parent 85352dae0f
commit 20c99484f3
4 changed files with 401 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
models.py
View File

@@ -15,7 +15,7 @@ from backbones import (
from targets import PAD_IDX
TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v2"
TRAJ_MIXER_ARCHITECTURE = "traj_mixer_v3"
def validate_traj_mixer_config(config: Mapping[str, object]) -> None:
@@ -29,6 +29,13 @@ 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.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.gate_proj",
"blocks.0.mlp.value_proj",