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

125
train_turn_head.py
View File

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