Add causal person fusion and MuVAP stage-3 baselines
This commit is contained in:
Xin Wang
671
muvap_egospeak_stage3.py
Normal file
671
muvap_egospeak_stage3.py
Normal file
@@ -0,0 +1,671 @@
|
||||
"""Shared data/model code for the MuVAP-only EgoSpeak Stage-3 baseline."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from easycom_head import (
|
||||
CLASS_NAMES,
|
||||
NUM_CLASSES,
|
||||
NUM_HORIZONS,
|
||||
STEP_SECONDS,
|
||||
_validated_binary_activity_targets,
|
||||
average_precision_metrics,
|
||||
marginal_probabilities_from_logits,
|
||||
normalize_sample_name,
|
||||
)
|
||||
MODEL_FRAME_HZ = 25.0
|
||||
MODEL_FRAMES_PER_TASK_STEP = 5
|
||||
MUVAP_EGOSPEAK_INPUT_FORMAT = "muvap-egospeak-stage3-paper-reconstruction-v1"
|
||||
MODEL_KIND = "muvap-egospeak-stage3-paper-reconstruction"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MuVAPSessionFeatures:
|
||||
name: str
|
||||
vap: torch.Tensor
|
||||
people: torch.Tensor
|
||||
person_valid: torch.BoolTensor
|
||||
targets: torch.Tensor
|
||||
task_frame_indices: torch.LongTensor
|
||||
label_indices: torch.LongTensor
|
||||
times: torch.Tensor
|
||||
contract: dict[str, Any]
|
||||
|
||||
@property
|
||||
def vap_dim(self) -> int:
|
||||
return int(self.vap.shape[-1])
|
||||
|
||||
@property
|
||||
def asd_dim(self) -> int:
|
||||
return int(self.people.shape[-1])
|
||||
|
||||
|
||||
def _validate_exact_grid(values: torch.Tensor, expected: torch.Tensor, name: str) -> None:
|
||||
if values.dtype != torch.float64 or values.shape != expected.shape:
|
||||
raise ValueError(f"Invalid {name} shape/dtype")
|
||||
if not torch.allclose(values, expected, atol=1e-8, rtol=0):
|
||||
raise ValueError(f"{name} does not follow the required right-edge grid")
|
||||
|
||||
|
||||
def load_muvap_session(
|
||||
feature_root: str | Path,
|
||||
target_root: str | Path,
|
||||
session: str | int,
|
||||
*,
|
||||
expected_contract: dict[str, Any] | None = None,
|
||||
allow_partial: bool = False,
|
||||
) -> MuVAPSessionFeatures:
|
||||
"""Load one 25 Hz frozen-backbone sidecar and fail closed on alignment."""
|
||||
|
||||
name = normalize_sample_name(session)
|
||||
feature_path = Path(feature_root) / f"{name}.pt"
|
||||
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 payload.get("format") != MUVAP_EGOSPEAK_INPUT_FORMAT:
|
||||
raise ValueError(f"Not a MuVAP EgoSpeak Stage-3 sidecar: {feature_path}")
|
||||
if payload.get("dataset") != "easycom" or payload.get("sample_id") != name:
|
||||
raise ValueError(f"Dataset/sample identity mismatch: {feature_path}")
|
||||
if bool(payload.get("partial", False)) and not allow_partial:
|
||||
raise ValueError(f"Partial sidecar cannot be used for training: {feature_path}")
|
||||
contract = {
|
||||
"format": payload["format"],
|
||||
"dataset": payload["dataset"],
|
||||
"config": payload.get("config"),
|
||||
}
|
||||
if not isinstance(contract["config"], dict):
|
||||
raise ValueError(f"Missing MuVAP extraction config: {feature_path}")
|
||||
if contract["config"].get("uses_omni_features") is not False:
|
||||
raise ValueError(f"MuVAP-only sidecar does not exclude Omni: {feature_path}")
|
||||
if contract["config"].get("uses_speech_labels") is not False:
|
||||
raise ValueError(f"Feature extractor consumed speech labels: {feature_path}")
|
||||
if expected_contract is not None and contract != expected_contract:
|
||||
raise ValueError(f"MuVAP feature contract differs across sessions: {feature_path}")
|
||||
|
||||
vap = payload.get("vap")
|
||||
if (
|
||||
not isinstance(vap, torch.Tensor)
|
||||
or vap.ndim != 2
|
||||
or not vap.is_floating_point()
|
||||
or not torch.isfinite(vap).all()
|
||||
):
|
||||
raise ValueError(f"vap must be finite floating [K,D]: {feature_path}")
|
||||
num_ticks = len(vap)
|
||||
frame_times = payload.get("frame_times")
|
||||
if not isinstance(frame_times, torch.Tensor):
|
||||
raise ValueError(f"Missing frame_times: {feature_path}")
|
||||
_validate_exact_grid(
|
||||
frame_times,
|
||||
torch.arange(1, num_ticks + 1, dtype=torch.float64) / MODEL_FRAME_HZ,
|
||||
"MuVAP frame_times",
|
||||
)
|
||||
|
||||
times = payload.get("times")
|
||||
label_indices = payload.get("label_indices")
|
||||
task_frame_indices = payload.get("task_frame_indices")
|
||||
if not all(isinstance(value, torch.Tensor) for value in (times, label_indices, task_frame_indices)):
|
||||
raise ValueError(f"Missing task alignment tensors: {feature_path}")
|
||||
num_steps = len(times)
|
||||
_validate_exact_grid(
|
||||
times,
|
||||
torch.arange(1, num_steps + 1, dtype=torch.float64) * STEP_SECONDS,
|
||||
"task times",
|
||||
)
|
||||
expected_labels = torch.arange(num_steps, dtype=torch.long)
|
||||
expected_frames = (expected_labels + 1) * MODEL_FRAMES_PER_TASK_STEP - 1
|
||||
if label_indices.dtype != torch.long or not torch.equal(label_indices, expected_labels):
|
||||
raise ValueError(f"label_indices must be dense and zero-based: {feature_path}")
|
||||
if task_frame_indices.dtype != torch.long or not torch.equal(
|
||||
task_frame_indices, expected_frames
|
||||
):
|
||||
raise ValueError(f"Invalid task-to-25Hz mapping: {feature_path}")
|
||||
if num_ticks != num_steps * MODEL_FRAMES_PER_TASK_STEP:
|
||||
raise ValueError(f"25Hz/5Hz length mismatch: {feature_path}")
|
||||
|
||||
person = payload.get("persons")
|
||||
if not isinstance(person, dict):
|
||||
raise ValueError(f"Missing persons block: {feature_path}")
|
||||
features = person.get("features")
|
||||
offsets = person.get("frame_offsets")
|
||||
records = person.get("records")
|
||||
if (
|
||||
not isinstance(features, torch.Tensor)
|
||||
or features.ndim != 2
|
||||
or not features.is_floating_point()
|
||||
or not torch.isfinite(features).all()
|
||||
):
|
||||
raise ValueError(f"Invalid fused ASD features: {feature_path}")
|
||||
if (
|
||||
not isinstance(offsets, torch.Tensor)
|
||||
or offsets.dtype != torch.long
|
||||
or offsets.shape != (num_ticks + 1,)
|
||||
or int(offsets[0]) != 0
|
||||
or int(offsets[-1]) != len(features)
|
||||
or bool((offsets[1:] < offsets[:-1]).any())
|
||||
):
|
||||
raise ValueError(f"Invalid person frame offsets: {feature_path}")
|
||||
if not isinstance(records, list) or len(records) != len(features):
|
||||
raise ValueError(f"Person records/features differ: {feature_path}")
|
||||
config = contract["config"]
|
||||
if config.get("vap_dim") != vap.shape[1] or config.get("asd_dim") != features.shape[1]:
|
||||
raise ValueError(f"Saved feature dimensions disagree with config: {feature_path}")
|
||||
if not config.get("vap_release_sha256") or not config.get("asd_release_sha256"):
|
||||
raise ValueError(f"Missing frozen release identity: {feature_path}")
|
||||
counts = offsets[1:] - offsets[:-1]
|
||||
max_people = int(counts.max()) if len(counts) else 0
|
||||
saved_max = payload.get("max_people")
|
||||
if isinstance(saved_max, bool) or not isinstance(saved_max, int) or saved_max != max_people:
|
||||
raise ValueError(f"Incorrect max_people metadata: {feature_path}")
|
||||
people = torch.zeros(num_ticks, max_people, features.shape[1], dtype=features.dtype)
|
||||
person_valid = torch.zeros(num_ticks, max_people, dtype=torch.bool)
|
||||
for tick in range(num_ticks):
|
||||
start, end = int(offsets[tick]), int(offsets[tick + 1])
|
||||
count = end - start
|
||||
if count:
|
||||
seen: set[str] = set()
|
||||
expected_edge = (tick + 1) / MODEL_FRAME_HZ
|
||||
for record in records[start:end]:
|
||||
if not isinstance(record, dict):
|
||||
raise ValueError(f"Invalid person record at tick {tick}: {feature_path}")
|
||||
person_id = str(record.get("id", ""))
|
||||
if person_id in seen:
|
||||
raise ValueError(f"Duplicate person ID at tick {tick}: {feature_path}")
|
||||
seen.add(person_id)
|
||||
if record.get("asd_model_tick") != tick:
|
||||
raise ValueError(f"Person record assigned to wrong tick: {feature_path}")
|
||||
edge = record.get("asd_interval_end")
|
||||
source_time = record.get("source_time")
|
||||
if (
|
||||
not isinstance(edge, (int, float))
|
||||
or not isinstance(source_time, (int, float))
|
||||
or abs(float(edge) - expected_edge) > 1e-8
|
||||
or float(source_time) >= expected_edge - 1e-8
|
||||
):
|
||||
raise ValueError(f"Non-causal person record at tick {tick}: {feature_path}")
|
||||
people[tick, :count] = features[start:end]
|
||||
person_valid[tick, :count] = True
|
||||
|
||||
target_array = np.load(target_path, allow_pickle=False)
|
||||
targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
|
||||
if targets.shape != (num_steps, NUM_CLASSES):
|
||||
raise ValueError(
|
||||
f"Target/feature length mismatch {targets.shape} vs ({num_steps},3): {name}"
|
||||
)
|
||||
_validated_binary_activity_targets(targets)
|
||||
return MuVAPSessionFeatures(
|
||||
name=name,
|
||||
vap=vap,
|
||||
people=people,
|
||||
person_valid=person_valid,
|
||||
targets=targets,
|
||||
task_frame_indices=task_frame_indices,
|
||||
label_indices=label_indices,
|
||||
times=times,
|
||||
contract=contract,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Stage3Window:
|
||||
name: str
|
||||
current_row: int
|
||||
vap: torch.Tensor
|
||||
people: torch.Tensor
|
||||
person_valid: torch.BoolTensor
|
||||
history_valid: torch.BoolTensor
|
||||
targets: torch.Tensor
|
||||
future_valid: torch.BoolTensor
|
||||
|
||||
|
||||
class MuVAPWindowDataset(Dataset[Stage3Window]):
|
||||
"""One fixed past-only 25 Hz window for every 5 Hz prediction cutoff."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sessions: Sequence[MuVAPSessionFeatures],
|
||||
*,
|
||||
history_seconds: float,
|
||||
require_full_future: bool,
|
||||
) -> None:
|
||||
if not sessions:
|
||||
raise ValueError("At least one session is required")
|
||||
history_frames = round(history_seconds * MODEL_FRAME_HZ)
|
||||
if history_frames <= 0 or not math.isclose(
|
||||
history_frames / MODEL_FRAME_HZ, history_seconds, abs_tol=1e-8
|
||||
):
|
||||
raise ValueError("history_seconds must be a positive 25 Hz multiple")
|
||||
self.sessions = list(sessions)
|
||||
self.history_frames = history_frames
|
||||
self.require_full_future = require_full_future
|
||||
self.examples: list[tuple[int, int]] = []
|
||||
for session_index, session in enumerate(self.sessions):
|
||||
stop = len(session.targets) - NUM_HORIZONS if require_full_future else len(session.targets)
|
||||
self.examples.extend((session_index, row) for row in range(max(0, stop)))
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.examples)
|
||||
|
||||
def __getitem__(self, index: int) -> Stage3Window:
|
||||
session_index, row = self.examples[index]
|
||||
session = self.sessions[session_index]
|
||||
end = int(session.task_frame_indices[row]) + 1
|
||||
start = max(0, end - self.history_frames)
|
||||
observed = end - start
|
||||
left = self.history_frames - observed
|
||||
vap = torch.zeros(self.history_frames, session.vap_dim, dtype=session.vap.dtype)
|
||||
vap[left:] = session.vap[start:end]
|
||||
people = torch.zeros(
|
||||
self.history_frames,
|
||||
session.people.shape[1],
|
||||
session.asd_dim,
|
||||
dtype=session.people.dtype,
|
||||
)
|
||||
person_valid = torch.zeros(
|
||||
self.history_frames, session.people.shape[1], dtype=torch.bool
|
||||
)
|
||||
if session.people.shape[1]:
|
||||
people[left:] = session.people[start:end]
|
||||
person_valid[left:] = session.person_valid[start:end]
|
||||
history_valid = torch.zeros(self.history_frames, dtype=torch.bool)
|
||||
history_valid[left:] = True
|
||||
targets = torch.zeros(NUM_HORIZONS, NUM_CLASSES, dtype=torch.float32)
|
||||
future_valid = torch.zeros(NUM_HORIZONS, dtype=torch.bool)
|
||||
available = min(NUM_HORIZONS, len(session.targets) - row - 1)
|
||||
if available:
|
||||
targets[:available] = session.targets[row + 1 : row + 1 + available]
|
||||
future_valid[:available] = True
|
||||
if self.require_full_future and not bool(future_valid.all()):
|
||||
raise RuntimeError("Training dataset exposed an incomplete future target")
|
||||
return Stage3Window(
|
||||
name=session.name,
|
||||
current_row=row,
|
||||
vap=vap,
|
||||
people=people,
|
||||
person_valid=person_valid,
|
||||
history_valid=history_valid,
|
||||
targets=targets,
|
||||
future_valid=future_valid,
|
||||
)
|
||||
|
||||
|
||||
def collate_stage3_windows(items: Sequence[Stage3Window]) -> dict[str, Any]:
|
||||
if not items:
|
||||
raise ValueError("Cannot collate an empty batch")
|
||||
max_people = max(item.people.shape[1] for item in items)
|
||||
batch_people = torch.zeros(
|
||||
len(items),
|
||||
items[0].people.shape[0],
|
||||
max_people,
|
||||
items[0].people.shape[-1],
|
||||
dtype=items[0].people.dtype,
|
||||
)
|
||||
batch_person_valid = torch.zeros(
|
||||
len(items), items[0].people.shape[0], max_people, dtype=torch.bool
|
||||
)
|
||||
for index, item in enumerate(items):
|
||||
count = item.people.shape[1]
|
||||
if count:
|
||||
batch_people[index, :, :count] = item.people
|
||||
batch_person_valid[index, :, :count] = item.person_valid
|
||||
return {
|
||||
"names": [item.name for item in items],
|
||||
"current_rows": torch.tensor([item.current_row for item in items], dtype=torch.long),
|
||||
"vap": torch.stack([item.vap for item in items]),
|
||||
"people": batch_people,
|
||||
"person_valid": batch_person_valid,
|
||||
"history_valid": torch.stack([item.history_valid for item in items]),
|
||||
"targets": torch.stack([item.targets for item in items]),
|
||||
"future_valid": torch.stack([item.future_valid for item in items]),
|
||||
}
|
||||
|
||||
|
||||
def sinusoidal_positions(length: int, dimension: int) -> torch.Tensor:
|
||||
position = torch.arange(length, dtype=torch.float32)[:, None]
|
||||
frequencies = torch.exp(
|
||||
torch.arange(0, dimension, 2, dtype=torch.float32)
|
||||
* (-math.log(10_000.0) / dimension)
|
||||
)
|
||||
result = torch.zeros(length, dimension, dtype=torch.float32)
|
||||
result[:, 0::2] = torch.sin(position * frequencies)
|
||||
result[:, 1::2] = torch.cos(position * frequencies[: result[:, 1::2].shape[1]])
|
||||
return result
|
||||
|
||||
|
||||
class MuVAPEgoSpeakStage3(nn.Module):
|
||||
"""Paper-inspired MuVAP frame/temporal fusion with an EgoSpeak head.
|
||||
|
||||
Published VAP and ASD states are inputs and remain outside this trainable
|
||||
module. The model trains only the Stage-3 projections, same-frame person
|
||||
cross-attention, causal temporal transformer and the task head.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
vap_dim: int,
|
||||
asd_dim: int,
|
||||
model_dim: int = 256,
|
||||
num_heads: int = 4,
|
||||
dropout: float = 0.1,
|
||||
temporal_layers: int = 1,
|
||||
history_frames: int = 200,
|
||||
num_classes: int = NUM_CLASSES,
|
||||
person_evidence: str = "aligned",
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if min(vap_dim, asd_dim, model_dim, num_heads, temporal_layers, history_frames) <= 0:
|
||||
raise ValueError("Stage-3 dimensions must be positive")
|
||||
if model_dim % num_heads:
|
||||
raise ValueError("model_dim must be divisible by num_heads")
|
||||
if person_evidence not in {"aligned", "none"}:
|
||||
raise ValueError("person_evidence must be 'aligned' or 'none'")
|
||||
self.vap_dim = vap_dim
|
||||
self.asd_dim = asd_dim
|
||||
self.model_dim = model_dim
|
||||
self.num_heads = num_heads
|
||||
self.dropout_probability = dropout
|
||||
self.temporal_layers = temporal_layers
|
||||
self.history_frames = history_frames
|
||||
self.num_classes = num_classes
|
||||
self.person_evidence = person_evidence
|
||||
self.vap_projection = nn.Sequential(
|
||||
nn.Linear(vap_dim, model_dim),
|
||||
nn.LayerNorm(model_dim),
|
||||
)
|
||||
self.asd_projection = nn.Sequential(
|
||||
nn.Linear(asd_dim, model_dim),
|
||||
nn.LayerNorm(model_dim),
|
||||
)
|
||||
self.null_person = nn.Parameter(torch.zeros(1, 1, model_dim))
|
||||
nn.init.normal_(self.null_person, std=0.02)
|
||||
self.frame_attention = nn.MultiheadAttention(
|
||||
model_dim,
|
||||
num_heads,
|
||||
dropout=dropout,
|
||||
batch_first=True,
|
||||
)
|
||||
self.frame_attention_norm = nn.LayerNorm(model_dim)
|
||||
self.frame_ff_norm = nn.LayerNorm(model_dim)
|
||||
self.frame_ff = nn.Sequential(
|
||||
nn.Linear(model_dim, model_dim * 4),
|
||||
nn.GELU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(model_dim * 4, model_dim),
|
||||
)
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
temporal_layer = nn.TransformerEncoderLayer(
|
||||
d_model=model_dim,
|
||||
nhead=num_heads,
|
||||
dim_feedforward=model_dim * 4,
|
||||
dropout=dropout,
|
||||
activation="gelu",
|
||||
batch_first=True,
|
||||
norm_first=True,
|
||||
)
|
||||
self.temporal = nn.TransformerEncoder(
|
||||
temporal_layer,
|
||||
num_layers=temporal_layers,
|
||||
norm=nn.LayerNorm(model_dim),
|
||||
)
|
||||
self.head = nn.Linear(model_dim, NUM_HORIZONS * num_classes)
|
||||
self.register_buffer(
|
||||
"position_encoding",
|
||||
sinusoidal_positions(history_frames, model_dim),
|
||||
persistent=False,
|
||||
)
|
||||
self.register_buffer(
|
||||
"causal_mask",
|
||||
torch.triu(torch.ones(history_frames, history_frames, dtype=torch.bool), diagonal=1),
|
||||
persistent=False,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
vap: torch.Tensor,
|
||||
people: torch.Tensor,
|
||||
person_valid: torch.BoolTensor,
|
||||
history_valid: torch.BoolTensor,
|
||||
*,
|
||||
return_embedding: bool = False,
|
||||
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
|
||||
if vap.ndim != 3 or vap.shape[-1] != self.vap_dim:
|
||||
raise ValueError(f"Expected VAP [B,H,{self.vap_dim}], got {tuple(vap.shape)}")
|
||||
batch, history, _ = vap.shape
|
||||
if history != self.history_frames:
|
||||
raise ValueError(f"Expected {self.history_frames} history frames, got {history}")
|
||||
if people.ndim != 4 or people.shape[:2] != (batch, history) or people.shape[-1] != self.asd_dim:
|
||||
raise ValueError(
|
||||
f"Expected people [B,H,N,{self.asd_dim}], got {tuple(people.shape)}"
|
||||
)
|
||||
if person_valid.dtype != torch.bool or person_valid.shape != people.shape[:3]:
|
||||
raise ValueError("person_valid must be bool [B,H,N]")
|
||||
if history_valid.dtype != torch.bool or history_valid.shape != (batch, history):
|
||||
raise ValueError("history_valid must be bool [B,H]")
|
||||
if not bool(history_valid[:, -1].all()):
|
||||
raise ValueError("Every Stage-3 window must end at an observed cutoff")
|
||||
if not torch.isfinite(vap).all() or not torch.isfinite(people).all():
|
||||
raise ValueError("Stage-3 inputs contain NaN/Inf")
|
||||
if self.person_evidence == "none":
|
||||
# Capacity-matched VAP-only control. The same frame and temporal
|
||||
# blocks remain trainable; only observed ASD tokens are hidden.
|
||||
person_valid = torch.zeros_like(person_valid)
|
||||
|
||||
query = self.vap_projection(vap).reshape(batch * history, 1, self.model_dim)
|
||||
count = people.shape[2]
|
||||
person_tokens = self.asd_projection(people).reshape(
|
||||
batch * history, count, self.model_dim
|
||||
)
|
||||
null = self.null_person.expand(batch * history, -1, -1)
|
||||
person_tokens = torch.cat((person_tokens, null), dim=1)
|
||||
key_padding = torch.cat(
|
||||
(
|
||||
~person_valid.reshape(batch * history, count),
|
||||
torch.zeros(batch * history, 1, dtype=torch.bool, device=vap.device),
|
||||
),
|
||||
dim=1,
|
||||
)
|
||||
selected, _ = self.frame_attention(
|
||||
query,
|
||||
person_tokens,
|
||||
person_tokens,
|
||||
key_padding_mask=key_padding,
|
||||
need_weights=False,
|
||||
)
|
||||
frame = self.frame_attention_norm(query + self.dropout(selected))
|
||||
frame = frame + self.dropout(self.frame_ff(self.frame_ff_norm(frame)))
|
||||
frame = frame.reshape(batch, history, self.model_dim)
|
||||
frame = frame + self.position_encoding[None].to(frame.dtype)
|
||||
|
||||
# Windows are left-padded so the final token is always the current
|
||||
# 5 Hz cutoff. A plain key-padding mask leaves the earliest padded
|
||||
# queries with no visible key under causality, which can yield NaNs
|
||||
# (and can contaminate later layers). Give each padded query only its
|
||||
# own diagonal while every observed query still masks every padded
|
||||
# key. The mask is repeated because MultiheadAttention accepts one
|
||||
# [H,H] matrix per batch/head pair.
|
||||
temporal_mask = self.causal_mask[None].expand(batch, -1, -1).clone()
|
||||
temporal_mask |= (~history_valid)[:, None, :]
|
||||
padded_batch, padded_query = torch.where(~history_valid)
|
||||
temporal_mask[padded_batch, padded_query, padded_query] = False
|
||||
temporal_mask = temporal_mask.repeat_interleave(self.num_heads, dim=0)
|
||||
temporal = self.temporal(
|
||||
frame,
|
||||
mask=temporal_mask,
|
||||
)
|
||||
embedding = temporal[:, -1]
|
||||
logits = self.head(embedding).reshape(batch, NUM_HORIZONS, self.num_classes)
|
||||
if return_embedding:
|
||||
return logits, embedding
|
||||
return logits
|
||||
|
||||
|
||||
def model_config(model: MuVAPEgoSpeakStage3) -> dict[str, Any]:
|
||||
return {
|
||||
"vap_dim": model.vap_dim,
|
||||
"asd_dim": model.asd_dim,
|
||||
"model_dim": model.model_dim,
|
||||
"num_heads": model.num_heads,
|
||||
"dropout": model.dropout_probability,
|
||||
"temporal_layers": model.temporal_layers,
|
||||
"history_frames": model.history_frames,
|
||||
"num_classes": model.num_classes,
|
||||
"person_evidence": model.person_evidence,
|
||||
}
|
||||
|
||||
|
||||
def build_model(config: dict[str, Any]) -> MuVAPEgoSpeakStage3:
|
||||
required = {
|
||||
"vap_dim",
|
||||
"asd_dim",
|
||||
"model_dim",
|
||||
"num_heads",
|
||||
"dropout",
|
||||
"temporal_layers",
|
||||
"history_frames",
|
||||
"num_classes",
|
||||
"person_evidence",
|
||||
}
|
||||
if set(config) != required:
|
||||
raise ValueError(f"Invalid Stage-3 model config keys: {sorted(set(config) ^ required)}")
|
||||
return MuVAPEgoSpeakStage3(**config)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict_loader(
|
||||
model: MuVAPEgoSpeakStage3,
|
||||
loader: Iterable[dict[str, Any]],
|
||||
*,
|
||||
device: torch.device,
|
||||
loss_name: str,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
|
||||
model.eval()
|
||||
probabilities, targets, rows, names = [], [], [], []
|
||||
for batch in loader:
|
||||
logits = model(
|
||||
batch["vap"].to(device=device, dtype=torch.float32, non_blocking=True),
|
||||
batch["people"].to(device=device, dtype=torch.float32, non_blocking=True),
|
||||
batch["person_valid"].to(device=device, non_blocking=True),
|
||||
batch["history_valid"].to(device=device, non_blocking=True),
|
||||
)
|
||||
probabilities.append(marginal_probabilities_from_logits(logits, loss_name).cpu())
|
||||
targets.append(batch["targets"])
|
||||
rows.append(batch["current_rows"])
|
||||
names.extend(batch["names"])
|
||||
return (
|
||||
torch.cat(probabilities),
|
||||
torch.cat(targets),
|
||||
torch.cat(rows),
|
||||
names,
|
||||
)
|
||||
|
||||
|
||||
def evaluate_valid_predictions(
|
||||
probabilities: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
) -> dict[str, Any]:
|
||||
if probabilities.shape != targets.shape or probabilities.shape[-2:] != (
|
||||
NUM_HORIZONS,
|
||||
NUM_CLASSES,
|
||||
):
|
||||
raise ValueError("Evaluation expects matching [N,10,3] arrays")
|
||||
results = []
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
result = average_precision_metrics(
|
||||
probabilities[:, horizon].numpy(),
|
||||
targets[:, horizon].numpy(),
|
||||
)
|
||||
result["horizon"] = horizon + 1
|
||||
result["seconds"] = (horizon + 1) * STEP_SECONDS
|
||||
results.append(result)
|
||||
return {
|
||||
"horizons": results,
|
||||
"mean_map": float(np.mean([result["map"] for result in results])),
|
||||
}
|
||||
|
||||
|
||||
def align_session_predictions(
|
||||
probabilities: np.ndarray,
|
||||
num_steps: int,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
if probabilities.shape != (num_steps, NUM_HORIZONS, NUM_CLASSES):
|
||||
raise ValueError(
|
||||
f"Expected [{num_steps},{NUM_HORIZONS},{NUM_CLASSES}], got "
|
||||
f"{probabilities.shape}"
|
||||
)
|
||||
predictions = np.zeros((num_steps, NUM_CLASSES, NUM_HORIZONS), dtype=np.float32)
|
||||
valid = np.zeros((num_steps, NUM_HORIZONS), dtype=bool)
|
||||
current = np.arange(num_steps)
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
future = current + horizon + 1
|
||||
usable = future < num_steps
|
||||
predictions[future[usable], :, horizon] = probabilities[usable, horizon]
|
||||
valid[future[usable], horizon] = True
|
||||
return predictions, valid
|
||||
|
||||
|
||||
def evaluate_aligned(
|
||||
session_predictions: dict[str, np.ndarray],
|
||||
session_valid: dict[str, np.ndarray],
|
||||
session_targets: dict[str, np.ndarray],
|
||||
) -> dict[str, Any]:
|
||||
protocols: dict[str, Any] = {}
|
||||
for protocol in ("egospeak_full_timeline", "causal_valid_only"):
|
||||
horizons = []
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
if protocol == "egospeak_full_timeline":
|
||||
scores = np.concatenate(
|
||||
[value[:, :, horizon] for value in session_predictions.values()]
|
||||
)
|
||||
targets = np.concatenate(list(session_targets.values()))
|
||||
else:
|
||||
scores = np.concatenate(
|
||||
[
|
||||
session_predictions[name][session_valid[name][:, horizon], :, horizon]
|
||||
for name in session_predictions
|
||||
]
|
||||
)
|
||||
targets = np.concatenate(
|
||||
[
|
||||
session_targets[name][session_valid[name][:, horizon]]
|
||||
for name in session_predictions
|
||||
]
|
||||
)
|
||||
result = average_precision_metrics(scores, targets)
|
||||
result["horizon"] = horizon + 1
|
||||
result["seconds"] = (horizon + 1) * STEP_SECONDS
|
||||
horizons.append(result)
|
||||
protocols[protocol] = {
|
||||
"horizons": horizons,
|
||||
"mean_map": float(np.mean([result["map"] for result in horizons])),
|
||||
}
|
||||
return protocols
|
||||
|
||||
|
||||
def print_metrics(metrics: dict[str, Any]) -> None:
|
||||
for protocol in ("egospeak_full_timeline", "causal_valid_only"):
|
||||
print(f"\n{protocol}:")
|
||||
for result in metrics[protocol]["horizons"]:
|
||||
values = ", ".join(
|
||||
f"{name}={result['class_ap'][name] * 100:.5f}" for name in CLASS_NAMES
|
||||
)
|
||||
print(
|
||||
f"Action anticipation ({result['seconds']:.2f}s) perframe mAP: "
|
||||
f"{result['map'] * 100:.5f}"
|
||||
)
|
||||
print(f" AP: {values}; n={result['num_steps']}")
|
||||
print(
|
||||
f"Action anticipation (mean) perframe mAP: "
|
||||
f"{metrics[protocol]['mean_map'] * 100:.5f}"
|
||||
)
|
||||
Reference in New Issue
Block a user