619 lines
22 KiB
Python
619 lines
22 KiB
Python
|
"""Frozen EgoSpeak MiniROADA plus a trainable Omni residual branch.
|
||
|
|
|
||
|
|
This module deliberately does not modify the vendored EgoSpeak checkout. The
|
||
|
|
``MiniROADABase`` parameter names match the released implementation so its raw
|
||
|
|
state dict can be loaded strictly. Stage-one training freezes that complete
|
||
|
|
base and learns only a zero-initialized residual over the ten anticipation
|
||
|
|
logits.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
from collections.abc import Sequence
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Literal
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
from torch.nn import functional as F
|
||
|
|
from torch.utils.data import Dataset
|
||
|
|
|
||
|
|
from easycom_head import NUM_CLASSES, NUM_HORIZONS, STEP_SECONDS
|
||
|
|
|
||
|
|
OmniLayerFusion = Literal["single", "last", "mean"]
|
||
|
|
|
||
|
|
|
||
|
|
def load_torch_payload(path: Path, *, weights_only: bool = False) -> Any:
|
||
|
|
"""Load a torch payload while retaining compatibility with older PyTorch."""
|
||
|
|
|
||
|
|
try:
|
||
|
|
return torch.load(path, map_location="cpu", weights_only=weights_only)
|
||
|
|
except TypeError: # pragma: no cover - only exercised by older PyTorch.
|
||
|
|
return torch.load(path, map_location="cpu")
|
||
|
|
|
||
|
|
|
||
|
|
def sha256_file(path: Path) -> str:
|
||
|
|
digest = hashlib.sha256()
|
||
|
|
with path.open("rb") as handle:
|
||
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
||
|
|
digest.update(chunk)
|
||
|
|
return digest.hexdigest()
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_session_name(value: str | int) -> str:
|
||
|
|
text = str(value).strip()
|
||
|
|
suffix = text.removeprefix("Session_")
|
||
|
|
if not suffix.isdigit():
|
||
|
|
raise ValueError(f"Invalid EasyCom session: {value!r}")
|
||
|
|
return f"Session_{int(suffix)}"
|
||
|
|
|
||
|
|
|
||
|
|
class MiniROADABase(nn.Module):
|
||
|
|
"""Clean-room parameter-compatible implementation of EgoSpeak MiniROADA."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
*,
|
||
|
|
rgb_dim: int = 2048,
|
||
|
|
audio_dim: int = 5120,
|
||
|
|
embedding_dim: int = 2048,
|
||
|
|
hidden_dim: int = 1024,
|
||
|
|
num_layers: int = 1,
|
||
|
|
num_horizons: int = NUM_HORIZONS,
|
||
|
|
num_classes: int = NUM_CLASSES,
|
||
|
|
dropout: float = 0.2,
|
||
|
|
) -> None:
|
||
|
|
super().__init__()
|
||
|
|
self.rgb_dim = int(rgb_dim)
|
||
|
|
self.audio_dim = int(audio_dim)
|
||
|
|
self.embedding_dim = int(embedding_dim)
|
||
|
|
self.hidden_dim = int(hidden_dim)
|
||
|
|
self.num_layers = int(num_layers)
|
||
|
|
self.num_horizons = int(num_horizons)
|
||
|
|
self.num_classes = int(num_classes)
|
||
|
|
|
||
|
|
self.layer1 = nn.Sequential(
|
||
|
|
nn.Linear(self.rgb_dim + self.audio_dim, self.embedding_dim),
|
||
|
|
nn.LayerNorm(self.embedding_dim),
|
||
|
|
nn.ReLU(),
|
||
|
|
nn.Dropout(p=dropout),
|
||
|
|
)
|
||
|
|
self.relu = nn.ReLU()
|
||
|
|
self.gru = nn.GRU(
|
||
|
|
self.embedding_dim,
|
||
|
|
self.hidden_dim,
|
||
|
|
self.num_layers,
|
||
|
|
batch_first=True,
|
||
|
|
)
|
||
|
|
self.f_classification = nn.Sequential(
|
||
|
|
nn.Linear(self.hidden_dim, self.num_classes)
|
||
|
|
)
|
||
|
|
self.anticipation_layer = nn.Sequential(
|
||
|
|
nn.Linear(
|
||
|
|
self.hidden_dim,
|
||
|
|
self.num_horizons * self.hidden_dim,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
def encode(
|
||
|
|
self,
|
||
|
|
rgb: torch.Tensor,
|
||
|
|
audio: torch.Tensor,
|
||
|
|
recurrent_state: torch.Tensor | None = None,
|
||
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
|
|
if rgb.ndim != 3 or rgb.shape[-1] != self.rgb_dim:
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected RGB [B,S,{self.rgb_dim}], got {tuple(rgb.shape)}"
|
||
|
|
)
|
||
|
|
if (
|
||
|
|
audio.ndim != 3
|
||
|
|
or audio.shape[:-1] != rgb.shape[:-1]
|
||
|
|
or audio.shape[-1] != self.audio_dim
|
||
|
|
):
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected audio [B,S,{self.audio_dim}] aligned with RGB, got "
|
||
|
|
f"{tuple(audio.shape)}"
|
||
|
|
)
|
||
|
|
embedded = self.layer1(torch.cat((rgb, audio), dim=-1))
|
||
|
|
hidden, next_state = self.gru(embedded, recurrent_state)
|
||
|
|
return self.relu(hidden), next_state
|
||
|
|
|
||
|
|
def anticipation_logits(self, hidden: torch.Tensor) -> torch.Tensor:
|
||
|
|
shape = hidden.shape[:-1]
|
||
|
|
anticipated = self.anticipation_layer(self.relu(hidden)).view(
|
||
|
|
*shape,
|
||
|
|
self.num_horizons,
|
||
|
|
self.hidden_dim,
|
||
|
|
)
|
||
|
|
return self.f_classification(self.relu(anticipated))
|
||
|
|
|
||
|
|
def current_logits(self, hidden: torch.Tensor) -> torch.Tensor:
|
||
|
|
return self.f_classification(self.relu(hidden))
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self, rgb: torch.Tensor, audio: torch.Tensor
|
||
|
|
) -> dict[str, torch.Tensor]:
|
||
|
|
hidden, _ = self.encode(rgb, audio)
|
||
|
|
return {
|
||
|
|
"logits": self.current_logits(hidden),
|
||
|
|
"anticipation_logits": self.anticipation_logits(hidden),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _unwrap_state_dict(payload: Any) -> dict[str, torch.Tensor]:
|
||
|
|
if not isinstance(payload, dict):
|
||
|
|
raise TypeError("EgoSpeak checkpoint must contain a state dict")
|
||
|
|
for key in ("state_dict", "model_state_dict"):
|
||
|
|
nested = payload.get(key)
|
||
|
|
if isinstance(nested, dict):
|
||
|
|
payload = nested
|
||
|
|
break
|
||
|
|
if not payload or not all(
|
||
|
|
isinstance(value, torch.Tensor) for value in payload.values()
|
||
|
|
):
|
||
|
|
raise ValueError("EgoSpeak checkpoint does not contain tensor parameters")
|
||
|
|
state = dict(payload)
|
||
|
|
if all(key.startswith("module.") for key in state):
|
||
|
|
state = {key.removeprefix("module."): value for key, value in state.items()}
|
||
|
|
return state
|
||
|
|
|
||
|
|
|
||
|
|
def load_frozen_miniroada(checkpoint: Path, **model_kwargs: Any) -> MiniROADABase:
|
||
|
|
checkpoint = checkpoint.expanduser().resolve()
|
||
|
|
if not checkpoint.is_file():
|
||
|
|
raise FileNotFoundError(checkpoint)
|
||
|
|
base = MiniROADABase(**model_kwargs)
|
||
|
|
state = _unwrap_state_dict(load_torch_payload(checkpoint, weights_only=True))
|
||
|
|
base.load_state_dict(state, strict=True)
|
||
|
|
base.requires_grad_(False)
|
||
|
|
base.eval()
|
||
|
|
return base
|
||
|
|
|
||
|
|
|
||
|
|
class EgoSpeakOmniResidual(nn.Module):
|
||
|
|
"""Feature fusion whose zero residual exactly recovers frozen MiniROADA."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
base: MiniROADABase,
|
||
|
|
*,
|
||
|
|
omni_dim: int,
|
||
|
|
fusion_dim: int = 256,
|
||
|
|
tail_embedding_dim: int = 32,
|
||
|
|
num_tail_bins: int = NUM_HORIZONS,
|
||
|
|
dropout: float = 0.2,
|
||
|
|
) -> None:
|
||
|
|
super().__init__()
|
||
|
|
if fusion_dim <= 0 or tail_embedding_dim <= 0 or num_tail_bins <= 0:
|
||
|
|
raise ValueError("Fusion and tail dimensions must be positive")
|
||
|
|
self.base = base
|
||
|
|
self.omni_dim = int(omni_dim)
|
||
|
|
self.fusion_dim = int(fusion_dim)
|
||
|
|
self.tail_embedding_dim = int(tail_embedding_dim)
|
||
|
|
self.num_tail_bins = int(num_tail_bins)
|
||
|
|
|
||
|
|
self.ego_projection = nn.Sequential(
|
||
|
|
nn.LayerNorm(base.hidden_dim),
|
||
|
|
nn.Linear(base.hidden_dim, fusion_dim),
|
||
|
|
nn.GELU(),
|
||
|
|
)
|
||
|
|
self.omni_projection = nn.Sequential(
|
||
|
|
nn.LayerNorm(omni_dim),
|
||
|
|
nn.Linear(omni_dim, fusion_dim),
|
||
|
|
nn.GELU(),
|
||
|
|
)
|
||
|
|
self.tail_embedding = nn.Embedding(num_tail_bins, tail_embedding_dim)
|
||
|
|
self.fusion = nn.Sequential(
|
||
|
|
nn.LayerNorm(3 * fusion_dim + tail_embedding_dim),
|
||
|
|
nn.Linear(3 * fusion_dim + tail_embedding_dim, fusion_dim),
|
||
|
|
nn.GELU(),
|
||
|
|
nn.Dropout(dropout),
|
||
|
|
)
|
||
|
|
self.residual_output = nn.Linear(
|
||
|
|
fusion_dim,
|
||
|
|
base.num_horizons * base.num_classes,
|
||
|
|
)
|
||
|
|
nn.init.zeros_(self.residual_output.weight)
|
||
|
|
nn.init.zeros_(self.residual_output.bias)
|
||
|
|
|
||
|
|
self.base.requires_grad_(False)
|
||
|
|
self.base.eval()
|
||
|
|
|
||
|
|
def train(self, mode: bool = True) -> EgoSpeakOmniResidual:
|
||
|
|
super().train(mode)
|
||
|
|
# The released base remains deterministic: in particular, its dropout
|
||
|
|
# must never reactivate during stage-one residual training.
|
||
|
|
self.base.eval()
|
||
|
|
return self
|
||
|
|
|
||
|
|
def trainable_parameters(self) -> list[nn.Parameter]:
|
||
|
|
return [parameter for parameter in self.parameters() if parameter.requires_grad]
|
||
|
|
|
||
|
|
def _residual_logits(
|
||
|
|
self,
|
||
|
|
ego_hidden: torch.Tensor,
|
||
|
|
omni: torch.Tensor,
|
||
|
|
tail_bin: torch.Tensor,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
if omni.shape != (*ego_hidden.shape[:-1], self.omni_dim):
|
||
|
|
raise ValueError(
|
||
|
|
f"Omni/hidden shape mismatch: {tuple(omni.shape)} versus "
|
||
|
|
f"{tuple(ego_hidden.shape)}"
|
||
|
|
)
|
||
|
|
if tail_bin.shape != ego_hidden.shape[:-1]:
|
||
|
|
raise ValueError(
|
||
|
|
f"Tail-bin shape mismatch: {tuple(tail_bin.shape)} versus "
|
||
|
|
f"{tuple(ego_hidden.shape[:-1])}"
|
||
|
|
)
|
||
|
|
if tail_bin.dtype != torch.long:
|
||
|
|
raise ValueError("tail_bin must be int64")
|
||
|
|
if bool(((tail_bin < 0) | (tail_bin >= self.num_tail_bins)).any()):
|
||
|
|
raise ValueError("tail_bin lies outside the configured block phase")
|
||
|
|
|
||
|
|
ego = self.ego_projection(ego_hidden)
|
||
|
|
global_state = self.omni_projection(omni)
|
||
|
|
fused = torch.cat(
|
||
|
|
(
|
||
|
|
ego,
|
||
|
|
global_state,
|
||
|
|
ego * global_state,
|
||
|
|
self.tail_embedding(tail_bin),
|
||
|
|
),
|
||
|
|
dim=-1,
|
||
|
|
)
|
||
|
|
residual = self.residual_output(self.fusion(fused))
|
||
|
|
return residual.view(
|
||
|
|
*ego_hidden.shape[:-1],
|
||
|
|
self.base.num_horizons,
|
||
|
|
self.base.num_classes,
|
||
|
|
)
|
||
|
|
|
||
|
|
def forward_last(
|
||
|
|
self,
|
||
|
|
rgb: torch.Tensor,
|
||
|
|
audio: torch.Tensor,
|
||
|
|
omni: torch.Tensor,
|
||
|
|
tail_bin: torch.Tensor,
|
||
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
||
|
|
with torch.no_grad():
|
||
|
|
hidden, _ = self.base.encode(rgb, audio)
|
||
|
|
ego_hidden = hidden[:, -1, :]
|
||
|
|
base_logits = self.base.anticipation_logits(ego_hidden)
|
||
|
|
residual = self._residual_logits(ego_hidden.detach(), omni, tail_bin)
|
||
|
|
return base_logits + residual, base_logits
|
||
|
|
|
||
|
|
def forward_sequence_chunk(
|
||
|
|
self,
|
||
|
|
rgb: torch.Tensor,
|
||
|
|
audio: torch.Tensor,
|
||
|
|
omni: torch.Tensor,
|
||
|
|
tail_bin: torch.Tensor,
|
||
|
|
recurrent_state: torch.Tensor | None,
|
||
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||
|
|
with torch.no_grad():
|
||
|
|
hidden, next_state = self.base.encode(rgb, audio, recurrent_state)
|
||
|
|
base_logits = self.base.anticipation_logits(hidden)
|
||
|
|
residual = self._residual_logits(hidden.detach(), omni, tail_bin)
|
||
|
|
return base_logits + residual, base_logits, next_state
|
||
|
|
|
||
|
|
|
||
|
|
def collapse_omni_states(
|
||
|
|
states: torch.Tensor,
|
||
|
|
mode: OmniLayerFusion,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
if not isinstance(states, torch.Tensor) or states.ndim not in (2, 3):
|
||
|
|
raise ValueError("Omni states must be [T,D] or [T,L,D]")
|
||
|
|
if not torch.is_floating_point(states) or not torch.isfinite(states).all():
|
||
|
|
raise ValueError("Omni states must be finite floating-point values")
|
||
|
|
if states.ndim == 2:
|
||
|
|
if mode not in ("single", "last", "mean"):
|
||
|
|
raise ValueError(f"Unsupported Omni layer fusion: {mode}")
|
||
|
|
return states.float().contiguous()
|
||
|
|
if mode == "single":
|
||
|
|
raise ValueError("Stacked Omni states require --omni-layer-fusion last or mean")
|
||
|
|
if mode == "last":
|
||
|
|
return states[:, -1, :].float().contiguous()
|
||
|
|
if mode == "mean":
|
||
|
|
return states.float().mean(dim=1).contiguous()
|
||
|
|
raise ValueError(f"Unsupported Omni layer fusion: {mode}")
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class EasyComSessionArrays:
|
||
|
|
name: str
|
||
|
|
rgb: np.ndarray
|
||
|
|
audio: np.ndarray
|
||
|
|
targets: np.ndarray
|
||
|
|
omni: torch.Tensor
|
||
|
|
tail_bins: torch.LongTensor
|
||
|
|
omni_metadata: dict[str, Any]
|
||
|
|
|
||
|
|
@property
|
||
|
|
def num_steps(self) -> int:
|
||
|
|
return len(self.targets)
|
||
|
|
|
||
|
|
|
||
|
|
def _require_feature_array(path: Path, expected_dim: int) -> np.ndarray:
|
||
|
|
if not path.is_file():
|
||
|
|
raise FileNotFoundError(path)
|
||
|
|
array = np.load(path, mmap_mode="r", allow_pickle=False)
|
||
|
|
if array.ndim != 2 or array.shape[1] != expected_dim:
|
||
|
|
raise ValueError(f"Expected [T,{expected_dim}], got {array.shape}: {path}")
|
||
|
|
if not np.issubdtype(array.dtype, np.floating):
|
||
|
|
raise ValueError(f"Expected floating-point features: {path}")
|
||
|
|
return array
|
||
|
|
|
||
|
|
|
||
|
|
def _require_target_array(path: Path, expected_dim: int) -> np.ndarray:
|
||
|
|
if not path.is_file():
|
||
|
|
raise FileNotFoundError(path)
|
||
|
|
array = np.load(path, mmap_mode="r", allow_pickle=False)
|
||
|
|
if array.ndim != 2 or array.shape[1] != expected_dim:
|
||
|
|
raise ValueError(f"Expected [T,{expected_dim}], got {array.shape}: {path}")
|
||
|
|
if not (
|
||
|
|
np.issubdtype(array.dtype, np.number) or np.issubdtype(array.dtype, np.bool_)
|
||
|
|
):
|
||
|
|
raise ValueError(f"Expected numeric activity targets: {path}")
|
||
|
|
return array
|
||
|
|
|
||
|
|
|
||
|
|
def load_easycom_session_arrays(
|
||
|
|
egospeak_feature_root: Path,
|
||
|
|
omni_feature_root: Path,
|
||
|
|
session: str | int,
|
||
|
|
*,
|
||
|
|
expected_omni_mode: str | None = "open-stream-replay",
|
||
|
|
omni_layer_fusion: OmniLayerFusion = "single",
|
||
|
|
rgb_dim: int = 2048,
|
||
|
|
audio_dim: int = 5120,
|
||
|
|
num_classes: int = NUM_CLASSES,
|
||
|
|
) -> EasyComSessionArrays:
|
||
|
|
name = normalize_session_name(session)
|
||
|
|
ego_root = egospeak_feature_root.expanduser().resolve()
|
||
|
|
omni_root = omni_feature_root.expanduser().resolve()
|
||
|
|
rgb = _require_feature_array(
|
||
|
|
ego_root / "rgb_kinetics_resnet50" / f"{name}.npy",
|
||
|
|
rgb_dim,
|
||
|
|
)
|
||
|
|
audio = _require_feature_array(
|
||
|
|
ego_root / "audio_wav2vec2" / f"{name}.npy",
|
||
|
|
audio_dim,
|
||
|
|
)
|
||
|
|
targets = _require_target_array(
|
||
|
|
ego_root / "target_perframe" / f"{name}.npy",
|
||
|
|
num_classes,
|
||
|
|
)
|
||
|
|
lengths = {len(rgb), len(audio), len(targets)}
|
||
|
|
if len(lengths) != 1:
|
||
|
|
raise ValueError(
|
||
|
|
f"{name}: unequal EgoSpeak RGB/audio/target lengths "
|
||
|
|
f"{len(rgb)}/{len(audio)}/{len(targets)}"
|
||
|
|
)
|
||
|
|
unique_targets = np.unique(np.asarray(targets))
|
||
|
|
if not np.isin(unique_targets, (0, 1)).all():
|
||
|
|
raise ValueError(
|
||
|
|
f"{name}: target_perframe contains values outside binary activity: "
|
||
|
|
f"{unique_targets.tolist()}"
|
||
|
|
)
|
||
|
|
|
||
|
|
omni_path = omni_root / f"{name}.pt"
|
||
|
|
if not omni_path.is_file():
|
||
|
|
raise FileNotFoundError(omni_path)
|
||
|
|
payload = load_torch_payload(omni_path)
|
||
|
|
if not isinstance(payload, dict):
|
||
|
|
raise TypeError(f"Invalid Omni payload: {omni_path}")
|
||
|
|
mode = payload.get("mode")
|
||
|
|
if expected_omni_mode is not None and mode != expected_omni_mode:
|
||
|
|
raise ValueError(
|
||
|
|
f"{name}: Omni mode is {mode!r}, expected {expected_omni_mode!r}"
|
||
|
|
)
|
||
|
|
omni = collapse_omni_states(payload.get("states"), omni_layer_fusion)
|
||
|
|
if len(omni) != len(targets):
|
||
|
|
raise ValueError(
|
||
|
|
f"{name}: Omni/target length mismatch {len(omni)} vs {len(targets)}"
|
||
|
|
)
|
||
|
|
if payload.get("label_indices") is None:
|
||
|
|
raise ValueError(f"{name}: Omni payload lacks label_indices")
|
||
|
|
label_indices = torch.as_tensor(payload["label_indices"], dtype=torch.long)
|
||
|
|
expected_indices = torch.arange(len(targets), dtype=torch.long)
|
||
|
|
if label_indices.ndim != 1 or not torch.equal(label_indices, expected_indices):
|
||
|
|
raise ValueError(f"{name}: Omni label_indices are not dense zero-origin rows")
|
||
|
|
|
||
|
|
config = payload.get("config")
|
||
|
|
if not isinstance(config, dict):
|
||
|
|
raise TypeError(f"{name}: Omni payload lacks extraction config")
|
||
|
|
step_seconds = float(config.get("step_seconds", float("nan")))
|
||
|
|
native_block_seconds = float(config.get("native_block_seconds", float("nan")))
|
||
|
|
if not np.isclose(step_seconds, STEP_SECONDS, rtol=0.0, atol=1e-9):
|
||
|
|
raise ValueError(
|
||
|
|
f"{name}: Omni step_seconds={step_seconds}, expected {STEP_SECONDS}"
|
||
|
|
)
|
||
|
|
block_steps_float = native_block_seconds / step_seconds
|
||
|
|
block_steps = round(block_steps_float)
|
||
|
|
if block_steps <= 0 or not np.isclose(
|
||
|
|
block_steps_float, block_steps, rtol=0.0, atol=1e-9
|
||
|
|
):
|
||
|
|
raise ValueError(
|
||
|
|
f"{name}: native block {native_block_seconds}s is not an integer "
|
||
|
|
f"multiple of {step_seconds}s"
|
||
|
|
)
|
||
|
|
if block_steps != NUM_HORIZONS:
|
||
|
|
raise ValueError(
|
||
|
|
f"{name}: expected ten 0.2-s tail bins in a 2-s block, got {block_steps}"
|
||
|
|
)
|
||
|
|
|
||
|
|
times = torch.as_tensor(payload.get("times"), dtype=torch.float64)
|
||
|
|
start = float(payload.get("start", 0.0))
|
||
|
|
expected_times = (
|
||
|
|
start + (torch.arange(len(targets), dtype=torch.float64) + 1) * step_seconds
|
||
|
|
)
|
||
|
|
if times.ndim != 1 or not torch.allclose(
|
||
|
|
times, expected_times, rtol=0.0, atol=1e-6
|
||
|
|
):
|
||
|
|
raise ValueError(f"{name}: Omni times do not match the dense 5 Hz grid")
|
||
|
|
if bool(payload.get("partial", False)):
|
||
|
|
raise ValueError(f"{name}: partial Omni smoke feature cannot train the model")
|
||
|
|
|
||
|
|
tail_bins = torch.remainder(expected_indices, block_steps)
|
||
|
|
metadata = {
|
||
|
|
"path": str(omni_path),
|
||
|
|
"mode": mode,
|
||
|
|
"model_id": payload.get("model_id"),
|
||
|
|
"config": config,
|
||
|
|
"video_sampling_grid": payload.get("video_sampling_grid"),
|
||
|
|
"media_interval_policy": payload.get("media_interval_policy"),
|
||
|
|
"open_stream_policy": payload.get("open_stream_policy"),
|
||
|
|
"state_shape_before_layer_fusion": list(payload["states"].shape),
|
||
|
|
"state_shape": list(omni.shape),
|
||
|
|
}
|
||
|
|
return EasyComSessionArrays(
|
||
|
|
name=name,
|
||
|
|
rgb=rgb,
|
||
|
|
audio=audio,
|
||
|
|
targets=targets,
|
||
|
|
omni=omni,
|
||
|
|
tail_bins=tail_bins,
|
||
|
|
omni_metadata=metadata,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class EgoSpeakOmniWindowDataset(Dataset[tuple[torch.Tensor, ...]]):
|
||
|
|
"""MiniROADA windows plus the aligned current Omni state."""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
sessions: Sequence[EasyComSessionArrays],
|
||
|
|
*,
|
||
|
|
window_size: int = 128,
|
||
|
|
stride: int = 4,
|
||
|
|
num_horizons: int = NUM_HORIZONS,
|
||
|
|
seed: int = 0,
|
||
|
|
max_examples: int | None = None,
|
||
|
|
) -> None:
|
||
|
|
if not sessions:
|
||
|
|
raise ValueError("At least one training session is required")
|
||
|
|
if window_size <= 0 or stride <= 0 or num_horizons <= 0:
|
||
|
|
raise ValueError("Window, stride, and horizon counts must be positive")
|
||
|
|
if max_examples is not None and max_examples <= 0:
|
||
|
|
raise ValueError("max_examples must be positive")
|
||
|
|
self.sessions = list(sessions)
|
||
|
|
self.window_size = int(window_size)
|
||
|
|
self.stride = int(stride)
|
||
|
|
self.num_horizons = int(num_horizons)
|
||
|
|
self.seed = int(seed)
|
||
|
|
self.max_examples = max_examples
|
||
|
|
self.indices: list[tuple[int, int]] = []
|
||
|
|
self.resample(0)
|
||
|
|
|
||
|
|
def resample(self, epoch: int) -> None:
|
||
|
|
rng = np.random.default_rng(self.seed + int(epoch))
|
||
|
|
indices: list[tuple[int, int]] = []
|
||
|
|
for session_index, session in enumerate(self.sessions):
|
||
|
|
last_source = session.num_steps - self.num_horizons - 1
|
||
|
|
if last_source < 0:
|
||
|
|
raise ValueError(
|
||
|
|
f"{session.name}: only {session.num_steps} rows for "
|
||
|
|
f"{self.num_horizons} horizons"
|
||
|
|
)
|
||
|
|
offset = int(rng.integers(self.stride))
|
||
|
|
indices.extend(
|
||
|
|
(session_index, source_index)
|
||
|
|
for source_index in range(offset, last_source + 1, self.stride)
|
||
|
|
)
|
||
|
|
if self.max_examples is not None:
|
||
|
|
rng.shuffle(indices)
|
||
|
|
indices = indices[: self.max_examples]
|
||
|
|
if not indices:
|
||
|
|
raise ValueError("Training resampling produced no examples")
|
||
|
|
self.indices = indices
|
||
|
|
|
||
|
|
def __len__(self) -> int:
|
||
|
|
return len(self.indices)
|
||
|
|
|
||
|
|
def __getitem__(self, index: int) -> tuple[torch.Tensor, ...]:
|
||
|
|
session_index, source_index = self.indices[index]
|
||
|
|
session = self.sessions[session_index]
|
||
|
|
first = max(0, source_index - self.window_size + 1)
|
||
|
|
length = source_index - first + 1
|
||
|
|
rgb = torch.zeros(
|
||
|
|
self.window_size,
|
||
|
|
session.rgb.shape[1],
|
||
|
|
dtype=torch.float32,
|
||
|
|
)
|
||
|
|
audio = torch.zeros(
|
||
|
|
self.window_size,
|
||
|
|
session.audio.shape[1],
|
||
|
|
dtype=torch.float32,
|
||
|
|
)
|
||
|
|
rgb[-length:] = torch.from_numpy(
|
||
|
|
np.asarray(session.rgb[first : source_index + 1], dtype=np.float32).copy()
|
||
|
|
)
|
||
|
|
audio[-length:] = torch.from_numpy(
|
||
|
|
np.asarray(session.audio[first : source_index + 1], dtype=np.float32).copy()
|
||
|
|
)
|
||
|
|
future = torch.from_numpy(
|
||
|
|
np.asarray(
|
||
|
|
session.targets[
|
||
|
|
source_index + 1 : source_index + self.num_horizons + 1
|
||
|
|
],
|
||
|
|
dtype=np.float32,
|
||
|
|
).copy()
|
||
|
|
)
|
||
|
|
return (
|
||
|
|
rgb,
|
||
|
|
audio,
|
||
|
|
session.omni[source_index].float(),
|
||
|
|
session.tail_bins[source_index],
|
||
|
|
future,
|
||
|
|
torch.tensor(session_index, dtype=torch.long),
|
||
|
|
torch.tensor(source_index, dtype=torch.long),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def miniroad_normalized_loss(
|
||
|
|
logits: torch.Tensor,
|
||
|
|
targets: torch.Tensor,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Released MiniROAD multi-positive anticipation objective."""
|
||
|
|
|
||
|
|
if logits.shape != targets.shape or logits.shape[-1] != NUM_CLASSES:
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected equal [...,{NUM_CLASSES}] logits/targets, got "
|
||
|
|
f"{tuple(logits.shape)} and {tuple(targets.shape)}"
|
||
|
|
)
|
||
|
|
normalized_targets = F.normalize(targets, dim=-1)
|
||
|
|
return torch.sum(
|
||
|
|
-normalized_targets * F.log_softmax(logits, dim=-1),
|
||
|
|
dim=-1,
|
||
|
|
).mean()
|
||
|
|
|
||
|
|
|
||
|
|
def align_source_predictions(
|
||
|
|
raw_prediction: np.ndarray,
|
||
|
|
target_steps: int,
|
||
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
||
|
|
"""Align MiniROADA [source,horizon,class] forecasts to target rows."""
|
||
|
|
|
||
|
|
raw_prediction = np.asarray(raw_prediction, dtype=np.float32)
|
||
|
|
expected = (target_steps - NUM_HORIZONS, NUM_HORIZONS, NUM_CLASSES)
|
||
|
|
if raw_prediction.shape != expected:
|
||
|
|
raise ValueError(
|
||
|
|
f"Expected raw prediction {expected}, got {raw_prediction.shape}"
|
||
|
|
)
|
||
|
|
if not np.isfinite(raw_prediction).all():
|
||
|
|
raise ValueError("Raw prediction contains NaN/Inf")
|
||
|
|
predictions = np.zeros(
|
||
|
|
(target_steps, NUM_CLASSES, NUM_HORIZONS),
|
||
|
|
dtype=np.float32,
|
||
|
|
)
|
||
|
|
valid = np.zeros((target_steps, NUM_HORIZONS), dtype=bool)
|
||
|
|
source_rows = np.arange(len(raw_prediction), dtype=np.int64)
|
||
|
|
for horizon in range(NUM_HORIZONS):
|
||
|
|
target_rows = source_rows + horizon + 1
|
||
|
|
predictions[target_rows, :, horizon] = raw_prediction[:, horizon, :]
|
||
|
|
valid[target_rows, horizon] = True
|
||
|
|
return predictions, valid
|