Add pooled local AV baselines and research notes

infer_turn_head.py

@@ -38,6 +38,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--target-root", type=Path, required=True)
     parser.add_argument("--checkpoint", type=Path, required=True)
     parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument(
+        "--local-feature-root", type=Path,
+        help="Required for a local-fusion checkpoint; modalities come from the checkpoint.",
+    )
     parser.add_argument(
         "--dataset",
         choices=DATASET_NAMES,
@@ -113,6 +117,7 @@ def _predict_states(
     description: str,
     show_progress: bool,
     loss_name: str,
+    local_features: torch.Tensor | None = None,
 ) -> tuple[np.ndarray, np.ndarray | None]:
     chunks: list[torch.Tensor] = []
     transition_chunks: list[torch.Tensor] = []
@@ -130,11 +135,15 @@ def _predict_states(
             dtype=torch.float32,
             non_blocking=True,
         )
+        side = (
+            local_features[start : start + batch_size].to(device=device, dtype=torch.float32)
+            if local_features is not None else None
+        )
         if head.transition_auxiliary:
-            activity_logits, transition_logits = head.forward_with_transition(batch)
+            activity_logits, transition_logits = head.forward_with_transition(batch, side)
             transition_chunks.append(transition_logits.sigmoid().cpu())
         else:
-            activity_logits = head(batch)
+            activity_logits = head(batch, side)
         chunks.append(
             marginal_probabilities_from_logits(activity_logits, loss_name).cpu()
         )
@@ -357,6 +366,13 @@ def main() -> None:
         )
     head_config = checkpoint_head_config(checkpoint_payload)
     loss_name = checkpoint_loss_name(checkpoint_payload)
+    local_contract = checkpoint_payload.get("local_contract")
+    if bool(head.local_input_size) != bool(local_contract):
+        raise ValueError("Checkpoint local dimensions/contract disagree")
+    if bool(head.local_input_size) != (args.local_feature_root is not None):
+        raise ValueError(
+            "--local-feature-root must be supplied exactly when the checkpoint uses local fusion"
+        )
 
     args.output_dir.mkdir(parents=True, exist_ok=True)
     prediction_dir = args.output_dir / "pred"
@@ -384,6 +400,9 @@ def main() -> None:
             name,
             expected_mode=expected_mode,
             dataset=args.dataset,
+            local_feature_root=args.local_feature_root,
+            local_modalities=local_contract["modalities"] if local_contract else "av",
+            local_contract=local_contract,
         )
         if data.hidden_size != int(head_config["hidden_size"]):
             raise ValueError(
@@ -403,6 +422,7 @@ def main() -> None:
             description=f"infer {name}",
             show_progress=not args.no_progress,
             loss_name=loss_name,
+            local_features=data.local_features,
         )
         predictions, valid = _align_to_egospeak_timeline(
             probabilities,
@@ -464,6 +484,7 @@ def main() -> None:
             "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
             "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
             "layer_fusion": head.layer_fusion,
+            "local_contract": local_contract,
             "transition_auxiliary": head.transition_auxiliary,
             "transition_objective": checkpoint_payload.get(
                 "training_config",