feat: add causal frame set attention fusion
This commit is contained in:
Xin Wang
@@ -1,4 +1,4 @@
|
||||
"""Validated, causal local-feature sidecars for the pooled B1 baseline.
|
||||
"""Validated, causal pooled and frame-set local-feature sidecars.
|
||||
|
||||
No model/decoder imports: these utilities are also used by CPU head training.
|
||||
Sidecars store encoder-native vectors, never a randomly initialized projection.
|
||||
@@ -6,6 +6,7 @@ Sidecars store encoder-native vectors, never a randomly initialized projection.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -16,6 +17,150 @@ LOCAL_FORMAT = "qwen-local-pooled-v1"
|
||||
MODALITIES = {"visual": ("visual",), "audio": ("audio",), "av": ("visual", "audio")}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FramePersonFeatures:
|
||||
"""Ragged person sets indexed by aligned 5 Hz rows without frame duplication."""
|
||||
|
||||
features: torch.Tensor
|
||||
frame_offsets: torch.LongTensor
|
||||
row_to_frame: torch.LongTensor
|
||||
chosen_frame_times: torch.Tensor
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if (
|
||||
not isinstance(self.features, torch.Tensor)
|
||||
or self.features.ndim != 2
|
||||
or not self.features.is_floating_point()
|
||||
or not torch.isfinite(self.features).all()
|
||||
or self.features.shape[1] <= 0
|
||||
):
|
||||
raise ValueError("Person features must be finite floating [N,D]")
|
||||
if (
|
||||
not isinstance(self.frame_offsets, torch.Tensor)
|
||||
or self.frame_offsets.dtype != torch.long
|
||||
or self.frame_offsets.ndim != 1
|
||||
or not len(self.frame_offsets)
|
||||
or int(self.frame_offsets[0]) != 0
|
||||
or int(self.frame_offsets[-1]) != len(self.features)
|
||||
or bool((self.frame_offsets[1:] < self.frame_offsets[:-1]).any())
|
||||
):
|
||||
raise ValueError("Invalid person frame_offsets")
|
||||
if (
|
||||
not isinstance(self.row_to_frame, torch.Tensor)
|
||||
or self.row_to_frame.dtype != torch.long
|
||||
or self.row_to_frame.ndim != 1
|
||||
):
|
||||
raise ValueError("row_to_frame must be int64 [T]")
|
||||
num_frames = len(self.frame_offsets) - 1
|
||||
if bool(((self.row_to_frame < -1) | (self.row_to_frame >= num_frames)).any()):
|
||||
raise ValueError("row_to_frame contains an invalid frame index")
|
||||
if (
|
||||
not isinstance(self.chosen_frame_times, torch.Tensor)
|
||||
or self.chosen_frame_times.dtype != torch.float64
|
||||
or self.chosen_frame_times.shape != self.row_to_frame.shape
|
||||
):
|
||||
raise ValueError("chosen_frame_times must be float64 [T]")
|
||||
missing = self.row_to_frame < 0
|
||||
if not torch.equal(torch.isnan(self.chosen_frame_times), missing):
|
||||
raise ValueError("Only rows without an available frame may have NaN frame time")
|
||||
if not torch.isfinite(self.chosen_frame_times[~missing]).all():
|
||||
raise ValueError("Chosen person frame times must be finite")
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.row_to_frame)
|
||||
|
||||
@property
|
||||
def dimension(self) -> int:
|
||||
return int(self.features.shape[1])
|
||||
|
||||
@property
|
||||
def max_people(self) -> int:
|
||||
counts = self.frame_offsets[1:] - self.frame_offsets[:-1]
|
||||
return int(counts.max()) if len(counts) else 0
|
||||
|
||||
def select_rows(self, indices: torch.Tensor) -> "FramePersonFeatures":
|
||||
indices = torch.as_tensor(indices)
|
||||
if indices.dtype != torch.long or indices.ndim != 1:
|
||||
raise ValueError("Person row selection must be int64 [N]")
|
||||
if len(indices) and (int(indices.min()) < 0 or int(indices.max()) >= len(self)):
|
||||
raise IndexError("Person row selection is out of range")
|
||||
return FramePersonFeatures(
|
||||
self.features,
|
||||
self.frame_offsets,
|
||||
self.row_to_frame[indices],
|
||||
self.chosen_frame_times[indices],
|
||||
)
|
||||
|
||||
def padded(
|
||||
self,
|
||||
row_indices: torch.Tensor | None = None,
|
||||
) -> tuple[torch.Tensor, torch.BoolTensor]:
|
||||
"""Materialize only one batch as [B,P,D] plus a True-for-real mask."""
|
||||
|
||||
if row_indices is None:
|
||||
frames = self.row_to_frame
|
||||
else:
|
||||
row_indices = torch.as_tensor(row_indices)
|
||||
if row_indices.dtype != torch.long or row_indices.ndim != 1:
|
||||
raise ValueError("Batch row indices must be int64 [B]")
|
||||
if len(row_indices) and (
|
||||
int(row_indices.min()) < 0 or int(row_indices.max()) >= len(self)
|
||||
):
|
||||
raise IndexError("Batch person row index is out of range")
|
||||
frames = self.row_to_frame[row_indices]
|
||||
valid_frames = frames >= 0
|
||||
safe_frames = frames.clamp_min(0)
|
||||
starts = torch.zeros_like(frames)
|
||||
counts = torch.zeros_like(frames)
|
||||
if bool(valid_frames.any()):
|
||||
starts[valid_frames] = self.frame_offsets[safe_frames[valid_frames]]
|
||||
counts[valid_frames] = (
|
||||
self.frame_offsets[safe_frames[valid_frames] + 1]
|
||||
- starts[valid_frames]
|
||||
)
|
||||
width = int(counts.max()) if len(counts) else 0
|
||||
mask = torch.arange(width)[None, :] < counts[:, None]
|
||||
if width == 0:
|
||||
return self.features.new_zeros((len(frames), 0, self.dimension)), mask
|
||||
person_indices = starts[:, None] + torch.arange(width)[None, :]
|
||||
safe_indices = person_indices.masked_fill(~mask, 0)
|
||||
values = self.features[safe_indices]
|
||||
return values.masked_fill(~mask[..., None], 0), mask
|
||||
|
||||
@classmethod
|
||||
def concatenate(
|
||||
cls,
|
||||
parts: list["FramePersonFeatures"],
|
||||
) -> "FramePersonFeatures":
|
||||
if not parts:
|
||||
raise ValueError("Cannot concatenate no person feature sets")
|
||||
dimension = parts[0].dimension
|
||||
if any(part.dimension != dimension for part in parts):
|
||||
raise ValueError("Person feature dimensions differ across sessions")
|
||||
feature_parts = []
|
||||
offset_parts = [torch.zeros(1, dtype=torch.long)]
|
||||
row_parts = []
|
||||
time_parts = []
|
||||
feature_base = 0
|
||||
frame_base = 0
|
||||
for part in parts:
|
||||
feature_parts.append(part.features)
|
||||
offset_parts.append(part.frame_offsets[1:] + feature_base)
|
||||
rows = part.row_to_frame.clone()
|
||||
valid = rows >= 0
|
||||
rows[valid] += frame_base
|
||||
row_parts.append(rows)
|
||||
time_parts.append(part.chosen_frame_times)
|
||||
feature_base += len(part.features)
|
||||
frame_base += len(part.frame_offsets) - 1
|
||||
return cls(
|
||||
torch.cat(feature_parts, dim=0),
|
||||
torch.cat(offset_parts, dim=0),
|
||||
torch.cat(row_parts, dim=0),
|
||||
torch.cat(time_parts, dim=0),
|
||||
)
|
||||
|
||||
|
||||
def reference_alignment(payload: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Require the same full, zero-origin 5 Hz timeline used by head training."""
|
||||
states = payload.get("states")
|
||||
@@ -34,16 +179,15 @@ def reference_alignment(payload: dict[str, Any]) -> tuple[torch.Tensor, torch.Te
|
||||
return times, indices
|
||||
|
||||
|
||||
def load_local_features(
|
||||
def _load_validated_payload(
|
||||
path: Path,
|
||||
reference: dict[str, Any],
|
||||
*,
|
||||
sample_id: str,
|
||||
dataset: str,
|
||||
modalities: str,
|
||||
expected_contract: dict[str, Any] | None = None,
|
||||
) -> tuple[torch.Tensor, dict[str, Any]]:
|
||||
"""Pack each requested modality as [raw vector, validity bit]. Fail closed."""
|
||||
expected_contract: dict[str, Any] | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], dict[str, Any], dict[str, Any]]:
|
||||
if modalities not in MODALITIES:
|
||||
raise ValueError(f"Unknown local modalities: {modalities}")
|
||||
times, indices = reference_alignment(reference)
|
||||
@@ -66,6 +210,27 @@ def load_local_features(
|
||||
contract = {"format": LOCAL_FORMAT, "modalities": modalities, "config": config}
|
||||
if expected_contract is not None and contract != expected_contract:
|
||||
raise ValueError(f"Local encoder/pooling contract differs across sessions/checkpoint: {path}")
|
||||
return times, indices, payload, config, contract
|
||||
|
||||
|
||||
def load_local_features(
|
||||
path: Path,
|
||||
reference: dict[str, Any],
|
||||
*,
|
||||
sample_id: str,
|
||||
dataset: str,
|
||||
modalities: str,
|
||||
expected_contract: dict[str, Any] | None = None,
|
||||
) -> tuple[torch.Tensor, dict[str, Any]]:
|
||||
"""Pack each requested modality as [raw vector, validity bit]. Fail closed."""
|
||||
times, _, payload, config, contract = _load_validated_payload(
|
||||
path,
|
||||
reference,
|
||||
sample_id=sample_id,
|
||||
dataset=dataset,
|
||||
modalities=modalities,
|
||||
expected_contract=expected_contract,
|
||||
)
|
||||
chunks = []
|
||||
for name in MODALITIES[modalities]:
|
||||
values = payload.get(name)
|
||||
@@ -85,6 +250,109 @@ def load_local_features(
|
||||
return torch.cat(chunks, dim=-1), contract
|
||||
|
||||
|
||||
def load_frame_set_features(
|
||||
path: Path,
|
||||
reference: dict[str, Any],
|
||||
*,
|
||||
sample_id: str,
|
||||
dataset: str,
|
||||
expected_contract: dict[str, Any] | None = None,
|
||||
) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
|
||||
"""Load direct audio and the latest causal per-frame person set per 5 Hz row."""
|
||||
|
||||
times, _, payload, config, contract = _load_validated_payload(
|
||||
path,
|
||||
reference,
|
||||
sample_id=sample_id,
|
||||
dataset=dataset,
|
||||
modalities="av",
|
||||
expected_contract=expected_contract,
|
||||
)
|
||||
audio = payload.get("audio")
|
||||
audio_valid = payload.get("audio_valid")
|
||||
audio_dim = config.get("audio_dim")
|
||||
if not isinstance(audio_dim, int) or audio_dim <= 0:
|
||||
raise ValueError(f"Missing audio_dim: {path}")
|
||||
if not isinstance(audio, torch.Tensor) or audio.shape != (len(times), audio_dim):
|
||||
raise ValueError(f"Invalid audio shape: {path}")
|
||||
if not audio.is_floating_point() or not torch.isfinite(audio).all():
|
||||
raise ValueError(f"Non-finite/non-floating audio features: {path}")
|
||||
if (
|
||||
not isinstance(audio_valid, torch.Tensor)
|
||||
or audio_valid.dtype != torch.bool
|
||||
or audio_valid.shape != (len(times),)
|
||||
):
|
||||
raise ValueError(f"Invalid audio_valid mask: {path}")
|
||||
if torch.count_nonzero(audio[~audio_valid]):
|
||||
raise ValueError(f"Invalid audio rows must be stored as zeros: {path}")
|
||||
|
||||
visual_dim = config.get("visual_dim")
|
||||
persons = payload.get("persons")
|
||||
if not isinstance(visual_dim, int) or visual_dim <= 0:
|
||||
raise ValueError(f"Missing visual_dim: {path}")
|
||||
if not isinstance(persons, dict):
|
||||
raise ValueError(f"Missing ragged person features: {path}")
|
||||
features = persons.get("features")
|
||||
frame_times = persons.get("frame_times")
|
||||
frame_offsets = persons.get("frame_offsets")
|
||||
records = persons.get("records")
|
||||
if (
|
||||
not isinstance(features, torch.Tensor)
|
||||
or features.ndim != 2
|
||||
or features.shape[1] != visual_dim
|
||||
or not features.is_floating_point()
|
||||
or not torch.isfinite(features).all()
|
||||
):
|
||||
raise ValueError(f"Invalid person feature shape/content: {path}")
|
||||
if (
|
||||
not isinstance(frame_times, torch.Tensor)
|
||||
or frame_times.ndim != 1
|
||||
or not frame_times.is_floating_point()
|
||||
):
|
||||
raise ValueError(f"Invalid person frame_times: {path}")
|
||||
frame_times = frame_times.to(dtype=torch.float64)
|
||||
if (
|
||||
not torch.isfinite(frame_times).all()
|
||||
or bool((frame_times < 0).any())
|
||||
or bool((frame_times[1:] <= frame_times[:-1]).any())
|
||||
):
|
||||
raise ValueError(f"Person frame_times must be finite, nonnegative and increasing: {path}")
|
||||
if (
|
||||
not isinstance(frame_offsets, torch.Tensor)
|
||||
or frame_offsets.dtype != torch.long
|
||||
or frame_offsets.shape != (len(frame_times) + 1,)
|
||||
or not len(frame_offsets)
|
||||
or int(frame_offsets[0]) != 0
|
||||
or int(frame_offsets[-1]) != len(features)
|
||||
or bool((frame_offsets[1:] < frame_offsets[:-1]).any())
|
||||
):
|
||||
raise ValueError(f"Invalid person frame_offsets: {path}")
|
||||
if not isinstance(records, list) or len(records) != len(features):
|
||||
raise ValueError(f"Person records/features length mismatch: {path}")
|
||||
if len(frame_times) and float(frame_times[-1]) > float(times[-1]) + 1e-8:
|
||||
raise ValueError(f"Person frame timeline extends beyond the reference: {path}")
|
||||
|
||||
# Match the sidecar's half-open policy: a frame stamped exactly at the
|
||||
# state cutoff is not yet part of that state's observed prefix.
|
||||
row_to_frame = torch.searchsorted(frame_times, times, right=False) - 1
|
||||
chosen = torch.full(times.shape, float("nan"), dtype=torch.float64)
|
||||
has_frame = row_to_frame >= 0
|
||||
chosen[has_frame] = frame_times[row_to_frame[has_frame]]
|
||||
if bool((chosen[has_frame] >= times[has_frame]).any()):
|
||||
raise ValueError(f"Person frame mapping selected future visual evidence: {path}")
|
||||
frame_sets = FramePersonFeatures(
|
||||
features,
|
||||
frame_offsets,
|
||||
row_to_frame.to(dtype=torch.long),
|
||||
chosen,
|
||||
)
|
||||
direct_audio = torch.cat(
|
||||
(audio.float(), audio_valid[:, None].float()),
|
||||
dim=-1,
|
||||
)
|
||||
return direct_audio, frame_sets, contract
|
||||
|
||||
|
||||
def local_head_dimensions(contract: dict[str, Any] | None) -> dict[str, int]:
|
||||
if contract is None:
|
||||
return {"local_visual_dim": 0, "local_audio_dim": 0}
|
||||
|
||||
Reference in New Issue
Block a user