1167 lines
41 KiB
Python
1167 lines
41 KiB
Python
"""Shared head, labels, and metrics for the EasyCom/Ego4D experiments."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import csv
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
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",
|
|
"Session_11",
|
|
"Session_10",
|
|
"Session_6",
|
|
"Session_9",
|
|
"Session_4",
|
|
"Session_8",
|
|
"Session_5",
|
|
"Session_12",
|
|
)
|
|
OFFICIAL_TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
|
|
CLASS_NAMES = (
|
|
"background",
|
|
"camera_wearer_speaking",
|
|
"other_people_speaking",
|
|
)
|
|
FOUR_STATE_NAMES = (
|
|
"silence",
|
|
"wearer_only",
|
|
"others_only",
|
|
"overlap",
|
|
)
|
|
STEP_SECONDS = 0.2
|
|
NUM_HORIZONS = 10
|
|
NUM_CLASSES = 3
|
|
NUM_FOUR_STATES = len(FOUR_STATE_NAMES)
|
|
LOSS_NAMES = (
|
|
"lstr-mce",
|
|
"miniroad-normalized",
|
|
"independent-bce",
|
|
"four-state-ce",
|
|
)
|
|
TRANSITION_OBJECTIVE_NAMES = (
|
|
"none",
|
|
"r1-supervised",
|
|
"r2-expected-reward",
|
|
)
|
|
DATASET_NAMES = ("easycom", "ego4d")
|
|
LAYER_FUSION_MODES = (
|
|
"single",
|
|
"fixed-mean",
|
|
"scalar-mixture",
|
|
"concat-mlp",
|
|
)
|
|
LOCAL_FUSION_MODES = ("pooled", "frame-set-attention")
|
|
DatasetName = Literal["easycom", "ego4d"]
|
|
|
|
|
|
class TurnPredictionHead(nn.Module):
|
|
"""Map one Omni state to configurable ``[horizon, output]`` logits."""
|
|
|
|
def __init__(
|
|
self,
|
|
hidden_size: int = 2048,
|
|
num_horizons: int = NUM_HORIZONS,
|
|
num_classes: int = NUM_CLASSES,
|
|
dropout: float = 0.1,
|
|
num_input_layers: int = 1,
|
|
learnable_layer_mixture: bool = False,
|
|
layer_fusion: str | None = None,
|
|
transition_auxiliary: bool = False,
|
|
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:
|
|
raise ValueError("num_input_layers must be positive")
|
|
if layer_fusion is None:
|
|
layer_fusion = (
|
|
"scalar-mixture" if learnable_layer_mixture else "single"
|
|
)
|
|
if layer_fusion not in LAYER_FUSION_MODES:
|
|
raise ValueError(
|
|
f"Unknown layer_fusion={layer_fusion!r}; expected one of "
|
|
f"{LAYER_FUSION_MODES}"
|
|
)
|
|
if learnable_layer_mixture and layer_fusion != "scalar-mixture":
|
|
raise ValueError(
|
|
"learnable_layer_mixture=True is only compatible with "
|
|
"layer_fusion='scalar-mixture'"
|
|
)
|
|
if layer_fusion != "single" and num_input_layers < 2:
|
|
raise ValueError(
|
|
f"layer_fusion={layer_fusion!r} requires at least two input layers"
|
|
)
|
|
if layer_fusion == "single" and num_input_layers != 1:
|
|
raise ValueError(
|
|
"Multiple input layers require a non-single layer_fusion"
|
|
)
|
|
self.hidden_size = hidden_size
|
|
self.num_horizons = num_horizons
|
|
self.num_classes = num_classes
|
|
self.dropout = dropout
|
|
self.num_input_layers = num_input_layers
|
|
self.layer_fusion = layer_fusion
|
|
self.learnable_layer_mixture = layer_fusion == "scalar-mixture"
|
|
self.transition_auxiliary = transition_auxiliary
|
|
if min(local_visual_dim, local_audio_dim, local_projection_dim) < 0:
|
|
raise ValueError("Local dimensions must be non-negative")
|
|
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.
|
|
self.layer_logits = nn.Parameter(torch.zeros(num_input_layers))
|
|
else:
|
|
self.register_parameter("layer_logits", None)
|
|
if layer_fusion == "concat-mlp":
|
|
fusion_input_size = num_input_layers * hidden_size
|
|
self.layer_fusion_net: nn.Module | None = nn.Sequential(
|
|
nn.LayerNorm(fusion_input_size),
|
|
nn.Linear(fusion_input_size, hidden_size),
|
|
nn.GELU(),
|
|
)
|
|
else:
|
|
self.layer_fusion_net = None
|
|
self.local_projections = nn.ModuleDict()
|
|
fused_size = hidden_size
|
|
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(
|
|
nn.LayerNorm(dimension),
|
|
nn.Linear(dimension, local_projection_dim),
|
|
nn.GELU(),
|
|
)
|
|
if local_projection_dim
|
|
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),
|
|
nn.GELU(),
|
|
nn.Dropout(dropout),
|
|
nn.Linear(hidden_size // 2, num_horizons * num_classes),
|
|
)
|
|
self.transition_head: nn.Module | None
|
|
if transition_auxiliary:
|
|
# Share the activity MLP trunk and add only a small binary readout
|
|
# for each future horizon.
|
|
self.transition_head = nn.Linear(hidden_size // 2, num_horizons)
|
|
else:
|
|
self.transition_head = None
|
|
|
|
def _fuse_layers(self, states: torch.Tensor) -> torch.Tensor:
|
|
if states.shape[-1] != self.hidden_size:
|
|
raise ValueError(
|
|
f"Expected hidden size {self.hidden_size}, got {states.shape[-1]}"
|
|
)
|
|
if self.layer_fusion != "single":
|
|
if states.ndim < 3 or states.shape[-2] != self.num_input_layers:
|
|
raise ValueError(
|
|
"Multi-layer input must end in "
|
|
f"[{self.num_input_layers},{self.hidden_size}], got "
|
|
f"{tuple(states.shape)}"
|
|
)
|
|
if self.layer_fusion == "fixed-mean":
|
|
states = states.mean(dim=-2)
|
|
elif self.layer_fusion == "scalar-mixture":
|
|
assert self.layer_logits is not None
|
|
weights = self.layer_logits.softmax(dim=0).to(states.dtype)
|
|
view_shape = (1,) * (states.ndim - 2) + (
|
|
self.num_input_layers,
|
|
1,
|
|
)
|
|
states = (states * weights.view(view_shape)).sum(dim=-2)
|
|
elif self.layer_fusion == "concat-mlp":
|
|
assert self.layer_fusion_net is not None
|
|
states = states.flatten(start_dim=-2)
|
|
states = self.layer_fusion_net(states)
|
|
return states
|
|
|
|
@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_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}"
|
|
)
|
|
pieces = [states]
|
|
offset = 0
|
|
for name, dimension in (
|
|
("visual", self.local_visual_dim),
|
|
("audio", self.local_audio_dim),
|
|
):
|
|
if dimension:
|
|
values = local[..., offset : offset + dimension]
|
|
valid = local[..., offset + dimension : offset + dimension + 1]
|
|
# Mask AFTER projection too: a learned bias is not an observation.
|
|
pieces.extend((self.local_projections[name](values) * valid, valid))
|
|
offset += dimension + 1
|
|
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,
|
|
person_features,
|
|
person_mask,
|
|
)
|
|
logits = self.net[-1](features)
|
|
return logits.view(
|
|
*features.shape[:-1], self.num_horizons, self.num_classes
|
|
)
|
|
|
|
def forward_with_transition(
|
|
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."""
|
|
|
|
if self.transition_head is None:
|
|
raise RuntimeError(
|
|
"This head was created without transition_auxiliary=True"
|
|
)
|
|
features = self._shared_features(
|
|
states,
|
|
local,
|
|
person_features,
|
|
person_mask,
|
|
)
|
|
activity_logits = self.net[-1](features).view(
|
|
*features.shape[:-1],
|
|
self.num_horizons,
|
|
self.num_classes,
|
|
)
|
|
transition_logits = self.transition_head(features)
|
|
return activity_logits, transition_logits
|
|
|
|
def layer_mixture_weights(self) -> torch.Tensor | None:
|
|
"""Return normalized oldest-to-newest layer weights, if enabled."""
|
|
|
|
if self.layer_logits is None:
|
|
return None
|
|
return self.layer_logits.softmax(dim=0)
|
|
|
|
|
|
def _validate_loss_inputs(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
num_output_classes: int = NUM_CLASSES,
|
|
) -> None:
|
|
if logits.shape != targets.shape:
|
|
raise ValueError(
|
|
f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
|
|
f"{tuple(targets.shape)}"
|
|
)
|
|
if logits.shape[-1] != num_output_classes:
|
|
raise ValueError(
|
|
f"Expected {num_output_classes} classes, got {logits.shape[-1]}"
|
|
)
|
|
|
|
|
|
def _validated_binary_activity_targets(
|
|
targets: torch.Tensor,
|
|
*,
|
|
allow_noncanonical_binary: bool = False,
|
|
) -> torch.Tensor:
|
|
"""Validate three-label activity rows.
|
|
|
|
EasyCom uses four explicit activity states. Ego4D's released target files
|
|
are consumed by EgoSpeak's MCE objective as binary/multi-positive weights,
|
|
so the compatibility path preserves any binary row instead of silently
|
|
rewriting it to an EasyCom state.
|
|
"""
|
|
|
|
if targets.ndim < 1 or targets.shape[-1] != NUM_CLASSES:
|
|
raise ValueError(
|
|
f"Expected activity targets [...,{NUM_CLASSES}], got "
|
|
f"{tuple(targets.shape)}"
|
|
)
|
|
if not torch.isfinite(targets).all():
|
|
raise ValueError("Activity targets contain NaN or Inf")
|
|
rounded = targets.round()
|
|
if not torch.allclose(targets, rounded, atol=1e-6, rtol=0.0):
|
|
raise ValueError("Activity targets must contain only binary 0/1 values")
|
|
expected_background = rounded[..., 1:].sum(dim=-1) == 0
|
|
valid = rounded[..., 0].bool() == expected_background
|
|
if allow_noncanonical_binary:
|
|
return rounded
|
|
if not bool(valid.all()):
|
|
invalid = torch.unique(rounded[~valid].reshape(-1, NUM_CLASSES), dim=0)
|
|
invalid_rows = ", ".join(
|
|
str([int(value) for value in row.tolist()])
|
|
for row in invalid[:8]
|
|
)
|
|
if len(invalid) > 8:
|
|
invalid_rows += ", ..."
|
|
raise ValueError(
|
|
"Activity targets must be one of [1,0,0], [0,1,0], "
|
|
"[0,0,1], or [0,1,1]; observed invalid "
|
|
f"row(s): {invalid_rows}"
|
|
)
|
|
return rounded
|
|
|
|
|
|
def activity_targets_to_four_states(targets: torch.Tensor) -> torch.LongTensor:
|
|
"""Map three-label activity targets to silence/wearer/others/overlap."""
|
|
|
|
targets = _validated_binary_activity_targets(targets)
|
|
return (targets[..., 1] + 2 * targets[..., 2]).long()
|
|
|
|
|
|
def transition_targets_from_activity(
|
|
activity_sequence: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Mark every activity-state change between adjacent 5 Hz target rows.
|
|
|
|
The input ends in [time, 3] and the output ends in [time - 1].
|
|
A silence/onset, offset, speaker handover, or overlap-state change is one
|
|
positive transition.
|
|
"""
|
|
|
|
activity_sequence = _validated_binary_activity_targets(activity_sequence)
|
|
if activity_sequence.ndim < 2 or activity_sequence.shape[-2] < 2:
|
|
raise ValueError(
|
|
"Activity sequence must contain at least two target rows, got "
|
|
f"{tuple(activity_sequence.shape)}"
|
|
)
|
|
changed = activity_sequence[..., 1:, :] != activity_sequence[..., :-1, :]
|
|
return changed.any(dim=-1).to(activity_sequence.dtype)
|
|
|
|
|
|
def weighted_transition_bce_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
pos_weight: float = 3.0,
|
|
) -> torch.Tensor:
|
|
"""Binary transition loss with an explicit positive-class weight."""
|
|
|
|
if logits.shape != targets.shape:
|
|
raise ValueError(
|
|
f"Transition logits/targets mismatch: {tuple(logits.shape)} vs "
|
|
f"{tuple(targets.shape)}"
|
|
)
|
|
if pos_weight <= 0:
|
|
raise ValueError("Transition pos_weight must be positive")
|
|
targets = targets.to(device=logits.device, dtype=logits.dtype)
|
|
if not torch.isfinite(targets).all() or not torch.all(
|
|
(targets == 0) | (targets == 1)
|
|
):
|
|
raise ValueError("Transition targets must contain only finite 0/1 values")
|
|
return nn.functional.binary_cross_entropy_with_logits(
|
|
logits,
|
|
targets,
|
|
pos_weight=logits.new_tensor(pos_weight),
|
|
)
|
|
|
|
|
|
def expected_transition_reward(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
pos_weight: float = 3.0,
|
|
) -> torch.Tensor:
|
|
"""Exact expected weighted correctness of a Bernoulli transition policy.
|
|
|
|
Correct transition/non-transition actions receive reward one and incorrect
|
|
actions zero. The expectation is evaluated analytically rather than by
|
|
sampling, and normalized by the sum of example weights to remain in [0, 1].
|
|
"""
|
|
|
|
if logits.shape != targets.shape:
|
|
raise ValueError(
|
|
f"Transition logits/targets mismatch: {tuple(logits.shape)} vs "
|
|
f"{tuple(targets.shape)}"
|
|
)
|
|
if pos_weight <= 0:
|
|
raise ValueError("Transition pos_weight must be positive")
|
|
targets = targets.to(device=logits.device, dtype=logits.dtype)
|
|
if not torch.isfinite(targets).all() or not torch.all(
|
|
(targets == 0) | (targets == 1)
|
|
):
|
|
raise ValueError("Transition targets must contain only finite 0/1 values")
|
|
probability = logits.sigmoid()
|
|
probability_correct = (
|
|
targets * probability + (1.0 - targets) * (1.0 - probability)
|
|
)
|
|
weights = 1.0 + (pos_weight - 1.0) * targets
|
|
return (weights * probability_correct).sum() / weights.sum()
|
|
|
|
|
|
def transition_aware_objective(
|
|
activity_loss: torch.Tensor,
|
|
transition_logits: torch.Tensor,
|
|
transition_targets: torch.Tensor,
|
|
*,
|
|
objective: str,
|
|
transition_weight: float = 0.2,
|
|
pos_weight: float = 3.0,
|
|
reward_weight: float = 0.1,
|
|
) -> dict[str, torch.Tensor]:
|
|
"""Compose the R1 or R2 transition-aware training objective."""
|
|
|
|
if objective not in TRANSITION_OBJECTIVE_NAMES[1:]:
|
|
raise ValueError(
|
|
f"Transition objective must be R1 or R2, got {objective!r}"
|
|
)
|
|
if transition_weight < 0 or reward_weight < 0:
|
|
raise ValueError("Transition and reward weights must be non-negative")
|
|
transition_loss = weighted_transition_bce_loss(
|
|
transition_logits,
|
|
transition_targets,
|
|
pos_weight=pos_weight,
|
|
)
|
|
expected_reward = expected_transition_reward(
|
|
transition_logits,
|
|
transition_targets,
|
|
pos_weight=pos_weight,
|
|
)
|
|
applied_reward = (
|
|
expected_reward
|
|
if objective == "r2-expected-reward"
|
|
else expected_reward.detach().new_zeros(())
|
|
)
|
|
total = (
|
|
activity_loss
|
|
+ transition_weight * transition_loss
|
|
- reward_weight * applied_reward
|
|
)
|
|
return {
|
|
"total": total,
|
|
"activity": activity_loss,
|
|
"transition": transition_loss,
|
|
"expected_reward": expected_reward,
|
|
"applied_reward": applied_reward,
|
|
}
|
|
|
|
|
|
def output_classes_for_loss(loss_name: str) -> int:
|
|
"""Return the number of head logits required by one training objective."""
|
|
|
|
if loss_name not in LOSS_NAMES:
|
|
raise ValueError(f"Unknown loss {loss_name!r}; expected one of {LOSS_NAMES}")
|
|
return NUM_FOUR_STATES if loss_name == "four-state-ce" else NUM_CLASSES
|
|
|
|
|
|
def output_space_for_loss(loss_name: str) -> str:
|
|
"""Return a stable checkpoint description of the head output space."""
|
|
|
|
output_classes_for_loss(loss_name)
|
|
if loss_name == "four-state-ce":
|
|
return "four-state-softmax"
|
|
if loss_name == "independent-bce":
|
|
return "three-label-sigmoid"
|
|
return "three-class-softmax"
|
|
|
|
|
|
def head_output_names_for_loss(loss_name: str) -> tuple[str, ...]:
|
|
output_classes_for_loss(loss_name)
|
|
return FOUR_STATE_NAMES if loss_name == "four-state-ce" else CLASS_NAMES
|
|
|
|
|
|
def egospeak_multihorizon_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
reduction: Literal["mean", "sum", "none"] = "mean",
|
|
) -> torch.Tensor:
|
|
"""Reproduce the LSTR ``MCE`` loss used by the supplied shell/config.
|
|
|
|
EasyCom targets may contain both active speaker classes as ``[0, 1, 1]``.
|
|
LSTR preserves those weights directly and sums their negative log-softmax
|
|
terms. It does not normalize the target row.
|
|
"""
|
|
|
|
_validate_loss_inputs(logits, targets)
|
|
targets = targets.to(device=logits.device, dtype=logits.dtype)
|
|
loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
|
|
if reduction == "mean":
|
|
return loss.mean()
|
|
if reduction == "sum":
|
|
return loss.sum()
|
|
if reduction == "none":
|
|
return loss
|
|
raise ValueError(f"Unknown reduction: {reduction}")
|
|
|
|
|
|
def miniroad_multihorizon_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
reduction: Literal["mean", "sum", "none"] = "mean",
|
|
) -> torch.Tensor:
|
|
"""Reproduce MiniROAD's L2-normalized multi-positive anticipation loss."""
|
|
|
|
_validate_loss_inputs(logits, targets)
|
|
targets = nn.functional.normalize(
|
|
targets.to(device=logits.device, dtype=logits.dtype),
|
|
p=2,
|
|
dim=-1,
|
|
)
|
|
loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
|
|
if reduction == "mean":
|
|
return loss.mean()
|
|
if reduction == "sum":
|
|
return loss.sum()
|
|
if reduction == "none":
|
|
return loss
|
|
raise ValueError(f"Unknown reduction: {reduction}")
|
|
|
|
|
|
def independent_bce_multihorizon_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
reduction: Literal["mean", "sum", "none"] = "mean",
|
|
) -> torch.Tensor:
|
|
"""Independent sigmoid/BCE loss for the three marginal activity labels."""
|
|
|
|
_validate_loss_inputs(logits, targets)
|
|
targets = targets.to(
|
|
device=logits.device,
|
|
dtype=logits.dtype,
|
|
)
|
|
element_loss = nn.functional.binary_cross_entropy_with_logits(
|
|
logits,
|
|
targets,
|
|
reduction="none",
|
|
)
|
|
if reduction == "mean":
|
|
return element_loss.mean()
|
|
if reduction == "sum":
|
|
return element_loss.sum()
|
|
if reduction == "none":
|
|
return element_loss
|
|
raise ValueError(f"Unknown reduction: {reduction}")
|
|
|
|
|
|
def four_state_multihorizon_loss(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
*,
|
|
reduction: Literal["mean", "sum", "none"] = "mean",
|
|
) -> torch.Tensor:
|
|
"""Four-state CE for silence, wearer-only, others-only, and overlap."""
|
|
|
|
if logits.ndim < 1 or logits.shape[-1] != NUM_FOUR_STATES:
|
|
raise ValueError(
|
|
f"Expected four-state logits [...,{NUM_FOUR_STATES}], got "
|
|
f"{tuple(logits.shape)}"
|
|
)
|
|
if logits.shape[:-1] != targets.shape[:-1] or targets.shape[-1] != NUM_CLASSES:
|
|
raise ValueError(
|
|
"Four-state logits/activity target shape mismatch: "
|
|
f"{tuple(logits.shape)} vs {tuple(targets.shape)}"
|
|
)
|
|
state_targets = (targets[..., 1] + 2 * targets[..., 2]).long().to(
|
|
logits.device
|
|
)
|
|
loss = nn.functional.cross_entropy(
|
|
logits.reshape(-1, NUM_FOUR_STATES),
|
|
state_targets.reshape(-1),
|
|
reduction="none",
|
|
).reshape(logits.shape[:-1])
|
|
if reduction == "mean":
|
|
return loss.mean()
|
|
if reduction == "sum":
|
|
return loss.sum()
|
|
if reduction == "none":
|
|
return loss
|
|
raise ValueError(f"Unknown reduction: {reduction}")
|
|
|
|
|
|
def marginal_probabilities_from_logits(
|
|
logits: torch.Tensor,
|
|
loss_name: str,
|
|
) -> torch.Tensor:
|
|
"""Convert any supported head output to three EgoSpeak marginals.
|
|
|
|
Four-state marginals deliberately are not renormalized: wearer and others
|
|
may both be near one when the overlap state is likely.
|
|
"""
|
|
|
|
expected_classes = output_classes_for_loss(loss_name)
|
|
if logits.ndim < 1 or logits.shape[-1] != expected_classes:
|
|
raise ValueError(
|
|
f"Loss {loss_name!r} expects {expected_classes} logits, got "
|
|
f"{tuple(logits.shape)}"
|
|
)
|
|
if loss_name == "independent-bce":
|
|
return logits.sigmoid()
|
|
probabilities = logits.softmax(dim=-1)
|
|
if loss_name != "four-state-ce":
|
|
return probabilities
|
|
return torch.stack(
|
|
(
|
|
probabilities[..., 0],
|
|
probabilities[..., 1] + probabilities[..., 3],
|
|
probabilities[..., 2] + probabilities[..., 3],
|
|
),
|
|
dim=-1,
|
|
)
|
|
|
|
|
|
def multihorizon_cross_entropy(
|
|
logits: torch.Tensor,
|
|
targets: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Backward-compatible alias for the LSTR/EgoSpeak loss."""
|
|
|
|
return egospeak_multihorizon_loss(logits, targets)
|
|
|
|
|
|
def normalize_session_name(value: str | int) -> str:
|
|
text = str(value)
|
|
if text.startswith("Session_"):
|
|
suffix = text.removeprefix("Session_")
|
|
else:
|
|
suffix = text
|
|
if not suffix.isdigit():
|
|
raise ValueError(f"Invalid EasyCom session: {value}")
|
|
return f"Session_{int(suffix)}"
|
|
|
|
|
|
def normalize_sample_name(
|
|
value: str | int,
|
|
dataset: DatasetName = "easycom",
|
|
) -> str:
|
|
"""Return the canonical EasyCom session name or Ego4D clip UUID."""
|
|
|
|
if dataset == "easycom":
|
|
return normalize_session_name(value)
|
|
if dataset != "ego4d":
|
|
raise ValueError(f"Unknown dataset: {dataset}")
|
|
|
|
text = str(value).strip()
|
|
try:
|
|
return str(UUID(text))
|
|
except (ValueError, AttributeError) as error:
|
|
raise ValueError(f"Invalid Ego4D clip UUID: {value}") from error
|
|
|
|
|
|
def resolve_feature_path(
|
|
feature_root: str | Path,
|
|
sample: str | int,
|
|
*,
|
|
dataset: DatasetName = "easycom",
|
|
) -> Path:
|
|
"""Read flat features or Ego4D train/test subdirectories without link farms.
|
|
|
|
Manifest membership, not directory discovery, determines the data split.
|
|
Multiple copies are rejected rather than silently choosing stale features.
|
|
"""
|
|
name = normalize_sample_name(sample, dataset)
|
|
root = Path(feature_root)
|
|
candidates = [root / f"{name}.pt"]
|
|
if dataset == "ego4d":
|
|
candidates += [root / split / f"{name}.pt" for split in ("train", "test")]
|
|
found = [path for path in candidates if path.is_file()]
|
|
if len(found) > 1:
|
|
raise ValueError(f"Ambiguous feature copies for {name}: {found}; select one feature root")
|
|
if not found:
|
|
raise FileNotFoundError(f"No features for {name}; checked: " + ", ".join(map(str, candidates)))
|
|
return found[0]
|
|
|
|
|
|
def load_ego4d_manifest_split(
|
|
manifest: str | Path,
|
|
split: Literal["train", "test", "all"],
|
|
) -> list[str]:
|
|
"""Load ordered UIDs from ``ego4d_egospeak.tsv``; all preserves row order."""
|
|
|
|
if split not in {"train", "test", "all"}:
|
|
raise ValueError(f"Invalid Ego4D split: {split}")
|
|
path = Path(manifest)
|
|
if not path.is_file():
|
|
raise FileNotFoundError(path)
|
|
|
|
selected: list[str] = []
|
|
seen: set[str] = set()
|
|
with path.open(newline="") as file:
|
|
reader = csv.DictReader(file, delimiter="\t")
|
|
columns = set(reader.fieldnames or [])
|
|
required = {"split", "uid"}
|
|
if not required.issubset(columns):
|
|
raise ValueError(
|
|
f"Manifest must contain {sorted(required)}, got "
|
|
f"{reader.fieldnames}: {path}"
|
|
)
|
|
for row_number, row in enumerate(reader, start=2):
|
|
row_split = row["split"].strip()
|
|
if row_split not in {"train", "test"}:
|
|
raise ValueError(
|
|
f"Invalid split {row_split!r} at {path}:{row_number}"
|
|
)
|
|
uid = normalize_sample_name(row["uid"], "ego4d")
|
|
if uid in seen:
|
|
raise ValueError(
|
|
f"Duplicate Ego4D UID {uid} at {path}:{row_number}"
|
|
)
|
|
seen.add(uid)
|
|
if split == "all" or row_split == split:
|
|
selected.append(uid)
|
|
|
|
if not selected:
|
|
raise ValueError(f"No {split!r} samples found in {path}")
|
|
return selected
|
|
|
|
|
|
@dataclass
|
|
class SessionData:
|
|
name: str
|
|
states: torch.Tensor
|
|
targets: torch.Tensor
|
|
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:
|
|
return int(self.states.shape[-1])
|
|
|
|
@property
|
|
def num_input_layers(self) -> int:
|
|
return 1 if self.states.ndim == 2 else int(self.states.shape[-2])
|
|
|
|
def future_examples(
|
|
self,
|
|
num_horizons: int = NUM_HORIZONS,
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
|
|
"""Return states, future labels, and target indices for valid rows."""
|
|
|
|
offsets = torch.arange(1, num_horizons + 1, dtype=torch.long)
|
|
future_indices = self.label_indices[:, None] + offsets[None, :]
|
|
valid = (future_indices >= 0).all(dim=1) & (
|
|
future_indices < len(self.targets)
|
|
).all(dim=1)
|
|
rows = future_indices[valid]
|
|
return self.states[valid], self.targets[rows], rows
|
|
|
|
|
|
def load_session_data(
|
|
feature_root: str | Path,
|
|
target_root: str | Path,
|
|
session: str | int,
|
|
*,
|
|
expected_mode: str | None = "causal-prefix",
|
|
dataset: DatasetName = "easycom",
|
|
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."""
|
|
|
|
name = normalize_sample_name(session, dataset)
|
|
feature_path = resolve_feature_path(feature_root, name, dataset=dataset)
|
|
target_path = Path(target_root) / f"{name}.npy"
|
|
if not feature_path.is_file():
|
|
raise FileNotFoundError(feature_path)
|
|
if not target_path.is_file():
|
|
raise FileNotFoundError(target_path)
|
|
|
|
payload = torch.load(feature_path, map_location="cpu", weights_only=False)
|
|
if not isinstance(payload, dict) or "states" not in payload:
|
|
raise ValueError(f"Invalid feature payload: {feature_path}")
|
|
states = payload["states"]
|
|
if not isinstance(states, torch.Tensor) or states.ndim not in {2, 3}:
|
|
raise ValueError(f"states must be [T,D] or [T,L,D] in {feature_path}")
|
|
if states.ndim == 3 and states.shape[1] < 2:
|
|
raise ValueError(
|
|
f"Stacked states must contain at least two layers: {feature_path}"
|
|
)
|
|
if not torch.is_floating_point(states):
|
|
raise ValueError(f"states must be floating point in {feature_path}")
|
|
if not torch.isfinite(states).all():
|
|
raise ValueError(f"states contain NaN/Inf in {feature_path}")
|
|
|
|
mode = payload.get("mode")
|
|
if expected_mode is not None and mode != expected_mode:
|
|
raise ValueError(
|
|
f"{feature_path} mode is {mode!r}, expected {expected_mode!r}"
|
|
)
|
|
label_indices = payload.get("label_indices")
|
|
if label_indices is None:
|
|
raise ValueError(f"Missing label_indices in {feature_path}")
|
|
label_indices = torch.as_tensor(label_indices, dtype=torch.long)
|
|
if label_indices.ndim != 1 or len(label_indices) != len(states):
|
|
raise ValueError(f"Invalid label_indices in {feature_path}")
|
|
if len(label_indices) and not torch.equal(
|
|
label_indices,
|
|
torch.arange(len(label_indices), dtype=torch.long),
|
|
):
|
|
raise ValueError(
|
|
f"{feature_path} does not use the verified zero-based dense alignment"
|
|
)
|
|
|
|
target_array = np.load(target_path, allow_pickle=False)
|
|
if target_array.ndim != 2 or target_array.shape[1] != NUM_CLASSES:
|
|
raise ValueError(
|
|
f"Expected target [T,{NUM_CLASSES}], got {target_array.shape}: "
|
|
f"{target_path}"
|
|
)
|
|
targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
|
|
try:
|
|
_validated_binary_activity_targets(
|
|
targets,
|
|
allow_noncanonical_binary=dataset == "ego4d",
|
|
)
|
|
except ValueError as error:
|
|
raise ValueError(f"{error}: {target_path}") from error
|
|
if len(label_indices) and int(label_indices[-1]) >= len(targets):
|
|
raise ValueError(
|
|
f"Last label index {int(label_indices[-1])} exceeds target length "
|
|
f"{len(targets)} in {name}"
|
|
)
|
|
if len(states) != len(targets):
|
|
raise ValueError(
|
|
f"{name}: feature/target length mismatch: {len(states)} vs "
|
|
f"{len(targets)}"
|
|
)
|
|
|
|
metadata = {
|
|
key: payload.get(key)
|
|
for key in (
|
|
"model_id",
|
|
"mode",
|
|
"config",
|
|
"start",
|
|
"end",
|
|
"num_steps",
|
|
"video_sampling_grid",
|
|
)
|
|
}
|
|
expected_background = targets[:, 1:].sum(dim=-1) == 0
|
|
canonical_rows = targets[:, 0].bool() == expected_background
|
|
unique_rows, row_counts = torch.unique(
|
|
targets.to(dtype=torch.int64),
|
|
dim=0,
|
|
return_counts=True,
|
|
)
|
|
metadata["num_noncanonical_target_rows"] = int(
|
|
(~canonical_rows).sum().item()
|
|
)
|
|
metadata["target_row_pattern_counts"] = {
|
|
"".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:
|
|
local_path = resolve_feature_path(
|
|
local_feature_root,
|
|
name,
|
|
dataset=dataset,
|
|
)
|
|
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")
|
|
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(
|
|
scores: np.ndarray,
|
|
targets: np.ndarray,
|
|
) -> dict[str, Any]:
|
|
"""Compute EgoSpeak frame AP for all three classes, including background."""
|
|
|
|
from sklearn.metrics import average_precision_score
|
|
|
|
scores = np.asarray(scores, dtype=np.float64)
|
|
targets = np.asarray(targets, dtype=np.float64)
|
|
if scores.shape != targets.shape or scores.ndim != 2:
|
|
raise ValueError(
|
|
f"scores/targets must have equal [N,C] shape, got "
|
|
f"{scores.shape} and {targets.shape}"
|
|
)
|
|
if scores.shape[1] != NUM_CLASSES:
|
|
raise ValueError(f"Expected {NUM_CLASSES} classes, got {scores.shape[1]}")
|
|
|
|
aps = average_precision_score(targets, scores, average=None)
|
|
return {
|
|
"num_steps": int(len(scores)),
|
|
"class_ap": {
|
|
name: float(value) for name, value in zip(CLASS_NAMES, aps)
|
|
},
|
|
"map": float(np.mean(aps)),
|
|
}
|
|
|
|
|
|
def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
|
|
"""Return a validated head config, including legacy-checkpoint defaults."""
|
|
|
|
config = dict(payload.get("head_config", {}))
|
|
config.setdefault("hidden_size", 2048)
|
|
config.setdefault("num_horizons", NUM_HORIZONS)
|
|
config.setdefault("num_classes", NUM_CLASSES)
|
|
config.setdefault("dropout", 0.1)
|
|
config.setdefault("num_input_layers", 1)
|
|
config.setdefault("learnable_layer_mixture", False)
|
|
config.setdefault("transition_auxiliary", False)
|
|
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",
|
|
(
|
|
"scalar-mixture"
|
|
if config["learnable_layer_mixture"]
|
|
else "single"
|
|
),
|
|
)
|
|
return config
|
|
|
|
|
|
def checkpoint_loss_name(payload: dict[str, Any]) -> str:
|
|
"""Read the objective from a checkpoint, defaulting legacy files to MCE."""
|
|
|
|
training_config = payload.get("training_config", {})
|
|
if not isinstance(training_config, dict):
|
|
raise ValueError("Checkpoint training_config must be a mapping")
|
|
loss_name = str(training_config.get("loss", "lstr-mce"))
|
|
expected_classes = output_classes_for_loss(loss_name)
|
|
actual_classes = int(checkpoint_head_config(payload)["num_classes"])
|
|
if actual_classes != expected_classes:
|
|
raise ValueError(
|
|
f"Checkpoint loss {loss_name!r} requires {expected_classes} head "
|
|
f"outputs, but head_config has {actual_classes}"
|
|
)
|
|
recorded_space = payload.get("output_space")
|
|
expected_space = output_space_for_loss(loss_name)
|
|
if recorded_space is not None and recorded_space != expected_space:
|
|
raise ValueError(
|
|
f"Checkpoint output_space is {recorded_space!r}, expected "
|
|
f"{expected_space!r} for loss {loss_name!r}"
|
|
)
|
|
return loss_name
|
|
|
|
|
|
def validate_disjoint_sessions(
|
|
first: Sequence[str],
|
|
second: Sequence[str],
|
|
) -> None:
|
|
overlap = set(first) & set(second)
|
|
if overlap:
|
|
raise ValueError(f"Session sets overlap: {sorted(overlap)}")
|