feat: add person fusion screening variants
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user