Add reproducible EasyCom causal head pipeline

2026-08-16 22:19:43 +08:00
parent bde72487bb
commit 080c45f049
8 changed files with 1451 additions and 470 deletions
--- a/easycom_head.py
+++ b/easycom_head.py
@@ -0,0 +1,299 @@
+"""Shared head, labels, and metrics for the EasyCom EgoSpeak experiment."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal, Sequence
+
+import numpy as np
+import torch
+from torch import nn
+
+
+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",
+)
+STEP_SECONDS = 0.2
+NUM_HORIZONS = 10
+NUM_CLASSES = 3
+
+
+class TurnPredictionHead(nn.Module):
+    """Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits."""
+
+    def __init__(
+        self,
+        hidden_size: int = 2048,
+        num_horizons: int = NUM_HORIZONS,
+        num_classes: int = NUM_CLASSES,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_horizons = num_horizons
+        self.num_classes = num_classes
+        self.dropout = dropout
+        self.net = nn.Sequential(
+            nn.LayerNorm(hidden_size),
+            nn.Linear(hidden_size, hidden_size // 2),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_size // 2, num_horizons * num_classes),
+        )
+
+    def forward(self, states: torch.Tensor) -> torch.Tensor:
+        logits = self.net(states)
+        return logits.view(
+            *states.shape[:-1], self.num_horizons, self.num_classes
+        )
+
+
+def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> 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_CLASSES:
+        raise ValueError(
+            f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}"
+        )
+
+
+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 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)}"
+
+
+@dataclass
+class SessionData:
+    name: str
+    states: torch.Tensor
+    targets: torch.Tensor
+    label_indices: torch.LongTensor
+    feature_metadata: dict[str, Any]
+
+    @property
+    def hidden_size(self) -> int:
+        return int(self.states.shape[-1])
+
+    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",
+) -> SessionData:
+    """Load and validate one saved Omni feature payload and target array."""
+
+    name = normalize_session_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 "states" not in payload:
+        raise ValueError(f"Invalid feature payload: {feature_path}")
+    states = payload["states"]
+    if not isinstance(states, torch.Tensor) or states.ndim != 2:
+        raise ValueError(f"states must be [T,D] in {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))
+    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",
+        )
+    }
+    return SessionData(name, states, targets, label_indices, metadata)
+
+
+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)
+    return config
+
+
+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)}")