Add pooled local AV baselines and research notes

This commit is contained in:
Xin Wang
2026-08-27 14:19:33 +08:00
parent d34725778f
commit 11f197724f
23 changed files with 1541 additions and 28 deletions

View File

@@ -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",