Add frozen EgoSpeak Omni residual fusion stage
This commit is contained in:
Xin Wang
@@ -135,6 +135,15 @@ result.
|
||||
described as upstream Stage-3 code because that entry point remains
|
||||
unpublished. See the
|
||||
[`runbook`](docs/easycom_muvap_stage3_baseline.md).
|
||||
- `egospeak-omni-zero-residual` is implemented but unmeasured. It strictly
|
||||
loads and freezes the released MiniROADA A+V/GRU/anticipation path, reuses
|
||||
existing causal 5 Hz Omni readouts, adds an explicit ten-bin partial-tail
|
||||
phase embedding, and learns only a zero-output residual over the original
|
||||
logits. The launcher verifies exact initial recovery and exports fused and
|
||||
matched frozen-EgoSpeak predictions together. Existing Omni states already
|
||||
combine completed history with the current partial tail; separating those
|
||||
states remains a future extraction ablation. See the
|
||||
[`runbook`](docs/easycom_egospeak_omni_residual.md).
|
||||
|
||||
### Measured oracle-information screen
|
||||
|
||||
|
||||
@@ -28,6 +28,7 @@ tests.
|
||||
| [`easycom_track_gru_residual.md`](easycom_track_gru_residual.md) | Matched pooled-to-track GRU experiment, causal identity contract and L20 commands |
|
||||
| [`easycom_causal_person_adapter_v2.md`](easycom_causal_person_adapter_v2.md) | Dynamic-gated, zero-output causal person-track residual and matched L20 screen |
|
||||
| [`easycom_egospeak_style_person_gru.md`](easycom_egospeak_style_person_gru.md) | Controlled V2 temporal-stem ablation using EgoSpeak MiniROADA's layer order |
|
||||
| [`easycom_egospeak_omni_residual.md`](easycom_egospeak_omni_residual.md) | Stage-one frozen released MiniROADA plus zero-initialized tail-aware Omni residual fusion |
|
||||
| [`easycom_highrate_person_av.md`](easycom_highrate_person_av.md) | Frozen 25 Hz causal ASD person evidence fused with 5 Hz global Omni context |
|
||||
| [`easycom_muvap_stage3_baseline.md`](easycom_muvap_stage3_baseline.md) | MuVAP-only frozen VAP+ASD Stage-3 paper reconstruction on EgoSpeak/EasyCom targets |
|
||||
| [`easycom_oracle_information.md`](easycom_oracle_information.md) | Causal oracle per-person speaking-history upper bound and Session-12 screening commands |
|
||||
|
||||
170
docs/easycom_egospeak_omni_residual.md
Normal file
170
docs/easycom_egospeak_omni_residual.md
Normal file
@@ -0,0 +1,170 @@
|
||||
# Frozen EgoSpeak + Omni tail-aware residual fusion
|
||||
|
||||
## Scope
|
||||
|
||||
This is the first-stage dual-stream screen. It preserves the released
|
||||
EgoSpeak MiniROADA RGB, wav2vec2-audio, GRU, and ten-horizon anticipation path
|
||||
as a frozen base. Existing causal 5 Hz Qwen2.5-Omni states provide a second
|
||||
branch. Only a tail-phase-aware residual fusion module is trained:
|
||||
|
||||
```text
|
||||
released EgoSpeak RGB + audio -> frozen MiniROADA -> h_ego -> z_ego ---------+
|
||||
existing causal Omni g_full -> projection -----------------------> fusion -> Δz
|
||||
|
|
||||
z_fused = z_ego + Δz
|
||||
```
|
||||
|
||||
The final residual projection is initialized to exactly zero. Before any
|
||||
optimization, the launcher verifies bit-exact equality between fused and base
|
||||
logits on a real training example. Unlike the earlier scalar-zero-gate track
|
||||
screen, the zero is placed only on the final residual output layer; after its
|
||||
first update, gradients can reach the fusion branch.
|
||||
|
||||
This is **baseline-recoverable initialization**, not a guarantee that the
|
||||
trained fusion will be non-inferior on unseen Sessions 1--3. The launcher
|
||||
therefore exports matched frozen-base and fused predictions from the same
|
||||
forward pass.
|
||||
|
||||
## Inputs
|
||||
|
||||
No Qwen re-extraction is needed for this stage. The Omni root must contain
|
||||
complete `Session_X.pt` files with:
|
||||
|
||||
- `mode=open-stream-replay` by default;
|
||||
- dense zero-origin `label_indices`;
|
||||
- one state per 5 Hz target row;
|
||||
- `step_seconds=0.2` and `native_block_seconds=2.0`;
|
||||
- a full-session `times` grid and no `partial=True` marker.
|
||||
|
||||
The saved state is already the readout over rolling completed blocks plus the
|
||||
current temporary partial block. The fusion branch derives the ten-value tail
|
||||
phase (`0.2, ..., 2.0` seconds) from the validated grid. It does not claim to
|
||||
separate a completed-history state from a partial-tail state. Such an
|
||||
ablation would require a new Omni extraction field.
|
||||
|
||||
The EgoSpeak feature root must contain the released directories:
|
||||
|
||||
```text
|
||||
rgb_kinetics_resnet50/Session_X.npy [T,2048]
|
||||
audio_wav2vec2/Session_X.npy [T,5120]
|
||||
target_perframe/Session_X.npy [T,3]
|
||||
```
|
||||
|
||||
The base checkpoint must be the released raw MiniROADA state dict. It is
|
||||
loaded strictly, and every base parameter remains frozen and in evaluation
|
||||
mode so its dropout cannot reactivate during residual training.
|
||||
|
||||
## Audit the inputs
|
||||
|
||||
Set machine-specific paths rather than committing `/data1` paths:
|
||||
|
||||
```bash
|
||||
EGO_ROOT=third_party/EgoSpeak/MiniROAD/data/EasyCom
|
||||
OMNI_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack
|
||||
BASE_CKPT=third_party/EgoSpeak/MiniROAD/checkpoints/miniroad_easycom_AV.pth
|
||||
```
|
||||
|
||||
The default `--omni-layer-fusion mean` uses a fixed mean for the existing
|
||||
last-four stack and is also valid for a single `[T,D]` state. Use `last` only
|
||||
for an explicit layer ablation.
|
||||
|
||||
Audit the released EgoSpeak arrays:
|
||||
|
||||
```bash
|
||||
python3 audit_egospeak_features.py \
|
||||
--feature-root "$EGO_ROOT"
|
||||
```
|
||||
|
||||
The training launcher independently validates every selected Omni and
|
||||
EgoSpeak session and aborts on a length, time-grid, mode, dimension, or
|
||||
zero-origin mismatch.
|
||||
|
||||
## Session 4 smoke test
|
||||
|
||||
Use a new output directory for every invocation; the launcher refuses to
|
||||
overwrite an existing nonempty result.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_residual.py \
|
||||
--egospeak-feature-root "$EGO_ROOT" \
|
||||
--omni-feature-root "$OMNI_ROOT" \
|
||||
--base-checkpoint "$BASE_CKPT" \
|
||||
--train-sessions Session_4 \
|
||||
--eval-sessions Session_4 \
|
||||
--epochs 1 \
|
||||
--max-train-examples 32 \
|
||||
--batch-size 8 \
|
||||
--num-workers 0 \
|
||||
--output-dir results/easycom_egospeak_omni_residual_smoke
|
||||
```
|
||||
|
||||
The required early output is:
|
||||
|
||||
```text
|
||||
Initialization audit: fused logits == frozen EgoSpeak logits exactly
|
||||
```
|
||||
|
||||
If a raw released Session 4 prediction is available, add:
|
||||
|
||||
```bash
|
||||
--reference-prediction-root \
|
||||
third_party/EgoSpeak/MiniROAD/output_cfg/miniroad_aa_easycom_session4_AV/eval
|
||||
```
|
||||
|
||||
This is a stronger implementation audit: the launcher checks that its frozen
|
||||
base reproduces the upstream raw `[source,horizon,class]` probabilities.
|
||||
|
||||
## Fixed four-epoch Sessions 4--12 run
|
||||
|
||||
After the smoke test passes, omit session lists to use the official EasyCom
|
||||
train/test split. Evaluation happens once after epoch 4; the test split is not
|
||||
used for checkpoint selection.
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_residual.py \
|
||||
--egospeak-feature-root "$EGO_ROOT" \
|
||||
--omni-feature-root "$OMNI_ROOT" \
|
||||
--base-checkpoint "$BASE_CKPT" \
|
||||
--epochs 4 \
|
||||
--batch-size 64 \
|
||||
--num-workers 4 \
|
||||
--seed 0 \
|
||||
--output-dir results/easycom_egospeak_omni_residual_seed0
|
||||
```
|
||||
|
||||
Do not add `--amp` to the first reproducibility run. Once the float32 base has
|
||||
been compared against released predictions, AMP can be screened separately.
|
||||
The default `--eval-chunk-size 0` processes each evaluation session in one GRU
|
||||
pass, matching the released evaluator. Set a positive value only if inference
|
||||
memory requires chunked recurrent execution, and repeat the reference audit.
|
||||
|
||||
The output contains:
|
||||
|
||||
```text
|
||||
checkpoint_last.pt self-contained model/checkpoint provenance
|
||||
metrics.json fused dense metrics plus matched frozen base
|
||||
raw_eval/ fused MiniROADA source-indexed probabilities
|
||||
pred/ and valid/ fused evaluator-ready aligned arrays
|
||||
frozen_egospeak_raw_eval/ matched frozen-base raw probabilities
|
||||
frozen_egospeak_pred/ matched frozen-base aligned arrays
|
||||
```
|
||||
|
||||
`metrics.json` keeps the fused `causal_valid_only` protocol at the top level,
|
||||
so it remains readable by `summarize_local_baselines.py`.
|
||||
|
||||
## Interpretation and next gate
|
||||
|
||||
This first stage uses only the released dense future-activity objective. It
|
||||
tests whether Omni adds useful context without discarding the task-specific
|
||||
EgoSpeak path. It does not train semantic ENTRY/YIELD heads, and Sessions
|
||||
1--3 LLM-assisted reviews must never be used for optimization or model
|
||||
selection.
|
||||
|
||||
Advance to selective MiniROADA unfreezing or semantic opportunity heads only
|
||||
if the frozen residual passes these audits:
|
||||
|
||||
1. exact zero-residual initialization and upstream base reproduction;
|
||||
2. no causal/time-grid mismatch;
|
||||
3. stable dense and boundary changes over seeds;
|
||||
4. opportunity-conditioned ENTRY/YIELD reported separately, without treating
|
||||
dense wearer probability as semantic intent.
|
||||
618
egospeak_omni_residual.py
Normal file
618
egospeak_omni_residual.py
Normal file
@@ -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
|
||||
169
test_egospeak_omni_residual.py
Normal file
169
test_egospeak_omni_residual.py
Normal file
@@ -0,0 +1,169 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from egospeak_omni_residual import (
|
||||
EgoSpeakOmniResidual,
|
||||
EgoSpeakOmniWindowDataset,
|
||||
MiniROADABase,
|
||||
align_source_predictions,
|
||||
load_easycom_session_arrays,
|
||||
load_frozen_miniroada,
|
||||
miniroad_normalized_loss,
|
||||
)
|
||||
|
||||
|
||||
class EgoSpeakOmniResidualTest(unittest.TestCase):
|
||||
def test_raw_miniroada_state_dict_loads_strictly_and_freezes(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
checkpoint = Path(directory) / "base.pth"
|
||||
source = MiniROADABase(
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
embedding_dim=8,
|
||||
hidden_dim=5,
|
||||
num_horizons=10,
|
||||
)
|
||||
torch.save(source.state_dict(), checkpoint)
|
||||
loaded = load_frozen_miniroada(
|
||||
checkpoint,
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
embedding_dim=8,
|
||||
hidden_dim=5,
|
||||
num_horizons=10,
|
||||
)
|
||||
self.assertFalse(loaded.training)
|
||||
self.assertTrue(
|
||||
all(not parameter.requires_grad for parameter in loaded.parameters())
|
||||
)
|
||||
for key, value in source.state_dict().items():
|
||||
torch.testing.assert_close(loaded.state_dict()[key], value)
|
||||
|
||||
def test_zero_residual_exactly_recovers_frozen_base(self) -> None:
|
||||
torch.manual_seed(7)
|
||||
base = MiniROADABase(
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
embedding_dim=8,
|
||||
hidden_dim=5,
|
||||
num_horizons=10,
|
||||
dropout=0.4,
|
||||
)
|
||||
model = EgoSpeakOmniResidual(
|
||||
base,
|
||||
omni_dim=7,
|
||||
fusion_dim=4,
|
||||
tail_embedding_dim=3,
|
||||
dropout=0.3,
|
||||
)
|
||||
model.train()
|
||||
self.assertFalse(model.base.training)
|
||||
self.assertTrue(
|
||||
all(not parameter.requires_grad for parameter in model.base.parameters())
|
||||
)
|
||||
|
||||
rgb = torch.randn(2, 9, 4)
|
||||
audio = torch.randn(2, 9, 6)
|
||||
omni = torch.randn(2, 7)
|
||||
tail = torch.tensor([0, 9], dtype=torch.long)
|
||||
fused, frozen = model.forward_last(rgb, audio, omni, tail)
|
||||
torch.testing.assert_close(fused, frozen, rtol=0.0, atol=0.0)
|
||||
|
||||
targets = torch.zeros_like(fused)
|
||||
targets[..., 1] = 1
|
||||
optimizer = torch.optim.SGD(model.trainable_parameters(), lr=0.1)
|
||||
loss = miniroad_normalized_loss(fused, targets)
|
||||
loss.backward()
|
||||
self.assertGreater(float(model.residual_output.weight.grad.abs().sum()), 0.0)
|
||||
optimizer.step()
|
||||
updated, frozen_again = model.forward_last(rgb, audio, omni, tail)
|
||||
self.assertGreater(
|
||||
float((updated - frozen_again).abs().max().detach()),
|
||||
0.0,
|
||||
)
|
||||
|
||||
def test_feature_contract_and_window_alignment(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
root = Path(directory)
|
||||
ego = root / "ego"
|
||||
omni = root / "omni"
|
||||
for subdir in (
|
||||
"rgb_kinetics_resnet50",
|
||||
"audio_wav2vec2",
|
||||
"target_perframe",
|
||||
):
|
||||
(ego / subdir).mkdir(parents=True, exist_ok=True)
|
||||
omni.mkdir()
|
||||
steps = 15
|
||||
rgb = np.arange(steps * 4, dtype=np.float32).reshape(steps, 4)
|
||||
audio = np.arange(steps * 6, dtype=np.float32).reshape(steps, 6)
|
||||
target = np.zeros((steps, 3), dtype=np.float32)
|
||||
target[:, 0] = 1
|
||||
np.save(ego / "rgb_kinetics_resnet50" / "Session_4.npy", rgb)
|
||||
np.save(ego / "audio_wav2vec2" / "Session_4.npy", audio)
|
||||
np.save(ego / "target_perframe" / "Session_4.npy", target)
|
||||
torch.save(
|
||||
{
|
||||
"mode": "open-stream-replay",
|
||||
"states": torch.arange(steps * 7, dtype=torch.float32).view(
|
||||
steps, 7
|
||||
),
|
||||
"times": (torch.arange(steps, dtype=torch.float64) + 1) / 5,
|
||||
"label_indices": torch.arange(steps, dtype=torch.long),
|
||||
"start": 0.0,
|
||||
"config": {
|
||||
"step_seconds": 0.2,
|
||||
"native_block_seconds": 2.0,
|
||||
},
|
||||
},
|
||||
omni / "Session_4.pt",
|
||||
)
|
||||
session = load_easycom_session_arrays(
|
||||
ego,
|
||||
omni,
|
||||
"Session_4",
|
||||
rgb_dim=4,
|
||||
audio_dim=6,
|
||||
)
|
||||
self.assertEqual(
|
||||
session.tail_bins.tolist(), list(range(10)) + list(range(5))
|
||||
)
|
||||
dataset = EgoSpeakOmniWindowDataset(
|
||||
[session],
|
||||
window_size=4,
|
||||
stride=1,
|
||||
num_horizons=3,
|
||||
seed=0,
|
||||
)
|
||||
first_index = dataset.indices.index((0, 0))
|
||||
rgb_window, audio_window, state, tail, future, _, source = dataset[
|
||||
first_index
|
||||
]
|
||||
self.assertEqual(int(source), 0)
|
||||
self.assertEqual(int(tail), 0)
|
||||
torch.testing.assert_close(rgb_window[-1], torch.from_numpy(rgb[0]))
|
||||
self.assertEqual(int(torch.count_nonzero(rgb_window[:-1])), 0)
|
||||
torch.testing.assert_close(audio_window[-1], torch.from_numpy(audio[0]))
|
||||
torch.testing.assert_close(state, session.omni[0])
|
||||
torch.testing.assert_close(future, torch.from_numpy(target[1:4]))
|
||||
|
||||
def test_source_prediction_alignment(self) -> None:
|
||||
target_steps = 14
|
||||
raw = np.arange(
|
||||
(target_steps - 10) * 10 * 3,
|
||||
dtype=np.float32,
|
||||
).reshape(target_steps - 10, 10, 3)
|
||||
pred, valid = align_source_predictions(raw, target_steps)
|
||||
np.testing.assert_array_equal(pred[1, :, 0], raw[0, 0, :])
|
||||
np.testing.assert_array_equal(pred[13, :, 9], raw[3, 9, :])
|
||||
self.assertTrue(valid[1, 0])
|
||||
self.assertFalse(valid[0, 0])
|
||||
self.assertTrue(valid[13, 9])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
631
train_egospeak_omni_residual.py
Executable file
631
train_egospeak_omni_residual.py
Executable file
@@ -0,0 +1,631 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Train stage-one frozen-EgoSpeak plus frozen-Omni residual fusion.
|
||||
|
||||
The released MiniROADA A+V checkpoint is loaded strictly and remains in eval
|
||||
mode with gradients disabled. Existing causal 5 Hz Omni states are fused with
|
||||
the MiniROADA GRU state, and only a zero-initialized residual anticipation head
|
||||
is optimized. Evaluation happens once after the requested fixed epoch budget;
|
||||
the script saves both fused and exactly matched frozen-base predictions.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
from collections.abc import Sequence
|
||||
from contextlib import nullcontext
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from easycom_head import (
|
||||
CLASS_NAMES,
|
||||
NUM_HORIZONS,
|
||||
OFFICIAL_TEST_SESSIONS,
|
||||
OFFICIAL_TRAIN_SESSIONS,
|
||||
STEP_SECONDS,
|
||||
average_precision_metrics,
|
||||
)
|
||||
from egospeak_omni_residual import (
|
||||
EasyComSessionArrays,
|
||||
EgoSpeakOmniResidual,
|
||||
EgoSpeakOmniWindowDataset,
|
||||
align_source_predictions,
|
||||
load_easycom_session_arrays,
|
||||
load_frozen_miniroada,
|
||||
miniroad_normalized_loss,
|
||||
normalize_session_name,
|
||||
sha256_file,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--egospeak-feature-root",
|
||||
type=Path,
|
||||
required=True,
|
||||
help=(
|
||||
"EasyCom root containing rgb_kinetics_resnet50/, "
|
||||
"audio_wav2vec2/, and target_perframe/."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--omni-feature-root",
|
||||
type=Path,
|
||||
required=True,
|
||||
help="Existing directory containing causal Session_X.pt Omni states.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--base-checkpoint",
|
||||
type=Path,
|
||||
required=True,
|
||||
help="Released EgoSpeak miniroad_easycom_AV.pth state dict.",
|
||||
)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument("--train-sessions", nargs="+", default=None)
|
||||
parser.add_argument("--eval-sessions", nargs="+", default=None)
|
||||
parser.add_argument(
|
||||
"--expected-omni-mode",
|
||||
default="open-stream-replay",
|
||||
help="Expected saved extraction mode; pass 'any' only for an audit run.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--omni-layer-fusion",
|
||||
choices=("single", "last", "mean"),
|
||||
default="mean",
|
||||
help="How to collapse saved [T,L,D] Thinker states.",
|
||||
)
|
||||
parser.add_argument("--epochs", type=int, default=4)
|
||||
parser.add_argument("--batch-size", type=int, default=64)
|
||||
parser.add_argument(
|
||||
"--eval-chunk-size",
|
||||
type=int,
|
||||
default=0,
|
||||
help=(
|
||||
"GRU inference chunk size; 0 processes each session in one pass "
|
||||
"to match released MiniROADA evaluation."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--num-workers", type=int, default=4)
|
||||
parser.add_argument("--window-size", type=int, default=128)
|
||||
parser.add_argument("--stride", type=int, default=4)
|
||||
parser.add_argument("--fusion-dim", type=int, default=256)
|
||||
parser.add_argument("--tail-embedding-dim", type=int, default=32)
|
||||
parser.add_argument("--dropout", type=float, default=0.2)
|
||||
parser.add_argument("--lr", type=float, default=1e-4)
|
||||
parser.add_argument("--weight-decay", type=float, default=0.05)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--amp", action="store_true")
|
||||
parser.add_argument(
|
||||
"--max-train-examples",
|
||||
type=int,
|
||||
help="Smoke-test cap; omit for a complete run.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--reference-prediction-root",
|
||||
type=Path,
|
||||
help=(
|
||||
"Optional released MiniROADA raw eval directory. The frozen base "
|
||||
"is compared against Session_X.npy files before results are accepted."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--reference-tolerance", type=float, default=1e-5)
|
||||
parser.add_argument("--no-progress", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def set_seed(seed: int) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.benchmark = False
|
||||
torch.backends.cudnn.deterministic = True
|
||||
|
||||
|
||||
def resolve_device(value: str) -> torch.device:
|
||||
device = torch.device(value)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError("CUDA was requested but is unavailable")
|
||||
return device
|
||||
|
||||
|
||||
def normalize_sessions(
|
||||
values: Sequence[str] | None,
|
||||
default: Sequence[str],
|
||||
) -> list[str]:
|
||||
sessions = [normalize_session_name(value) for value in (values or default)]
|
||||
if len(sessions) != len(set(sessions)):
|
||||
raise ValueError("Duplicate sessions were supplied")
|
||||
return sessions
|
||||
|
||||
|
||||
def load_sessions(
|
||||
args: argparse.Namespace,
|
||||
sessions: Sequence[str],
|
||||
) -> list[EasyComSessionArrays]:
|
||||
expected_mode = (
|
||||
None if args.expected_omni_mode == "any" else args.expected_omni_mode
|
||||
)
|
||||
loaded = []
|
||||
for session in sessions:
|
||||
item = load_easycom_session_arrays(
|
||||
args.egospeak_feature_root,
|
||||
args.omni_feature_root,
|
||||
session,
|
||||
expected_omni_mode=expected_mode,
|
||||
omni_layer_fusion=args.omni_layer_fusion,
|
||||
)
|
||||
loaded.append(item)
|
||||
print(
|
||||
f"{item.name}: steps={item.num_steps}, "
|
||||
f"omni={tuple(item.omni.shape)}, "
|
||||
f"mode={item.omni_metadata['mode']}"
|
||||
)
|
||||
dims = {int(item.omni.shape[-1]) for item in loaded}
|
||||
if len(dims) != 1:
|
||||
raise ValueError(f"Omni dimensions differ across sessions: {sorted(dims)}")
|
||||
signatures = {
|
||||
(
|
||||
item.omni_metadata["model_id"],
|
||||
item.omni_metadata["mode"],
|
||||
item.omni_metadata["config"].get("memory_seconds"),
|
||||
item.omni_metadata["config"].get("video_fps"),
|
||||
item.omni_metadata["config"].get("readout_last_n_layers"),
|
||||
item.omni_metadata["config"].get("readout_layer_fusion"),
|
||||
tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]),
|
||||
)
|
||||
for item in loaded
|
||||
}
|
||||
if len(signatures) != 1:
|
||||
raise ValueError(
|
||||
"Selected sessions mix incompatible Omni extraction contracts: "
|
||||
f"{sorted(map(str, signatures))}"
|
||||
)
|
||||
return loaded
|
||||
|
||||
|
||||
def omni_contract_signature(item: EasyComSessionArrays) -> tuple[Any, ...]:
|
||||
config = item.omni_metadata["config"]
|
||||
return (
|
||||
item.omni_metadata["model_id"],
|
||||
item.omni_metadata["mode"],
|
||||
config.get("memory_seconds"),
|
||||
config.get("video_fps"),
|
||||
config.get("readout_last_n_layers"),
|
||||
config.get("readout_layer_fusion"),
|
||||
tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]),
|
||||
)
|
||||
|
||||
|
||||
def autocast_context(device: torch.device, enabled: bool):
|
||||
if not enabled:
|
||||
return nullcontext()
|
||||
if device.type != "cuda":
|
||||
raise ValueError("--amp currently requires CUDA")
|
||||
return torch.autocast(device_type="cuda", dtype=torch.float16)
|
||||
|
||||
|
||||
def make_loader(
|
||||
dataset: EgoSpeakOmniWindowDataset,
|
||||
args: argparse.Namespace,
|
||||
epoch: int,
|
||||
) -> DataLoader:
|
||||
generator = torch.Generator()
|
||||
generator.manual_seed(args.seed + epoch)
|
||||
return DataLoader(
|
||||
dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=torch.cuda.is_available(),
|
||||
persistent_workers=False,
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def verify_zero_initialization(
|
||||
model: EgoSpeakOmniResidual,
|
||||
dataset: EgoSpeakOmniWindowDataset,
|
||||
device: torch.device,
|
||||
) -> float:
|
||||
model.eval()
|
||||
rgb, audio, omni, tail_bin, _, _, _ = dataset[0]
|
||||
fused, base = model.forward_last(
|
||||
rgb[None].to(device),
|
||||
audio[None].to(device),
|
||||
omni[None].to(device),
|
||||
tail_bin[None].to(device),
|
||||
)
|
||||
maximum = float((fused - base).abs().max().cpu())
|
||||
if maximum != 0.0:
|
||||
raise RuntimeError(
|
||||
"Zero-initialized residual does not exactly recover MiniROADA: "
|
||||
f"max_abs_difference={maximum}"
|
||||
)
|
||||
print("Initialization audit: fused logits == frozen EgoSpeak logits exactly")
|
||||
return maximum
|
||||
|
||||
|
||||
def train_one_epoch(
|
||||
model: EgoSpeakOmniResidual,
|
||||
loader: DataLoader,
|
||||
optimizer: torch.optim.Optimizer,
|
||||
scaler: torch.cuda.amp.GradScaler,
|
||||
device: torch.device,
|
||||
*,
|
||||
amp: bool,
|
||||
show_progress: bool,
|
||||
) -> dict[str, float]:
|
||||
model.train()
|
||||
total_loss = 0.0
|
||||
total_examples = 0
|
||||
iterator = tqdm(loader, desc="train", disable=not show_progress, dynamic_ncols=True)
|
||||
for rgb, audio, omni, tail_bin, targets, _, _ in iterator:
|
||||
rgb = rgb.to(device=device, dtype=torch.float32, non_blocking=True)
|
||||
audio = audio.to(device=device, dtype=torch.float32, non_blocking=True)
|
||||
omni = omni.to(device=device, dtype=torch.float32, non_blocking=True)
|
||||
tail_bin = tail_bin.to(device=device, dtype=torch.long, non_blocking=True)
|
||||
targets = targets.to(device=device, dtype=torch.float32, non_blocking=True)
|
||||
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
with autocast_context(device, amp):
|
||||
logits, _ = model.forward_last(rgb, audio, omni, tail_bin)
|
||||
loss = miniroad_normalized_loss(logits.float(), targets.float())
|
||||
scaler.scale(loss).backward()
|
||||
scaler.step(optimizer)
|
||||
scaler.update()
|
||||
|
||||
batch_size = len(rgb)
|
||||
total_examples += batch_size
|
||||
total_loss += float(loss.detach().cpu()) * batch_size
|
||||
iterator.set_postfix_str(f"loss={total_loss / total_examples:.5f}")
|
||||
return {
|
||||
"loss": total_loss / total_examples,
|
||||
"num_examples": total_examples,
|
||||
}
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict_session(
|
||||
model: EgoSpeakOmniResidual,
|
||||
session: EasyComSessionArrays,
|
||||
device: torch.device,
|
||||
*,
|
||||
chunk_size: int,
|
||||
amp: bool,
|
||||
show_progress: bool,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
model.eval()
|
||||
num_sources = session.num_steps - NUM_HORIZONS
|
||||
if num_sources <= 0:
|
||||
raise ValueError(f"{session.name}: no valid ten-horizon source rows")
|
||||
if chunk_size == 0:
|
||||
chunk_size = num_sources
|
||||
fused_chunks: list[torch.Tensor] = []
|
||||
base_chunks: list[torch.Tensor] = []
|
||||
recurrent_state = None
|
||||
starts = range(0, num_sources, chunk_size)
|
||||
for start in tqdm(
|
||||
starts,
|
||||
total=math.ceil(num_sources / chunk_size),
|
||||
desc=f"infer {session.name}",
|
||||
disable=not show_progress,
|
||||
dynamic_ncols=True,
|
||||
):
|
||||
end = min(start + chunk_size, num_sources)
|
||||
rgb = torch.from_numpy(
|
||||
np.asarray(session.rgb[start:end], dtype=np.float32).copy()
|
||||
)[None].to(device)
|
||||
audio = torch.from_numpy(
|
||||
np.asarray(session.audio[start:end], dtype=np.float32).copy()
|
||||
)[None].to(device)
|
||||
omni = session.omni[start:end][None].to(device=device, dtype=torch.float32)
|
||||
tail_bin = session.tail_bins[start:end][None].to(device)
|
||||
with autocast_context(device, amp):
|
||||
fused, base, recurrent_state = model.forward_sequence_chunk(
|
||||
rgb,
|
||||
audio,
|
||||
omni,
|
||||
tail_bin,
|
||||
recurrent_state,
|
||||
)
|
||||
recurrent_state = recurrent_state.detach()
|
||||
fused_chunks.append(torch.softmax(fused.float(), dim=-1)[0].cpu())
|
||||
base_chunks.append(torch.softmax(base.float(), dim=-1)[0].cpu())
|
||||
return torch.cat(fused_chunks).numpy(), torch.cat(base_chunks).numpy()
|
||||
|
||||
|
||||
def evaluate_predictions(
|
||||
predictions: dict[str, np.ndarray],
|
||||
valid: dict[str, np.ndarray],
|
||||
targets: dict[str, np.ndarray],
|
||||
) -> dict[str, Any]:
|
||||
full_results = []
|
||||
causal_results = []
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
full_scores = np.concatenate(
|
||||
[item[:, :, horizon] for item in predictions.values()], axis=0
|
||||
)
|
||||
full_targets = np.concatenate(list(targets.values()), axis=0)
|
||||
full = average_precision_metrics(full_scores, full_targets)
|
||||
full.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS)
|
||||
full_results.append(full)
|
||||
|
||||
causal_scores = np.concatenate(
|
||||
[
|
||||
predictions[name][valid[name][:, horizon], :, horizon]
|
||||
for name in predictions
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
causal_targets = np.concatenate(
|
||||
[targets[name][valid[name][:, horizon]] for name in predictions],
|
||||
axis=0,
|
||||
)
|
||||
causal = average_precision_metrics(causal_scores, causal_targets)
|
||||
causal.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS)
|
||||
causal_results.append(causal)
|
||||
return {
|
||||
"egospeak_full_timeline": {
|
||||
"horizons": full_results,
|
||||
"mean_map": float(np.mean([item["map"] for item in full_results])),
|
||||
},
|
||||
"causal_valid_only": {
|
||||
"horizons": causal_results,
|
||||
"mean_map": float(np.mean([item["map"] for item in causal_results])),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def print_metrics(label: str, metrics: dict[str, Any]) -> None:
|
||||
print(f"\n=== {label} ===")
|
||||
for protocol_name in ("egospeak_full_timeline", "causal_valid_only"):
|
||||
protocol = metrics[protocol_name]
|
||||
print(f"{protocol_name}:")
|
||||
for item in protocol["horizons"]:
|
||||
class_text = ", ".join(
|
||||
f"{name}={value * 100:.3f}" for name, value in item["class_ap"].items()
|
||||
)
|
||||
print(f" {item['seconds']:.1f}s mAP={item['map'] * 100:.3f}; {class_text}")
|
||||
print(f" mean mAP={protocol['mean_map'] * 100:.3f}")
|
||||
|
||||
|
||||
def prepare_output(path: Path) -> Path:
|
||||
path = path.expanduser().resolve()
|
||||
if path.exists() and any(path.iterdir()):
|
||||
raise FileExistsError(f"Refusing to overwrite nonempty output: {path}")
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
for name in (
|
||||
"raw_eval",
|
||||
"pred",
|
||||
"valid",
|
||||
"frozen_egospeak_raw_eval",
|
||||
"frozen_egospeak_pred",
|
||||
):
|
||||
(path / name).mkdir()
|
||||
return path
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.epochs <= 0 or args.batch_size <= 0 or args.eval_chunk_size < 0:
|
||||
raise ValueError(
|
||||
"Epoch and batch sizes must be positive; evaluation chunk size "
|
||||
"must be nonnegative"
|
||||
)
|
||||
if args.reference_tolerance < 0:
|
||||
raise ValueError("--reference-tolerance must be nonnegative")
|
||||
set_seed(args.seed)
|
||||
device = resolve_device(args.device)
|
||||
|
||||
train_names = normalize_sessions(args.train_sessions, OFFICIAL_TRAIN_SESSIONS)
|
||||
eval_names = normalize_sessions(args.eval_sessions, OFFICIAL_TEST_SESSIONS)
|
||||
train_sessions = load_sessions(args, train_names)
|
||||
eval_sessions = load_sessions(args, eval_names)
|
||||
omni_dims = {int(item.omni.shape[-1]) for item in (*train_sessions, *eval_sessions)}
|
||||
if len(omni_dims) != 1:
|
||||
raise ValueError(f"Omni dimensions differ: {sorted(omni_dims)}")
|
||||
contracts = {
|
||||
omni_contract_signature(item) for item in (*train_sessions, *eval_sessions)
|
||||
}
|
||||
if len(contracts) != 1:
|
||||
raise ValueError(
|
||||
"Train/evaluation sessions use incompatible Omni extraction "
|
||||
f"contracts: {sorted(map(str, contracts))}"
|
||||
)
|
||||
omni_dim = omni_dims.pop()
|
||||
|
||||
base = load_frozen_miniroada(args.base_checkpoint)
|
||||
model = EgoSpeakOmniResidual(
|
||||
base,
|
||||
omni_dim=omni_dim,
|
||||
fusion_dim=args.fusion_dim,
|
||||
tail_embedding_dim=args.tail_embedding_dim,
|
||||
dropout=args.dropout,
|
||||
).to(device)
|
||||
trainable = model.trainable_parameters()
|
||||
print(f"device: {device}")
|
||||
print(f"train sessions: {', '.join(train_names)}")
|
||||
print(f"evaluation sessions: {', '.join(eval_names)}")
|
||||
print(
|
||||
f"parameters: total={sum(p.numel() for p in model.parameters()):,}, "
|
||||
f"trainable={sum(p.numel() for p in trainable):,}, "
|
||||
"MiniROADA frozen=True"
|
||||
)
|
||||
|
||||
dataset = EgoSpeakOmniWindowDataset(
|
||||
train_sessions,
|
||||
window_size=args.window_size,
|
||||
stride=args.stride,
|
||||
seed=args.seed,
|
||||
max_examples=args.max_train_examples,
|
||||
)
|
||||
initialization_difference = verify_zero_initialization(model, dataset, device)
|
||||
output_dir = prepare_output(args.output_dir)
|
||||
optimizer = torch.optim.AdamW(
|
||||
trainable,
|
||||
lr=args.lr,
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
scaler = torch.cuda.amp.GradScaler(enabled=args.amp)
|
||||
history = []
|
||||
for epoch in range(1, args.epochs + 1):
|
||||
dataset.resample(epoch - 1)
|
||||
loader = make_loader(dataset, args, epoch)
|
||||
result = train_one_epoch(
|
||||
model,
|
||||
loader,
|
||||
optimizer,
|
||||
scaler,
|
||||
device,
|
||||
amp=args.amp,
|
||||
show_progress=not args.no_progress,
|
||||
)
|
||||
result["epoch"] = epoch
|
||||
history.append(result)
|
||||
print(
|
||||
f"epoch {epoch}/{args.epochs}: loss={result['loss']:.6f}, "
|
||||
f"examples={result['num_examples']}"
|
||||
)
|
||||
|
||||
checkpoint_path = output_dir / "checkpoint_last.pt"
|
||||
checkpoint_payload = {
|
||||
"format": "egospeak-omni-zero-residual-v1",
|
||||
"model_state_dict": model.state_dict(),
|
||||
"model_config": {
|
||||
"omni_dim": omni_dim,
|
||||
"fusion_dim": args.fusion_dim,
|
||||
"tail_embedding_dim": args.tail_embedding_dim,
|
||||
"num_tail_bins": NUM_HORIZONS,
|
||||
"dropout": args.dropout,
|
||||
},
|
||||
"base_checkpoint": str(args.base_checkpoint.expanduser().resolve()),
|
||||
"base_checkpoint_sha256": sha256_file(
|
||||
args.base_checkpoint.expanduser().resolve()
|
||||
),
|
||||
"train_sessions": train_names,
|
||||
"eval_sessions": eval_names,
|
||||
"training_config": {
|
||||
"epochs": args.epochs,
|
||||
"batch_size": args.batch_size,
|
||||
"window_size": args.window_size,
|
||||
"stride": args.stride,
|
||||
"lr": args.lr,
|
||||
"weight_decay": args.weight_decay,
|
||||
"seed": args.seed,
|
||||
"amp": args.amp,
|
||||
"omni_layer_fusion": args.omni_layer_fusion,
|
||||
"expected_omni_mode": args.expected_omni_mode,
|
||||
"max_train_examples": args.max_train_examples,
|
||||
},
|
||||
"initialization_max_abs_logit_difference": initialization_difference,
|
||||
"history": history,
|
||||
}
|
||||
torch.save(checkpoint_payload, checkpoint_path)
|
||||
print(f"Saved checkpoint: {checkpoint_path}")
|
||||
|
||||
fused_predictions: dict[str, np.ndarray] = {}
|
||||
base_predictions: dict[str, np.ndarray] = {}
|
||||
valid: dict[str, np.ndarray] = {}
|
||||
targets: dict[str, np.ndarray] = {}
|
||||
reference_audit: dict[str, Any] = {}
|
||||
for session in eval_sessions:
|
||||
fused_raw, base_raw = predict_session(
|
||||
model,
|
||||
session,
|
||||
device,
|
||||
chunk_size=args.eval_chunk_size,
|
||||
amp=args.amp,
|
||||
show_progress=not args.no_progress,
|
||||
)
|
||||
fused_aligned, session_valid = align_source_predictions(
|
||||
fused_raw, session.num_steps
|
||||
)
|
||||
base_aligned, base_valid = align_source_predictions(base_raw, session.num_steps)
|
||||
if not np.array_equal(session_valid, base_valid):
|
||||
raise RuntimeError("Fused/base causal-valid masks unexpectedly differ")
|
||||
np.save(output_dir / "raw_eval" / f"{session.name}.npy", fused_raw)
|
||||
np.save(output_dir / "pred" / f"{session.name}.npy", fused_aligned)
|
||||
np.save(output_dir / "valid" / f"{session.name}.npy", session_valid)
|
||||
np.save(
|
||||
output_dir / "frozen_egospeak_raw_eval" / f"{session.name}.npy",
|
||||
base_raw,
|
||||
)
|
||||
np.save(
|
||||
output_dir / "frozen_egospeak_pred" / f"{session.name}.npy",
|
||||
base_aligned,
|
||||
)
|
||||
fused_predictions[session.name] = fused_aligned
|
||||
base_predictions[session.name] = base_aligned
|
||||
valid[session.name] = session_valid
|
||||
targets[session.name] = np.asarray(session.targets, dtype=np.float32)
|
||||
|
||||
if args.reference_prediction_root is not None:
|
||||
reference_path = (
|
||||
args.reference_prediction_root.expanduser().resolve()
|
||||
/ f"{session.name}.npy"
|
||||
)
|
||||
if not reference_path.is_file():
|
||||
raise FileNotFoundError(reference_path)
|
||||
reference = np.load(reference_path, allow_pickle=False)
|
||||
if reference.shape != base_raw.shape:
|
||||
raise ValueError(
|
||||
f"{session.name}: reference shape {reference.shape} != "
|
||||
f"{base_raw.shape}"
|
||||
)
|
||||
difference = np.abs(reference.astype(np.float32) - base_raw)
|
||||
maximum = float(difference.max())
|
||||
reference_audit[session.name] = {
|
||||
"path": str(reference_path),
|
||||
"max_abs_difference": maximum,
|
||||
"mean_abs_difference": float(difference.mean()),
|
||||
}
|
||||
print(
|
||||
f"{session.name}: frozen-base/reference max abs "
|
||||
f"difference={maximum:.8g}"
|
||||
)
|
||||
if maximum > args.reference_tolerance:
|
||||
raise RuntimeError(
|
||||
f"{session.name}: frozen MiniROADA did not reproduce the "
|
||||
f"reference within {args.reference_tolerance}: {maximum}"
|
||||
)
|
||||
|
||||
fused_metrics = evaluate_predictions(fused_predictions, valid, targets)
|
||||
base_metrics = evaluate_predictions(base_predictions, valid, targets)
|
||||
metrics = {
|
||||
**fused_metrics,
|
||||
"format": "egospeak-omni-zero-residual-eval-v1",
|
||||
"classes": list(CLASS_NAMES),
|
||||
"sessions": eval_names,
|
||||
"prediction_layout": "[target_time, class, future_horizon]",
|
||||
"raw_prediction_layout": "[source_time, future_horizon, class]",
|
||||
"checkpoint": str(checkpoint_path),
|
||||
"train_sessions": train_names,
|
||||
"frozen_egospeak_baseline": base_metrics,
|
||||
"reference_prediction_audit": reference_audit,
|
||||
"initialization_max_abs_logit_difference": initialization_difference,
|
||||
"omni_sources": {
|
||||
item.name: item.omni_metadata for item in (*train_sessions, *eval_sessions)
|
||||
},
|
||||
}
|
||||
metrics_path = output_dir / "metrics.json"
|
||||
metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
|
||||
print_metrics("fused", fused_metrics)
|
||||
print_metrics("frozen EgoSpeak", base_metrics)
|
||||
print(f"\nSaved metrics: {metrics_path}")
|
||||
print(f"Fused evaluator input: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user