feat: add causal frame set attention fusion

train_turn_head.py

@@ -24,6 +24,7 @@ from easycom_head import (
     CLASS_NAMES,
     DATASET_NAMES,
     LAYER_FUSION_MODES,
+    LOCAL_FUSION_MODES,
     LOSS_NAMES,
     NUM_HORIZONS,
     OFFICIAL_TRAIN_SESSIONS,
@@ -46,6 +47,7 @@ from easycom_head import (
     transition_targets_from_activity,
     validate_disjoint_sessions,
 )
+from local_turn_features import FramePersonFeatures
 
 
 def parse_args() -> argparse.Namespace:
@@ -58,10 +60,18 @@ def parse_args() -> argparse.Namespace:
         help="Optional frozen local-feature sidecars; omit for the original baseline.",
     )
     parser.add_argument("--local-modalities", choices=("visual", "audio", "av"), default="av")
+    parser.add_argument(
+        "--local-fusion",
+        choices=LOCAL_FUSION_MODES,
+        default="pooled",
+        help="pooled keeps the existing baseline; frame-set-attention queries the latest causal person set.",
+    )
     parser.add_argument(
         "--local-projection-dim", type=int, default=256,
         help="Trainable projection per local modality; 0 retains encoder dimensions.",
     )
+    parser.add_argument("--person-attention-dim", type=int, default=256)
+    parser.add_argument("--person-attention-heads", type=int, default=4)
     parser.add_argument(
         "--dataset",
         choices=DATASET_NAMES,
@@ -165,6 +175,17 @@ def parse_args() -> argparse.Namespace:
 def _validate_args(args: argparse.Namespace) -> None:
     if args.local_projection_dim < 0:
         raise ValueError("--local-projection-dim must be non-negative")
+    if args.person_attention_dim <= 0 or args.person_attention_heads <= 0:
+        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 (
+        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"
+        )
     if args.epochs <= 0:
         raise ValueError("--epochs must be positive")
     if args.batch_size <= 0:
@@ -225,6 +246,7 @@ def _load_examples(
     local_modalities: str = "av",
     local_contract: dict[str, Any] | None = None,
     compute_transition_targets: bool = True,
+    local_fusion: str = "pooled",
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -233,11 +255,13 @@ def _load_examples(
     int,
     list[dict[str, Any]],
     torch.Tensor,
+    FramePersonFeatures | None,
 ]:
     all_states: list[torch.Tensor] = []
     all_targets: list[torch.Tensor] = []
     all_transition_targets: list[torch.Tensor] = []
     all_local: list[torch.Tensor] = []
+    all_people: list[FramePersonFeatures] = []
     metadata: list[dict[str, Any]] = []
     hidden_size: int | None = None
     num_input_layers: int | None = None
@@ -251,6 +275,7 @@ def _load_examples(
             local_feature_root=local_feature_root,
             local_modalities=local_modalities,
             local_contract=local_contract,
+            local_fusion=local_fusion,
         )
         local_contract = data.feature_metadata.get("local_contract")
         states, targets, future_indices = data.future_examples(NUM_HORIZONS)
@@ -260,6 +285,11 @@ def _load_examples(
             if data.local_features is not None
             else torch.empty(len(states), 0)
         )
+        people = (
+            data.local_person_features.select_rows(current_indices)
+            if data.local_person_features is not None
+            else None
+        )
         if compute_transition_targets:
             activity_sequence = torch.cat(
                 (data.targets[current_indices, None, :], targets),
@@ -286,6 +316,8 @@ def _load_examples(
             targets = targets[keep]
             transition_targets = transition_targets[keep]
             local = local[keep]
+            if people is not None:
+                people = people.select_rows(torch.nonzero(keep).flatten())
             if not len(states):
                 raise ValueError(f"Sparse anchor selection kept no examples for {name}")
         if hidden_size is None:
@@ -306,6 +338,8 @@ def _load_examples(
         all_targets.append(targets)
         all_transition_targets.append(transition_targets)
         all_local.append(local)
+        if people is not None:
+            all_people.append(people)
         metadata.append(
             {
                 "session": name,
@@ -332,6 +366,8 @@ def _load_examples(
         )
     if hidden_size is None or num_input_layers is None:
         raise ValueError("No sessions selected")
+    if all_people and len(all_people) != len(all_states):
+        raise ValueError("Person sets are missing for some sessions")
     return (
         torch.cat(all_states, dim=0),
         torch.cat(all_targets, dim=0),
@@ -340,6 +376,7 @@ def _load_examples(
         num_input_layers,
         metadata,
         torch.cat(all_local, dim=0),
+        FramePersonFeatures.concatenate(all_people) if all_people else None,
     )
 
 
@@ -407,6 +444,7 @@ def _evaluate_valid_horizons(
     batch_size: int,
     loss_name: str,
     local: torch.Tensor | None = None,
+    people: FramePersonFeatures | None = None,
 ) -> dict[str, Any]:
     head.eval()
     chunks: list[torch.Tensor] = []
@@ -420,7 +458,19 @@ def _evaluate_valid_horizons(
             local[start : start + batch_size].to(device=device, dtype=torch.float32)
             if local is not None else None
         )
-        chunks.append(marginal_probabilities_from_logits(head(x, side), loss_name).cpu())
+        person_batch = None
+        person_mask = None
+        if people is not None:
+            rows = torch.arange(start, min(start + batch_size, len(states)))
+            person_batch, person_mask = people.padded(rows)
+            person_batch = person_batch.to(device=device, dtype=torch.float32)
+            person_mask = person_mask.to(device=device)
+        chunks.append(
+            marginal_probabilities_from_logits(
+                head(x, side, person_batch, person_mask),
+                loss_name,
+            ).cpu()
+        )
     probabilities = torch.cat(chunks, dim=0).numpy()
     target_array = targets.numpy()
 
@@ -455,7 +505,7 @@ def _checkpoint(
     history: list[dict[str, Any]],
 ) -> dict[str, Any]:
     return {
-        "format_version": 4,
+        "format_version": 5,
         "head": head.state_dict(),
         "head_config": {
             "hidden_size": hidden_size,
@@ -469,6 +519,9 @@ def _checkpoint(
             "local_visual_dim": head.local_visual_dim,
             "local_audio_dim": head.local_audio_dim,
             "local_projection_dim": head.local_projection_dim,
+            "local_fusion": head.local_fusion,
+            "person_attention_dim": head.person_attention_dim,
+            "person_attention_heads": head.person_attention_heads,
         },
         "local_contract": feature_metadata[0].get("local_contract"),
         "alignment": {
@@ -513,7 +566,10 @@ def _checkpoint(
             "layer_fusion": args.layer_fusion,
             "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_projection_dim": args.local_projection_dim,
+            "person_attention_dim": args.person_attention_dim,
+            "person_attention_heads": args.person_attention_heads,
             "learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
             "supervision_policy": (
                 "transition-sparse-manifest"
@@ -605,6 +661,7 @@ def main() -> None:
         num_input_layers,
         feature_metadata,
         local_features,
+        person_features,
     ) = _load_examples(
         args.feature_root,
         args.target_root,
@@ -615,6 +672,7 @@ def main() -> None:
         args.local_feature_root,
         args.local_modalities,
         compute_transition_targets=args.transition_objective != "none",
+        local_fusion=args.local_fusion,
     )
     local_contract = feature_metadata[0].get("local_contract")
     if not len(states):
@@ -645,6 +703,7 @@ def main() -> None:
     dev_states: torch.Tensor | None = None
     dev_targets: torch.Tensor | None = None
     dev_local: torch.Tensor | None = None
+    dev_people: FramePersonFeatures | None = None
     if dev_sessions:
         (
             dev_states,
@@ -654,6 +713,7 @@ def main() -> None:
             dev_num_input_layers,
             _,
             dev_local,
+            dev_people,
         ) = _load_examples(
             args.feature_root,
             args.target_root,
@@ -665,6 +725,7 @@ def main() -> None:
             args.local_modalities,
             local_contract,
             compute_transition_targets=args.transition_objective != "none",
+            local_fusion=args.local_fusion,
         )
         if dev_hidden_size != hidden_size:
             raise ValueError(
@@ -678,7 +739,13 @@ def main() -> None:
             )
         print(f"development examples: {len(dev_states)}")
 
-    dataset = TensorDataset(states, targets, transition_targets, local_features)
+    dataset = TensorDataset(
+        states,
+        targets,
+        transition_targets,
+        local_features,
+        torch.arange(len(states), dtype=torch.long),
+    )
     generator = torch.Generator().manual_seed(args.seed)
     loader = DataLoader(
         dataset,
@@ -702,13 +769,26 @@ def main() -> None:
         layer_fusion=args.layer_fusion,
         transition_auxiliary=args.transition_objective != "none",
         local_projection_dim=args.local_projection_dim,
+        local_fusion=args.local_fusion,
+        person_attention_dim=args.person_attention_dim,
+        person_attention_heads=args.person_attention_heads,
         **local_head_dimensions(local_contract),
     ).to(device)
     if local_contract is not None:
         print(
-            f"local fusion: {args.local_modalities}; raw V={head.local_visual_dim}, "
-            f"A={head.local_audio_dim}; projection={args.local_projection_dim or 'identity'}"
+            f"local fusion: {args.local_fusion}/{args.local_modalities}; "
+            f"raw V={head.local_visual_dim}, A={head.local_audio_dim}; "
+            f"projection={args.local_projection_dim or 'identity'}"
         )
+        if person_features is not None:
+            print(
+                "frame sets: "
+                f"{len(person_features.frame_offsets) - 1} source frames, "
+                f"{len(person_features.features)} person observations, "
+                f"max_people={person_features.max_people}, "
+                f"attention={args.person_attention_dim}D/"
+                f"{args.person_attention_heads} heads"
+            )
     trainable_parameters = sum(
         parameter.numel() for parameter in head.parameters() if parameter.requires_grad
     )
@@ -782,7 +862,13 @@ def main() -> None:
             unit="batch",
             disable=args.no_progress,
         )
-        for state_batch, target_batch, transition_target_batch, local_batch in progress:
+        for (
+            state_batch,
+            target_batch,
+            transition_target_batch,
+            local_batch,
+            example_index_batch,
+        ) in progress:
             local_batch = local_batch.to(
                 device=device, dtype=torch.float32, non_blocking=True,
             )
@@ -801,15 +887,37 @@ def main() -> None:
                 dtype=torch.float32,
                 non_blocking=True,
             )
+            person_batch = None
+            person_mask = None
+            if person_features is not None:
+                person_batch, person_mask = person_features.padded(
+                    example_index_batch
+                )
+                person_batch = person_batch.to(
+                    device=device,
+                    dtype=torch.float32,
+                    non_blocking=True,
+                )
+                person_mask = person_mask.to(device=device, non_blocking=True)
             optimizer.zero_grad(set_to_none=True)
             if args.transition_objective == "none":
-                logits = head(state_batch, local_batch)
+                logits = head(
+                    state_batch,
+                    local_batch,
+                    person_batch,
+                    person_mask,
+                )
                 activity_loss = loss_function(logits, target_batch)
                 loss = activity_loss
                 transition_loss = activity_loss.detach().new_zeros(())
                 expected_reward = activity_loss.detach().new_zeros(())
             else:
-                logits, transition_logits = head.forward_with_transition(state_batch, local_batch)
+                logits, transition_logits = head.forward_with_transition(
+                    state_batch,
+                    local_batch,
+                    person_batch,
+                    person_mask,
+                )
                 activity_loss = loss_function(logits, target_batch)
                 components = transition_aware_objective(
                     activity_loss,
@@ -890,6 +998,7 @@ def main() -> None:
                 args.batch_size,
                 args.loss,
                 dev_local,
+                dev_people,
             )
             dev_map = float(dev_metrics["mean_map"])
             record["dev_valid_mean_map"] = dev_map