wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add causal frame set attention fusion

Browse Source
This commit is contained in:
Xin Wang
2026-09-02 18:48:26 +08:00
parent e769094b3d
commit 8e70281330
8 changed files with 837 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

226
easycom_head.py
View File

@@ -5,13 +5,16 @@ from __future__ import annotations
import csv
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal, Sequence
from typing import TYPE_CHECKING, Any, Literal, Sequence
from uuid import UUID
import numpy as np
import torch
from torch import nn
if TYPE_CHECKING:
from local_turn_features import FramePersonFeatures
OFFICIAL_TRAIN_SESSIONS = (
"Session_7",
@@ -58,6 +61,7 @@ LAYER_FUSION_MODES = (
"scalar-mixture",
"concat-mlp",
)
LOCAL_FUSION_MODES = ("pooled", "frame-set-attention")
DatasetName = Literal["easycom", "ego4d"]
@@ -77,6 +81,9 @@ class TurnPredictionHead(nn.Module):
local_visual_dim: int = 0,
local_audio_dim: int = 0,
local_projection_dim: int = 256,
local_fusion: str = "pooled",
person_attention_dim: int = 256,
person_attention_heads: int = 4,
) -> None:
super().__init__()
if num_input_layers <= 0:
@@ -116,6 +123,24 @@ class TurnPredictionHead(nn.Module):
self.local_visual_dim = local_visual_dim
self.local_audio_dim = local_audio_dim
self.local_projection_dim = local_projection_dim
if local_fusion not in LOCAL_FUSION_MODES:
raise ValueError(
f"Unknown local_fusion={local_fusion!r}; expected one of "
f"{LOCAL_FUSION_MODES}"
)
if person_attention_dim <= 0 or person_attention_heads <= 0:
raise ValueError("Person attention dimension/heads must be positive")
if person_attention_dim % person_attention_heads:
raise ValueError("Person attention dimension must be divisible by heads")
if local_fusion == "frame-set-attention" and (
local_visual_dim <= 0 or local_audio_dim <= 0
):
raise ValueError(
"Frame-set Attention requires both local visual people and audio"
)
self.local_fusion = local_fusion
self.person_attention_dim = person_attention_dim
self.person_attention_heads = person_attention_heads
if self.learnable_layer_mixture:
# Zero initialization exactly reproduces a uniform layer mean at
# step zero while allowing training to select useful query layers.
@@ -133,7 +158,12 @@ class TurnPredictionHead(nn.Module):
self.layer_fusion_net = None
self.local_projections = nn.ModuleDict()
fused_size = hidden_size
for name, dimension in (("visual", local_visual_dim), ("audio", local_audio_dim)):
projected_modalities = (
(("audio", local_audio_dim),)
if local_fusion == "frame-set-attention"
else (("visual", local_visual_dim), ("audio", local_audio_dim))
)
for name, dimension in projected_modalities:
if dimension:
self.local_projections[name] = (
nn.Sequential(
@@ -145,6 +175,36 @@ class TurnPredictionHead(nn.Module):
else nn.Identity()
)
fused_size += (local_projection_dim or dimension) + 1
self.person_input_projection: nn.Module | None = None
self.person_state_query: nn.Module | None = None
self.person_audio_query: nn.Module | None = None
self.person_attention: nn.MultiheadAttention | None = None
self.null_person: nn.Parameter | None = None
if local_fusion == "frame-set-attention":
self.person_input_projection = nn.Sequential(
nn.LayerNorm(local_visual_dim),
nn.Linear(local_visual_dim, person_attention_dim),
nn.GELU(),
)
self.person_state_query = nn.Sequential(
nn.LayerNorm(hidden_size),
nn.Linear(hidden_size, person_attention_dim),
)
self.person_audio_query = nn.Sequential(
nn.LayerNorm(local_audio_dim),
nn.Linear(local_audio_dim, person_attention_dim),
)
self.person_attention = nn.MultiheadAttention(
person_attention_dim,
person_attention_heads,
dropout=dropout,
batch_first=True,
)
self.null_person = nn.Parameter(
torch.zeros(1, 1, person_attention_dim)
)
nn.init.normal_(self.null_person, std=0.02)
fused_size += person_attention_dim + 1
self.net = nn.Sequential(
nn.LayerNorm(fused_size),
nn.Linear(fused_size, hidden_size // 2),
@@ -190,17 +250,99 @@ class TurnPredictionHead(nn.Module):
@property
def local_input_size(self) -> int:
if self.local_fusion == "frame-set-attention":
return self.local_audio_dim + 1
return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)
def _frame_set_context(
self,
states: torch.Tensor,
local: torch.Tensor,
person_features: torch.Tensor | None,
person_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
if states.ndim != 2:
raise ValueError("Frame-set Attention expects batched [B,D] states")
if local.shape != (len(states), self.local_audio_dim + 1):
raise ValueError(
"Frame-set Attention requires aligned [audio,valid] local features"
)
if (
person_features is None
or person_features.ndim != 3
or person_features.shape[0] != len(states)
or person_features.shape[2] != self.local_visual_dim
):
raise ValueError(
"Frame-set Attention requires person features [B,P,visual_dim]"
)
if (
person_mask is None
or person_mask.dtype != torch.bool
or person_mask.shape != person_features.shape[:2]
):
raise ValueError("Frame-set Attention requires a boolean person mask [B,P]")
if not torch.isfinite(person_features).all():
raise ValueError("Person features contain NaN/Inf")
assert self.person_input_projection is not None
assert self.person_state_query is not None
assert self.person_audio_query is not None
assert self.person_attention is not None
assert self.null_person is not None
audio_values = local[:, : self.local_audio_dim]
audio_valid = local[:, self.local_audio_dim : self.local_audio_dim + 1]
audio_projected = self.local_projections["audio"](audio_values) * audio_valid
query = (
self.person_state_query(states)
+ self.person_audio_query(audio_values) * audio_valid
)
people = self.person_input_projection(person_features)
null = self.null_person.expand(len(states), -1, -1)
memory = torch.cat((null, people), dim=1)
null_mask = torch.ones(
(len(states), 1),
dtype=torch.bool,
device=person_mask.device,
)
memory_valid = torch.cat((null_mask, person_mask), dim=1)
context, _ = self.person_attention(
query[:, None, :],
memory,
memory,
key_padding_mask=~memory_valid,
need_weights=False,
)
person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype)
return audio_projected, audio_valid, torch.cat(
(context[:, 0, :], person_valid),
dim=-1,
)
def _shared_features(
self,
states: torch.Tensor,
local: torch.Tensor | None = None,
person_features: torch.Tensor | None = None,
person_mask: torch.Tensor | None = None,
) -> torch.Tensor:
"""Return the activity MLP representation shared by both readouts."""
states = self._fuse_layers(states)
if self.local_input_size:
if self.local_fusion == "frame-set-attention":
if local is None:
raise ValueError("Frame-set Attention requires local audio features")
audio, audio_valid, person_context = self._frame_set_context(
states,
local,
person_features,
person_mask,
)
states = torch.cat(
(states, audio, audio_valid, person_context),
dim=-1,
)
elif self.local_input_size:
if local is None or local.shape != (*states.shape[:-1], self.local_input_size):
raise ValueError(
f"Head requires aligned local features of width {self.local_input_size}"
@@ -220,14 +362,25 @@ class TurnPredictionHead(nn.Module):
states = torch.cat(pieces, dim=-1)
elif local is not None and local.shape[-1] != 0:
raise ValueError("Local features supplied to a baseline head")
if self.local_fusion != "frame-set-attention" and (
person_features is not None or person_mask is not None
):
raise ValueError("Person sets supplied to a non-frame-set head")
return self.net[:-1](states)
def forward(
self,
states: torch.Tensor,
local: torch.Tensor | None = None,
person_features: torch.Tensor | None = None,
person_mask: torch.Tensor | None = None,
) -> torch.Tensor:
features = self._shared_features(states, local)
features = self._shared_features(
states,
local,
person_features,
person_mask,
)
logits = self.net[-1](features)
return logits.view(
*features.shape[:-1], self.num_horizons, self.num_classes
@@ -237,6 +390,8 @@ class TurnPredictionHead(nn.Module):
self,
states: torch.Tensor,
local: torch.Tensor | None = None,
person_features: torch.Tensor | None = None,
person_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Return activity logits and binary transition logits."""
@@ -244,7 +399,12 @@ class TurnPredictionHead(nn.Module):
raise RuntimeError(
"This head was created without transition_auxiliary=True"
)
features = self._shared_features(states, local)
features = self._shared_features(
states,
local,
person_features,
person_mask,
)
activity_logits = self.net[-1](features).view(
*features.shape[:-1],
self.num_horizons,
@@ -734,6 +894,7 @@ class SessionData:
label_indices: torch.LongTensor
feature_metadata: dict[str, Any]
local_features: torch.Tensor | None = None
local_person_features: "FramePersonFeatures | None" = None
@property
def hidden_size(self) -> int:
@@ -768,6 +929,7 @@ def load_session_data(
local_feature_root: str | Path | None = None,
local_modalities: str = "av",
local_contract: dict[str, Any] | None = None,
local_fusion: str = "pooled",
) -> SessionData:
"""Load and validate one saved Omni feature payload and target array."""
@@ -864,22 +1026,55 @@ def load_session_data(
"".join(str(int(value)) for value in row.tolist()): int(count)
for row, count in zip(unique_rows, row_counts)
}
if local_fusion not in LOCAL_FUSION_MODES:
raise ValueError(f"Unknown local fusion: {local_fusion}")
local = None
local_people = None
if local_feature_root is not None:
from local_turn_features import load_local_features
local, contract = load_local_features(
resolve_feature_path(local_feature_root, name, dataset=dataset),
payload,
sample_id=name,
local_path = resolve_feature_path(
local_feature_root,
name,
dataset=dataset,
modalities=local_modalities,
expected_contract=local_contract,
)
if local_fusion == "frame-set-attention":
if local_modalities != "av":
raise ValueError(
"Frame-set Attention requires --local-modalities av"
)
from local_turn_features import load_frame_set_features
local, local_people, contract = load_frame_set_features(
local_path,
payload,
sample_id=name,
dataset=dataset,
expected_contract=local_contract,
)
else:
from local_turn_features import load_local_features
local, contract = load_local_features(
local_path,
payload,
sample_id=name,
dataset=dataset,
modalities=local_modalities,
expected_contract=local_contract,
)
metadata["local_contract"] = contract
elif local_contract is not None:
raise ValueError("--local-feature-root is required for this checkpoint")
return SessionData(name, states, targets, label_indices, metadata, local)
elif local_fusion != "pooled":
raise ValueError("Frame-set Attention requires --local-feature-root")
return SessionData(
name,
states,
targets,
label_indices,
metadata,
local,
local_people,
)
def average_precision_metrics(
@@ -924,6 +1119,9 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
config.setdefault("local_visual_dim", 0)
config.setdefault("local_audio_dim", 0)
config.setdefault("local_projection_dim", 256)
config.setdefault("local_fusion", "pooled")
config.setdefault("person_attention_dim", 256)
config.setdefault("person_attention_heads", 4)
config.setdefault(
"layer_fusion",
(