feat: add person fusion screening variants

This commit is contained in:
Xin Wang
2026-09-06 21:48:32 +08:00
parent 8e70281330
commit b91a24f308
8 changed files with 447 additions and 47 deletions

View File

@@ -25,6 +25,7 @@ from easycom_head import (
DATASET_NAMES,
LAYER_FUSION_MODES,
LOCAL_FUSION_MODES,
LOCAL_FUSION_SCOPES,
LOSS_NAMES,
NUM_HORIZONS,
OFFICIAL_TRAIN_SESSIONS,
@@ -64,7 +65,11 @@ def parse_args() -> argparse.Namespace:
"--local-fusion",
choices=LOCAL_FUSION_MODES,
default="pooled",
help="pooled keeps the existing baseline; frame-set-attention queries the latest causal person set.",
help=(
"pooled keeps the existing baseline; frame-set-attention is F1; "
"audio-person-frame-set and gated-person-state-adapter are "
"current-sidecar F5/F6 screening variants, not full F5/F6."
),
)
parser.add_argument(
"--local-projection-dim", type=int, default=256,
@@ -179,12 +184,12 @@ def _validate_args(args: argparse.Namespace) -> None:
raise ValueError("Person attention dimension/heads must be positive")
if args.person_attention_dim % args.person_attention_heads:
raise ValueError("--person-attention-dim must be divisible by --person-attention-heads")
if args.local_fusion == "frame-set-attention" and (
if args.local_fusion != "pooled" and (
args.local_feature_root is None or args.local_modalities != "av"
):
raise ValueError(
"--local-fusion frame-set-attention requires --local-feature-root "
"and --local-modalities av"
"Person-set --local-fusion modes require --local-feature-root and "
"--local-modalities av"
)
if args.epochs <= 0:
raise ValueError("--epochs must be positive")
@@ -505,7 +510,7 @@ def _checkpoint(
history: list[dict[str, Any]],
) -> dict[str, Any]:
return {
"format_version": 5,
"format_version": 6,
"head": head.state_dict(),
"head_config": {
"hidden_size": hidden_size,
@@ -567,6 +572,7 @@ def _checkpoint(
"local_feature_root": str(args.local_feature_root) if args.local_feature_root else None,
"local_modalities": args.local_modalities if args.local_feature_root else None,
"local_fusion": args.local_fusion,
"local_fusion_scope": LOCAL_FUSION_SCOPES[args.local_fusion],
"local_projection_dim": args.local_projection_dim,
"person_attention_dim": args.person_attention_dim,
"person_attention_heads": args.person_attention_heads,
@@ -807,6 +813,10 @@ def main() -> None:
f"fusion_parameters={fusion_parameters:,}, "
f"head_trainable_parameters={trainable_parameters:,}"
)
print(
"local_fusion_scope="
f"{LOCAL_FUSION_SCOPES[args.local_fusion]}"
)
if head.layer_logits is None:
optimizer_parameters: Any = head.parameters()
else:
@@ -967,6 +977,9 @@ def main() -> None:
record["layer_mixture_weights_oldest_to_newest"] = (
layer_weights.detach().float().cpu().tolist()
)
interaction_gates = head.person_interaction_gates()
if interaction_gates:
record["person_interaction_gates"] = interaction_gates
weight_text = (
""
if layer_weights is None
@@ -987,6 +1000,12 @@ def main() -> None:
f"{record['train_expected_transition_reward']:.6f}"
)
)
gate_text = (
""
if not interaction_gates
else ", person_gates="
+ json.dumps(interaction_gates, sort_keys=True)
)
if dev_sessions:
assert dev_states is not None and dev_targets is not None
@@ -1005,7 +1024,7 @@ def main() -> None:
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
f"dev_valid_mean_mAP={dev_map * 100:.3f}"
f"{transition_text}{weight_text}"
f"{transition_text}{weight_text}{gate_text}"
)
if dev_map > float(best_dev_map):
best_dev_map = dev_map
@@ -1031,7 +1050,7 @@ def main() -> None:
else:
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}"
f"{transition_text}{weight_text}"
f"{transition_text}{weight_text}{gate_text}"
)
history.append(record)