Add frozen EgoSpeak Omni residual fusion stage
This commit is contained in:
Xin Wang
@@ -0,0 +1,618 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user