diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index 5397280..edca672 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -43,6 +43,13 @@ result. - Ego4D batch and 8-GPU local-feature launchers, with strict frame audits. - Timing-derived wearer floor-state and event diagnostics. +### Implemented, awaiting remote runtime verification + +- F1 Frame-set Attention loader/head/train/infer path. It uses fused G+A as + one query over the latest half-open causal person set, retains direct A in + the head, and dynamically pads people per batch. Do not report it as a + measured method until the L20 smoke and full controls pass. + ### Current blocker `evaluate_floor_state_changes.py` counts a wearer transition anywhere in a diff --git a/docs/local_feature_baseline.md b/docs/local_feature_baseline.md index 8a49d99..b417d17 100644 --- a/docs/local_feature_baseline.md +++ b/docs/local_feature_baseline.md @@ -174,6 +174,27 @@ warmup_factor=0.3、dropout=0.2、LSTR-MCE、不加 transition objective。 Loader 会检查 sample、时间、label_indices、维度、valid mask 和 encoder/pooling 配置。 训练与评测必须用同一种 AP 协议作比较;不要混用 full-timeline 与 causal-valid 的表格。 +## 4. F1 Frame-set Attention + +已有完整局部 sidecar 的 persons 字段保存了逐帧、逐人的原生 2048D 特征,因此 F1 +不需要重新提取图片或音频。新增训练参数为: + + --local-feature-root "$LOCAL_ROOT" \ + --local-modalities av \ + --local-fusion frame-set-attention \ + --local-projection-dim 256 \ + --person-attention-dim 256 \ + --person-attention-heads 4 + +F1 用融合后的 G 和直接局部音频 A 构造一个 query,读取当前最新半开因果帧中的可变人物 +集合。人物集合只在 batch 内动态 padding;NULL person 处理没有有效框的时刻。A 仍直接 +进入最终 head。第一版不使用 bbox、track ID、跨帧 GRU 或双向 AV attention。 + +训练输出会打印 source frame 数、person observation 数、max_people 和 attention 配置。 +推理只需继续提供 local feature root;融合方式和 attention 参数从 checkpoint 恢复。 +正式比较需要重新训练同 split、seed 和优化设置的 pooled AV 对照,不能直接把不同开发 +协议的旧 checkpoint 当作唯一基线。 + ## Ego4D 接口 完整命令见 [Ego4D 四组局部增强实验](ego4d_local_baselines.md):新增 manifest 批量提取、 diff --git a/docs/progressive_feature_fusion_route_2026-09-02.md b/docs/progressive_feature_fusion_route_2026-09-02.md index ab87492..63b72c3 100644 --- a/docs/progressive_feature_fusion_route_2026-09-02.md +++ b/docs/progressive_feature_fusion_route_2026-09-02.md @@ -21,14 +21,17 @@ 融合实验从 F1 Frame-set Attention 开始。后续模块只有在前一级通过预先定义的 gate 后 才实施。 -这是一份实验设计,不代表 F1--F6 或 D1 已经实现或验证。 +实现状态(2026-09-02):F1 的 loader、head、train/infer 路径和回归测试已经加入代码, +但尚未在远程 PyTorch/L20 环境完成 runtime smoke,不能写成已验证结果。F2--F6 和 D1 +仍然只是计划。 ## 2. 不可改变的实验条件 所有阶段共同遵守: - 目标与全局 state 均为 5 Hz;第 t 行只能使用截止到 (t+1)/5 秒的媒体。 -- 视频人物特征若为 2 FPS,在 5 Hz 预测点选择时间不晚于 cutoff 的最新人物集合。 +- 视频人物特征若为 2 FPS,在 5 Hz 预测点选择严格早于 cutoff 的最新人物集合,与已有 + half-open local audio/visual 契约一致。 - 不允许用未来帧插值;同一 2 FPS 人物集合可以供多个 5 Hz query 读取,但不能伪装成 新观测去重复更新时序状态。 - EasyCom Session 4--12 为官方训练集,Session 1--3 为最终测试集。模型和阈值选择使用 @@ -128,11 +131,11 @@ F1 暂不加入 bbox、track GRU、人物自注意力或 audio-to-person 反向 对每个 5 Hz cutoff c_t: -1. 找到 frame_time 小于或等于 c_t 的最新 local frame; +1. 找到 frame_time 严格小于 c_t 的最新 local frame; 2. 读取该帧 frame_offsets 范围内全部逐人 feature; 3. 如果尚无帧,使用仅含 NULL person 的集合; 4. 不允许 nearest-neighbor 选到 c_t 之后的帧; -5. 记录 chosen_frame_time,测试所有 chosen_frame_time 小于或等于 cutoff。 +5. 记录 chosen_frame_time,测试所有 chosen_frame_time 严格小于 cutoff。 ### 5.5 F1 必跑对照 diff --git a/easycom_head.py b/easycom_head.py index 95cd950..bdeff9e 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -5,13 +5,16 @@ from __future__ import annotations import csv from dataclasses import dataclass from pathlib import Path -from typing import Any, Literal, Sequence +from typing import TYPE_CHECKING, Any, Literal, Sequence from uuid import UUID import numpy as np import torch from torch import nn +if TYPE_CHECKING: + from local_turn_features import FramePersonFeatures + OFFICIAL_TRAIN_SESSIONS = ( "Session_7", @@ -58,6 +61,7 @@ LAYER_FUSION_MODES = ( "scalar-mixture", "concat-mlp", ) +LOCAL_FUSION_MODES = ("pooled", "frame-set-attention") DatasetName = Literal["easycom", "ego4d"] @@ -77,6 +81,9 @@ class TurnPredictionHead(nn.Module): local_visual_dim: int = 0, local_audio_dim: int = 0, local_projection_dim: int = 256, + local_fusion: str = "pooled", + person_attention_dim: int = 256, + person_attention_heads: int = 4, ) -> None: super().__init__() if num_input_layers <= 0: @@ -116,6 +123,24 @@ class TurnPredictionHead(nn.Module): self.local_visual_dim = local_visual_dim self.local_audio_dim = local_audio_dim self.local_projection_dim = local_projection_dim + if local_fusion not in LOCAL_FUSION_MODES: + raise ValueError( + f"Unknown local_fusion={local_fusion!r}; expected one of " + f"{LOCAL_FUSION_MODES}" + ) + if person_attention_dim <= 0 or person_attention_heads <= 0: + raise ValueError("Person attention dimension/heads must be positive") + if person_attention_dim % person_attention_heads: + raise ValueError("Person attention dimension must be divisible by heads") + if local_fusion == "frame-set-attention" and ( + local_visual_dim <= 0 or local_audio_dim <= 0 + ): + raise ValueError( + "Frame-set Attention requires both local visual people and audio" + ) + self.local_fusion = local_fusion + self.person_attention_dim = person_attention_dim + self.person_attention_heads = person_attention_heads if self.learnable_layer_mixture: # Zero initialization exactly reproduces a uniform layer mean at # step zero while allowing training to select useful query layers. @@ -133,7 +158,12 @@ class TurnPredictionHead(nn.Module): self.layer_fusion_net = None self.local_projections = nn.ModuleDict() fused_size = hidden_size - for name, dimension in (("visual", local_visual_dim), ("audio", local_audio_dim)): + projected_modalities = ( + (("audio", local_audio_dim),) + if local_fusion == "frame-set-attention" + else (("visual", local_visual_dim), ("audio", local_audio_dim)) + ) + for name, dimension in projected_modalities: if dimension: self.local_projections[name] = ( nn.Sequential( @@ -145,6 +175,36 @@ class TurnPredictionHead(nn.Module): else nn.Identity() ) fused_size += (local_projection_dim or dimension) + 1 + self.person_input_projection: nn.Module | None = None + self.person_state_query: nn.Module | None = None + self.person_audio_query: nn.Module | None = None + self.person_attention: nn.MultiheadAttention | None = None + self.null_person: nn.Parameter | None = None + if local_fusion == "frame-set-attention": + self.person_input_projection = nn.Sequential( + nn.LayerNorm(local_visual_dim), + nn.Linear(local_visual_dim, person_attention_dim), + nn.GELU(), + ) + self.person_state_query = nn.Sequential( + nn.LayerNorm(hidden_size), + nn.Linear(hidden_size, person_attention_dim), + ) + self.person_audio_query = nn.Sequential( + nn.LayerNorm(local_audio_dim), + nn.Linear(local_audio_dim, person_attention_dim), + ) + self.person_attention = nn.MultiheadAttention( + person_attention_dim, + person_attention_heads, + dropout=dropout, + batch_first=True, + ) + self.null_person = nn.Parameter( + torch.zeros(1, 1, person_attention_dim) + ) + nn.init.normal_(self.null_person, std=0.02) + fused_size += person_attention_dim + 1 self.net = nn.Sequential( nn.LayerNorm(fused_size), nn.Linear(fused_size, hidden_size // 2), @@ -190,17 +250,99 @@ class TurnPredictionHead(nn.Module): @property def local_input_size(self) -> int: + if self.local_fusion == "frame-set-attention": + return self.local_audio_dim + 1 return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d) + def _frame_set_context( + self, + states: torch.Tensor, + local: torch.Tensor, + person_features: torch.Tensor | None, + person_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if states.ndim != 2: + raise ValueError("Frame-set Attention expects batched [B,D] states") + if local.shape != (len(states), self.local_audio_dim + 1): + raise ValueError( + "Frame-set Attention requires aligned [audio,valid] local features" + ) + if ( + person_features is None + or person_features.ndim != 3 + or person_features.shape[0] != len(states) + or person_features.shape[2] != self.local_visual_dim + ): + raise ValueError( + "Frame-set Attention requires person features [B,P,visual_dim]" + ) + if ( + person_mask is None + or person_mask.dtype != torch.bool + or person_mask.shape != person_features.shape[:2] + ): + raise ValueError("Frame-set Attention requires a boolean person mask [B,P]") + if not torch.isfinite(person_features).all(): + raise ValueError("Person features contain NaN/Inf") + assert self.person_input_projection is not None + assert self.person_state_query is not None + assert self.person_audio_query is not None + assert self.person_attention is not None + assert self.null_person is not None + + audio_values = local[:, : self.local_audio_dim] + audio_valid = local[:, self.local_audio_dim : self.local_audio_dim + 1] + audio_projected = self.local_projections["audio"](audio_values) * audio_valid + query = ( + self.person_state_query(states) + + self.person_audio_query(audio_values) * audio_valid + ) + people = self.person_input_projection(person_features) + null = self.null_person.expand(len(states), -1, -1) + memory = torch.cat((null, people), dim=1) + null_mask = torch.ones( + (len(states), 1), + dtype=torch.bool, + device=person_mask.device, + ) + memory_valid = torch.cat((null_mask, person_mask), dim=1) + context, _ = self.person_attention( + query[:, None, :], + memory, + memory, + key_padding_mask=~memory_valid, + need_weights=False, + ) + person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype) + return audio_projected, audio_valid, torch.cat( + (context[:, 0, :], person_valid), + dim=-1, + ) + def _shared_features( self, states: torch.Tensor, local: torch.Tensor | None = None, + person_features: torch.Tensor | None = None, + person_mask: torch.Tensor | None = None, ) -> torch.Tensor: """Return the activity MLP representation shared by both readouts.""" states = self._fuse_layers(states) - if self.local_input_size: + if self.local_fusion == "frame-set-attention": + if local is None: + raise ValueError("Frame-set Attention requires local audio features") + audio, audio_valid, person_context = self._frame_set_context( + states, + local, + person_features, + person_mask, + ) + states = torch.cat( + (states, audio, audio_valid, person_context), + dim=-1, + ) + elif self.local_input_size: if local is None or local.shape != (*states.shape[:-1], self.local_input_size): raise ValueError( f"Head requires aligned local features of width {self.local_input_size}" @@ -220,14 +362,25 @@ class TurnPredictionHead(nn.Module): states = torch.cat(pieces, dim=-1) elif local is not None and local.shape[-1] != 0: raise ValueError("Local features supplied to a baseline head") + if self.local_fusion != "frame-set-attention" and ( + person_features is not None or person_mask is not None + ): + raise ValueError("Person sets supplied to a non-frame-set head") return self.net[:-1](states) def forward( self, states: torch.Tensor, local: torch.Tensor | None = None, + person_features: torch.Tensor | None = None, + person_mask: torch.Tensor | None = None, ) -> torch.Tensor: - features = self._shared_features(states, local) + features = self._shared_features( + states, + local, + person_features, + person_mask, + ) logits = self.net[-1](features) return logits.view( *features.shape[:-1], self.num_horizons, self.num_classes @@ -237,6 +390,8 @@ class TurnPredictionHead(nn.Module): self, states: torch.Tensor, local: torch.Tensor | None = None, + person_features: torch.Tensor | None = None, + person_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Return activity logits and binary transition logits.""" @@ -244,7 +399,12 @@ class TurnPredictionHead(nn.Module): raise RuntimeError( "This head was created without transition_auxiliary=True" ) - features = self._shared_features(states, local) + features = self._shared_features( + states, + local, + person_features, + person_mask, + ) activity_logits = self.net[-1](features).view( *features.shape[:-1], self.num_horizons, @@ -734,6 +894,7 @@ class SessionData: label_indices: torch.LongTensor feature_metadata: dict[str, Any] local_features: torch.Tensor | None = None + local_person_features: "FramePersonFeatures | None" = None @property def hidden_size(self) -> int: @@ -768,6 +929,7 @@ def load_session_data( local_feature_root: str | Path | None = None, local_modalities: str = "av", local_contract: dict[str, Any] | None = None, + local_fusion: str = "pooled", ) -> SessionData: """Load and validate one saved Omni feature payload and target array.""" @@ -864,22 +1026,55 @@ def load_session_data( "".join(str(int(value)) for value in row.tolist()): int(count) for row, count in zip(unique_rows, row_counts) } + if local_fusion not in LOCAL_FUSION_MODES: + raise ValueError(f"Unknown local fusion: {local_fusion}") local = None + local_people = None if local_feature_root is not None: - from local_turn_features import load_local_features - - local, contract = load_local_features( - resolve_feature_path(local_feature_root, name, dataset=dataset), - payload, - sample_id=name, + local_path = resolve_feature_path( + local_feature_root, + name, dataset=dataset, - modalities=local_modalities, - expected_contract=local_contract, ) + if local_fusion == "frame-set-attention": + if local_modalities != "av": + raise ValueError( + "Frame-set Attention requires --local-modalities av" + ) + from local_turn_features import load_frame_set_features + + local, local_people, contract = load_frame_set_features( + local_path, + payload, + sample_id=name, + dataset=dataset, + expected_contract=local_contract, + ) + else: + from local_turn_features import load_local_features + + local, contract = load_local_features( + local_path, + payload, + sample_id=name, + dataset=dataset, + modalities=local_modalities, + expected_contract=local_contract, + ) metadata["local_contract"] = contract elif local_contract is not None: raise ValueError("--local-feature-root is required for this checkpoint") - return SessionData(name, states, targets, label_indices, metadata, local) + elif local_fusion != "pooled": + raise ValueError("Frame-set Attention requires --local-feature-root") + return SessionData( + name, + states, + targets, + label_indices, + metadata, + local, + local_people, + ) def average_precision_metrics( @@ -924,6 +1119,9 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]: config.setdefault("local_visual_dim", 0) config.setdefault("local_audio_dim", 0) config.setdefault("local_projection_dim", 256) + config.setdefault("local_fusion", "pooled") + config.setdefault("person_attention_dim", 256) + config.setdefault("person_attention_heads", 4) config.setdefault( "layer_fusion", ( diff --git a/infer_turn_head.py b/infer_turn_head.py index 4dded80..05e2e19 100644 --- a/infer_turn_head.py +++ b/infer_turn_head.py @@ -30,6 +30,7 @@ from easycom_head import ( output_space_for_loss, transition_targets_from_activity, ) +from local_turn_features import FramePersonFeatures def parse_args() -> argparse.Namespace: @@ -118,6 +119,7 @@ def _predict_states( show_progress: bool, loss_name: str, local_features: torch.Tensor | None = None, + local_person_features: FramePersonFeatures | None = None, ) -> tuple[np.ndarray, np.ndarray | None]: chunks: list[torch.Tensor] = [] transition_chunks: list[torch.Tensor] = [] @@ -139,11 +141,35 @@ def _predict_states( local_features[start : start + batch_size].to(device=device, dtype=torch.float32) if local_features is not None else None ) + person_batch = None + person_mask = None + if local_person_features is not None: + rows = torch.arange( + start, + min(start + batch_size, len(states)), + ) + person_batch, person_mask = local_person_features.padded(rows) + person_batch = person_batch.to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + person_mask = person_mask.to(device=device, non_blocking=True) if head.transition_auxiliary: - activity_logits, transition_logits = head.forward_with_transition(batch, side) + activity_logits, transition_logits = head.forward_with_transition( + batch, + side, + person_batch, + person_mask, + ) transition_chunks.append(transition_logits.sigmoid().cpu()) else: - activity_logits = head(batch, side) + activity_logits = head( + batch, + side, + person_batch, + person_mask, + ) chunks.append( marginal_probabilities_from_logits(activity_logits, loss_name).cpu() ) @@ -403,6 +429,7 @@ def main() -> None: local_feature_root=args.local_feature_root, local_modalities=local_contract["modalities"] if local_contract else "av", local_contract=local_contract, + local_fusion=head.local_fusion, ) if data.hidden_size != int(head_config["hidden_size"]): raise ValueError( @@ -423,6 +450,7 @@ def main() -> None: show_progress=not args.no_progress, loss_name=loss_name, local_features=data.local_features, + local_person_features=data.local_person_features, ) predictions, valid = _align_to_egospeak_timeline( probabilities, @@ -484,6 +512,7 @@ def main() -> None: "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"), "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"), "layer_fusion": head.layer_fusion, + "local_fusion": head.local_fusion, "local_contract": local_contract, "transition_auxiliary": head.transition_auxiliary, "transition_objective": checkpoint_payload.get( diff --git a/local_turn_features.py b/local_turn_features.py index 921d21d..d6c2f6f 100644 --- a/local_turn_features.py +++ b/local_turn_features.py @@ -1,4 +1,4 @@ -"""Validated, causal local-feature sidecars for the pooled B1 baseline. +"""Validated, causal pooled and frame-set local-feature sidecars. No model/decoder imports: these utilities are also used by CPU head training. Sidecars store encoder-native vectors, never a randomly initialized projection. @@ -6,6 +6,7 @@ Sidecars store encoder-native vectors, never a randomly initialized projection. from __future__ import annotations +from dataclasses import dataclass import math from pathlib import Path from typing import Any @@ -16,6 +17,150 @@ LOCAL_FORMAT = "qwen-local-pooled-v1" MODALITIES = {"visual": ("visual",), "audio": ("audio",), "av": ("visual", "audio")} +@dataclass(frozen=True) +class FramePersonFeatures: + """Ragged person sets indexed by aligned 5 Hz rows without frame duplication.""" + + features: torch.Tensor + frame_offsets: torch.LongTensor + row_to_frame: torch.LongTensor + chosen_frame_times: torch.Tensor + + def __post_init__(self) -> None: + if ( + not isinstance(self.features, torch.Tensor) + or self.features.ndim != 2 + or not self.features.is_floating_point() + or not torch.isfinite(self.features).all() + or self.features.shape[1] <= 0 + ): + raise ValueError("Person features must be finite floating [N,D]") + if ( + not isinstance(self.frame_offsets, torch.Tensor) + or self.frame_offsets.dtype != torch.long + or self.frame_offsets.ndim != 1 + or not len(self.frame_offsets) + or int(self.frame_offsets[0]) != 0 + or int(self.frame_offsets[-1]) != len(self.features) + or bool((self.frame_offsets[1:] < self.frame_offsets[:-1]).any()) + ): + raise ValueError("Invalid person frame_offsets") + if ( + not isinstance(self.row_to_frame, torch.Tensor) + or self.row_to_frame.dtype != torch.long + or self.row_to_frame.ndim != 1 + ): + raise ValueError("row_to_frame must be int64 [T]") + num_frames = len(self.frame_offsets) - 1 + if bool(((self.row_to_frame < -1) | (self.row_to_frame >= num_frames)).any()): + raise ValueError("row_to_frame contains an invalid frame index") + if ( + not isinstance(self.chosen_frame_times, torch.Tensor) + or self.chosen_frame_times.dtype != torch.float64 + or self.chosen_frame_times.shape != self.row_to_frame.shape + ): + raise ValueError("chosen_frame_times must be float64 [T]") + missing = self.row_to_frame < 0 + if not torch.equal(torch.isnan(self.chosen_frame_times), missing): + raise ValueError("Only rows without an available frame may have NaN frame time") + if not torch.isfinite(self.chosen_frame_times[~missing]).all(): + raise ValueError("Chosen person frame times must be finite") + + def __len__(self) -> int: + return len(self.row_to_frame) + + @property + def dimension(self) -> int: + return int(self.features.shape[1]) + + @property + def max_people(self) -> int: + counts = self.frame_offsets[1:] - self.frame_offsets[:-1] + return int(counts.max()) if len(counts) else 0 + + def select_rows(self, indices: torch.Tensor) -> "FramePersonFeatures": + indices = torch.as_tensor(indices) + if indices.dtype != torch.long or indices.ndim != 1: + raise ValueError("Person row selection must be int64 [N]") + if len(indices) and (int(indices.min()) < 0 or int(indices.max()) >= len(self)): + raise IndexError("Person row selection is out of range") + return FramePersonFeatures( + self.features, + self.frame_offsets, + self.row_to_frame[indices], + self.chosen_frame_times[indices], + ) + + def padded( + self, + row_indices: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.BoolTensor]: + """Materialize only one batch as [B,P,D] plus a True-for-real mask.""" + + if row_indices is None: + frames = self.row_to_frame + else: + row_indices = torch.as_tensor(row_indices) + if row_indices.dtype != torch.long or row_indices.ndim != 1: + raise ValueError("Batch row indices must be int64 [B]") + if len(row_indices) and ( + int(row_indices.min()) < 0 or int(row_indices.max()) >= len(self) + ): + raise IndexError("Batch person row index is out of range") + frames = self.row_to_frame[row_indices] + valid_frames = frames >= 0 + safe_frames = frames.clamp_min(0) + starts = torch.zeros_like(frames) + counts = torch.zeros_like(frames) + if bool(valid_frames.any()): + starts[valid_frames] = self.frame_offsets[safe_frames[valid_frames]] + counts[valid_frames] = ( + self.frame_offsets[safe_frames[valid_frames] + 1] + - starts[valid_frames] + ) + width = int(counts.max()) if len(counts) else 0 + mask = torch.arange(width)[None, :] < counts[:, None] + if width == 0: + return self.features.new_zeros((len(frames), 0, self.dimension)), mask + person_indices = starts[:, None] + torch.arange(width)[None, :] + safe_indices = person_indices.masked_fill(~mask, 0) + values = self.features[safe_indices] + return values.masked_fill(~mask[..., None], 0), mask + + @classmethod + def concatenate( + cls, + parts: list["FramePersonFeatures"], + ) -> "FramePersonFeatures": + if not parts: + raise ValueError("Cannot concatenate no person feature sets") + dimension = parts[0].dimension + if any(part.dimension != dimension for part in parts): + raise ValueError("Person feature dimensions differ across sessions") + feature_parts = [] + offset_parts = [torch.zeros(1, dtype=torch.long)] + row_parts = [] + time_parts = [] + feature_base = 0 + frame_base = 0 + for part in parts: + feature_parts.append(part.features) + offset_parts.append(part.frame_offsets[1:] + feature_base) + rows = part.row_to_frame.clone() + valid = rows >= 0 + rows[valid] += frame_base + row_parts.append(rows) + time_parts.append(part.chosen_frame_times) + feature_base += len(part.features) + frame_base += len(part.frame_offsets) - 1 + return cls( + torch.cat(feature_parts, dim=0), + torch.cat(offset_parts, dim=0), + torch.cat(row_parts, dim=0), + torch.cat(time_parts, dim=0), + ) + + def reference_alignment(payload: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]: """Require the same full, zero-origin 5 Hz timeline used by head training.""" states = payload.get("states") @@ -34,16 +179,15 @@ def reference_alignment(payload: dict[str, Any]) -> tuple[torch.Tensor, torch.Te return times, indices -def load_local_features( +def _load_validated_payload( path: Path, reference: dict[str, Any], *, sample_id: str, dataset: str, modalities: str, - expected_contract: dict[str, Any] | None = None, -) -> tuple[torch.Tensor, dict[str, Any]]: - """Pack each requested modality as [raw vector, validity bit]. Fail closed.""" + expected_contract: dict[str, Any] | None, +) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any], dict[str, Any], dict[str, Any]]: if modalities not in MODALITIES: raise ValueError(f"Unknown local modalities: {modalities}") times, indices = reference_alignment(reference) @@ -66,6 +210,27 @@ def load_local_features( contract = {"format": LOCAL_FORMAT, "modalities": modalities, "config": config} if expected_contract is not None and contract != expected_contract: raise ValueError(f"Local encoder/pooling contract differs across sessions/checkpoint: {path}") + return times, indices, payload, config, contract + + +def load_local_features( + path: Path, + reference: dict[str, Any], + *, + sample_id: str, + dataset: str, + modalities: str, + expected_contract: dict[str, Any] | None = None, +) -> tuple[torch.Tensor, dict[str, Any]]: + """Pack each requested modality as [raw vector, validity bit]. Fail closed.""" + times, _, payload, config, contract = _load_validated_payload( + path, + reference, + sample_id=sample_id, + dataset=dataset, + modalities=modalities, + expected_contract=expected_contract, + ) chunks = [] for name in MODALITIES[modalities]: values = payload.get(name) @@ -85,6 +250,109 @@ def load_local_features( return torch.cat(chunks, dim=-1), contract +def load_frame_set_features( + path: Path, + reference: dict[str, Any], + *, + sample_id: str, + dataset: str, + expected_contract: dict[str, Any] | None = None, +) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]: + """Load direct audio and the latest causal per-frame person set per 5 Hz row.""" + + times, _, payload, config, contract = _load_validated_payload( + path, + reference, + sample_id=sample_id, + dataset=dataset, + modalities="av", + expected_contract=expected_contract, + ) + audio = payload.get("audio") + audio_valid = payload.get("audio_valid") + audio_dim = config.get("audio_dim") + if not isinstance(audio_dim, int) or audio_dim <= 0: + raise ValueError(f"Missing audio_dim: {path}") + if not isinstance(audio, torch.Tensor) or audio.shape != (len(times), audio_dim): + raise ValueError(f"Invalid audio shape: {path}") + if not audio.is_floating_point() or not torch.isfinite(audio).all(): + raise ValueError(f"Non-finite/non-floating audio features: {path}") + if ( + not isinstance(audio_valid, torch.Tensor) + or audio_valid.dtype != torch.bool + or audio_valid.shape != (len(times),) + ): + raise ValueError(f"Invalid audio_valid mask: {path}") + if torch.count_nonzero(audio[~audio_valid]): + raise ValueError(f"Invalid audio rows must be stored as zeros: {path}") + + visual_dim = config.get("visual_dim") + persons = payload.get("persons") + if not isinstance(visual_dim, int) or visual_dim <= 0: + raise ValueError(f"Missing visual_dim: {path}") + if not isinstance(persons, dict): + raise ValueError(f"Missing ragged person features: {path}") + features = persons.get("features") + frame_times = persons.get("frame_times") + frame_offsets = persons.get("frame_offsets") + records = persons.get("records") + if ( + not isinstance(features, torch.Tensor) + or features.ndim != 2 + or features.shape[1] != visual_dim + or not features.is_floating_point() + or not torch.isfinite(features).all() + ): + raise ValueError(f"Invalid person feature shape/content: {path}") + if ( + not isinstance(frame_times, torch.Tensor) + or frame_times.ndim != 1 + or not frame_times.is_floating_point() + ): + raise ValueError(f"Invalid person frame_times: {path}") + frame_times = frame_times.to(dtype=torch.float64) + if ( + not torch.isfinite(frame_times).all() + or bool((frame_times < 0).any()) + or bool((frame_times[1:] <= frame_times[:-1]).any()) + ): + raise ValueError(f"Person frame_times must be finite, nonnegative and increasing: {path}") + if ( + not isinstance(frame_offsets, torch.Tensor) + or frame_offsets.dtype != torch.long + or frame_offsets.shape != (len(frame_times) + 1,) + or not len(frame_offsets) + or int(frame_offsets[0]) != 0 + or int(frame_offsets[-1]) != len(features) + or bool((frame_offsets[1:] < frame_offsets[:-1]).any()) + ): + raise ValueError(f"Invalid person frame_offsets: {path}") + if not isinstance(records, list) or len(records) != len(features): + raise ValueError(f"Person records/features length mismatch: {path}") + if len(frame_times) and float(frame_times[-1]) > float(times[-1]) + 1e-8: + raise ValueError(f"Person frame timeline extends beyond the reference: {path}") + + # Match the sidecar's half-open policy: a frame stamped exactly at the + # state cutoff is not yet part of that state's observed prefix. + row_to_frame = torch.searchsorted(frame_times, times, right=False) - 1 + chosen = torch.full(times.shape, float("nan"), dtype=torch.float64) + has_frame = row_to_frame >= 0 + chosen[has_frame] = frame_times[row_to_frame[has_frame]] + if bool((chosen[has_frame] >= times[has_frame]).any()): + raise ValueError(f"Person frame mapping selected future visual evidence: {path}") + frame_sets = FramePersonFeatures( + features, + frame_offsets, + row_to_frame.to(dtype=torch.long), + chosen, + ) + direct_audio = torch.cat( + (audio.float(), audio_valid[:, None].float()), + dim=-1, + ) + return direct_audio, frame_sets, contract + + def local_head_dimensions(contract: dict[str, Any] | None) -> dict[str, int]: if contract is None: return {"local_visual_dim": 0, "local_audio_dim": 0} diff --git a/test_local_turn_features.py b/test_local_turn_features.py index 788bcf5..7334243 100644 --- a/test_local_turn_features.py +++ b/test_local_turn_features.py @@ -28,8 +28,9 @@ from extract_turn_local_features import ( main as extract_main, ) from local_turn_features import ( - LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes, load_local_features, - local_head_dimensions, pooled_recent_visual, reference_alignment, + LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes, + load_frame_set_features, load_local_features, local_head_dimensions, + pooled_recent_visual, reference_alignment, ) from train_turn_head import _load_examples @@ -47,6 +48,17 @@ def fixture(n=20, name="Session_4", dataset="easycom"): "causal_policy": "half-open-trailing-window", "uses_speech_labels": False}} side["visual_valid"][0] = False side["visual"][0] = 0 + frame_times = torch.arange(max(1, (n + 2) // 3), dtype=torch.float64) * .5 + counts = [2 if index % 3 == 0 else (0 if index % 3 == 1 else 1) + for index in range(len(frame_times))] + offsets = torch.tensor([0, *np.cumsum(counts).tolist()], dtype=torch.long) + person_features = torch.randn(int(offsets[-1]), 6) + side["persons"] = { + "features": person_features, + "frame_times": frame_times, + "frame_offsets": offsets, + "records": [{"id": str(index)} for index in range(len(person_features))], + } labels = np.tile(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], dtype=np.float32), (n // 4 + 1, 1))[:n] return reference, side, labels @@ -280,6 +292,148 @@ class LocalTests(unittest.TestCase): with self.assertRaisesRegex(ValueError, "times"): reference_alignment({**reference, "times": reference["times"] + .2}) + def test_frame_set_loader_is_causal_ragged_and_fail_closed(self): + reference, side, _ = fixture() + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "Session_4.pt" + torch.save(side, path) + audio, people, contract = load_frame_set_features( + path, + reference, + sample_id="Session_4", + dataset="easycom", + ) + self.assertEqual(audio.shape, (20, 6)) + self.assertEqual(contract["modalities"], "av") + self.assertEqual(people.dimension, 6) + self.assertTrue( + torch.all( + people.chosen_frame_times[ + torch.isfinite(people.chosen_frame_times) + ] + < reference["times"][ + torch.isfinite(people.chosen_frame_times) + ] + ) + ) + # At 0.2/0.4s, both rows reuse the 0.0s frame with two people. + values, mask = people.padded(torch.tensor([0, 1, 2])) + self.assertEqual(values.shape, (3, 2, 6)) + self.assertEqual(mask.tolist(), [[True, True], [True, True], [False, False]]) + self.assertTrue(torch.equal(values[0], values[1])) + self.assertEqual(torch.count_nonzero(values[2]), 0) + # The frame stamped exactly 1.0s is excluded at the 1.0s cutoff. + self.assertEqual(people.chosen_frame_times[4].item(), .5) + + broken_cases = [] + broken = copy.deepcopy(side) + broken["persons"]["frame_offsets"][-1] -= 1 + broken_cases.append(broken) + broken = copy.deepcopy(side) + broken["persons"]["frame_times"][1] = broken["persons"]["frame_times"][0] + broken_cases.append(broken) + broken = copy.deepcopy(side) + broken["persons"]["records"] = [] + broken_cases.append(broken) + broken = copy.deepcopy(side) + broken["persons"]["features"][0, 0] = float("nan") + broken_cases.append(broken) + for index, broken in enumerate(broken_cases): + torch.save(broken, path) + with self.subTest(index=index), self.assertRaises(ValueError): + load_frame_set_features( + path, + reference, + sample_id="Session_4", + dataset="easycom", + ) + + def test_frame_person_concatenation_selection_and_empty_padding(self): + features = torch.tensor([[1., 2.], [3., 4.]]) + a = FramePersonFeatures( + features, + torch.tensor([0, 2], dtype=torch.long), + torch.tensor([0, -1], dtype=torch.long), + torch.tensor([0., float("nan")], dtype=torch.float64), + ) + b = FramePersonFeatures( + torch.empty(0, 2), + torch.tensor([0, 0], dtype=torch.long), + torch.tensor([0], dtype=torch.long), + torch.tensor([.5], dtype=torch.float64), + ) + combined = FramePersonFeatures.concatenate([a, b]) + self.assertEqual(len(combined), 3) + self.assertEqual(combined.row_to_frame.tolist(), [0, -1, 1]) + selected = combined.select_rows(torch.tensor([2, 0])) + values, mask = selected.padded() + self.assertEqual(mask.tolist(), [[False, False], [True, True]]) + self.assertEqual(torch.count_nonzero(values[0]), 0) + self.assertTrue(torch.equal(values[1], features)) + + def test_frame_set_head_is_permutation_invariant_and_masks_padding(self): + torch.manual_seed(7) + head = TurnPredictionHead( + hidden_size=8, + num_input_layers=4, + layer_fusion="concat-mlp", + local_visual_dim=6, + local_audio_dim=5, + local_projection_dim=3, + local_fusion="frame-set-attention", + person_attention_dim=4, + person_attention_heads=2, + transition_auxiliary=True, + dropout=0, + ).eval() + states = torch.randn(3, 4, 8) + local = torch.randn(3, 6) + local[:, -1] = 1 + people = torch.randn(3, 3, 6) + mask = torch.tensor( + [[True, True, False], [True, False, False], [False, False, False]] + ) + people[~mask] = 0 + logits, transitions = head.forward_with_transition( + states, + local, + people, + mask, + ) + self.assertEqual(logits.shape, (3, 10, 3)) + self.assertEqual(transitions.shape, (3, 10)) + (logits.square().mean() + transitions.square().mean()).backward() + self.assertIsNotNone(head.person_attention.in_proj_weight.grad) + + permutation = torch.tensor([1, 0, 2]) + permuted = head(states, local, people[:, permutation], mask[:, permutation]) + self.assertTrue(torch.allclose(logits, permuted, atol=1e-6, rtol=1e-6)) + changed = people.clone() + changed[~mask] = 1000 + self.assertTrue(torch.allclose(logits, head(states, local, changed, mask))) + with self.assertRaisesRegex(ValueError, "person features"): + head(states, local) + + config = { + "hidden_size": 8, + "num_input_layers": 4, + "layer_fusion": "concat-mlp", + "local_visual_dim": 6, + "local_audio_dim": 5, + "local_projection_dim": 3, + "local_fusion": "frame-set-attention", + "person_attention_dim": 4, + "person_attention_heads": 2, + "transition_auxiliary": True, + "dropout": 0, + } + restored = TurnPredictionHead( + **checkpoint_head_config({"head_config": config}) + ) + restored.load_state_dict(head.state_dict()) + restored.eval() + self.assertTrue(torch.equal(restored(states, local, people, mask), logits)) + def test_head_shapes_gradients_missing_and_roundtrip(self): states = torch.randn(3, 4, 8) for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)): @@ -338,7 +492,8 @@ class LocalTests(unittest.TestCase): np.save(root / "targets/Session_4.npy", targets) data = _load_examples(root / "features", root / "targets", ["Session_4"], "open-stream-replay", "easycom", {"Session_4": {1, 4}}, root / "local", "audio") - self.assertTrue(torch.equal(data[-1][:, :5], side["audio"][[1, 4]])) + self.assertTrue(torch.equal(data[-2][:, :5], side["audio"][[1, 4]])) + self.assertIsNone(data[-1]) self.assertTrue(torch.equal(data[0], reference["states"][[1, 4]])) self.assertTrue(torch.equal(data[1][:, 0], torch.from_numpy(targets[[2, 5]]))) @@ -562,7 +717,7 @@ class RoundTripTests(unittest.TestCase): np.save(root / "targets" / f"{name}.npy", targets) common = ["--feature-root", str(root / "features"), "--target-root", str(root / "targets"), "--device", "cpu", "--no-progress", "--expected-mode", "open-stream-replay"] - for mode in ("baseline", "audio", "visual", "av", "ego4d"): + for mode in ("baseline", "audio", "visual", "av", "frame-set", "ego4d"): checkpoint = root / f"{mode}.pth" train = [sys.executable, str(repo / "train_turn_head.py"), *common, "--epochs", "1", "--warmup-epochs", "0", "--batch-size", "4", @@ -579,8 +734,15 @@ class RoundTripTests(unittest.TestCase): train += ["--dev-sessions", "5", "--local-projection-dim", "0"] if mode == "av": train += ["--transition-objective", "r1-supervised"] + if mode == "frame-set": + train += [ + "--local-fusion", "frame-set-attention", + "--person-attention-dim", "4", + "--person-attention-heads", "2", + ] if mode != "baseline": - train += ["--local-feature-root", str(root / "local"), "--local-modalities", "av" if mode == "ego4d" else mode] + modality = "av" if mode in ("ego4d", "frame-set") else mode + train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality] infer += ["--local-feature-root", str(root / "local")] for command in (train, infer): completed = subprocess.run(command, cwd=repo, env=env, capture_output=True, text=True, timeout=60) @@ -590,6 +752,8 @@ class RoundTripTests(unittest.TestCase): self.assertEqual(bool(metrics["local_contract"]), mode != "baseline") if mode == "av": self.assertTrue(metrics["transition_auxiliary"]) + if mode == "frame-set": + self.assertEqual(metrics["local_fusion"], "frame-set-attention") # The reporting command uses the same metric keys as inference. completed = subprocess.run( [sys.executable, str(repo / "summarize_local_baselines.py"), "--root", str(root)], diff --git a/train_turn_head.py b/train_turn_head.py index 9c850d4..a57a80b 100644 --- a/train_turn_head.py +++ b/train_turn_head.py @@ -24,6 +24,7 @@ from easycom_head import ( CLASS_NAMES, DATASET_NAMES, LAYER_FUSION_MODES, + LOCAL_FUSION_MODES, LOSS_NAMES, NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS, @@ -46,6 +47,7 @@ from easycom_head import ( transition_targets_from_activity, validate_disjoint_sessions, ) +from local_turn_features import FramePersonFeatures def parse_args() -> argparse.Namespace: @@ -58,10 +60,18 @@ def parse_args() -> argparse.Namespace: help="Optional frozen local-feature sidecars; omit for the original baseline.", ) parser.add_argument("--local-modalities", choices=("visual", "audio", "av"), default="av") + parser.add_argument( + "--local-fusion", + choices=LOCAL_FUSION_MODES, + default="pooled", + help="pooled keeps the existing baseline; frame-set-attention queries the latest causal person set.", + ) parser.add_argument( "--local-projection-dim", type=int, default=256, help="Trainable projection per local modality; 0 retains encoder dimensions.", ) + parser.add_argument("--person-attention-dim", type=int, default=256) + parser.add_argument("--person-attention-heads", type=int, default=4) parser.add_argument( "--dataset", choices=DATASET_NAMES, @@ -165,6 +175,17 @@ def parse_args() -> argparse.Namespace: def _validate_args(args: argparse.Namespace) -> None: if args.local_projection_dim < 0: raise ValueError("--local-projection-dim must be non-negative") + if args.person_attention_dim <= 0 or args.person_attention_heads <= 0: + raise ValueError("Person attention dimension/heads must be positive") + if args.person_attention_dim % args.person_attention_heads: + raise ValueError("--person-attention-dim must be divisible by --person-attention-heads") + if args.local_fusion == "frame-set-attention" and ( + args.local_feature_root is None or args.local_modalities != "av" + ): + raise ValueError( + "--local-fusion frame-set-attention requires --local-feature-root " + "and --local-modalities av" + ) if args.epochs <= 0: raise ValueError("--epochs must be positive") if args.batch_size <= 0: @@ -225,6 +246,7 @@ def _load_examples( local_modalities: str = "av", local_contract: dict[str, Any] | None = None, compute_transition_targets: bool = True, + local_fusion: str = "pooled", ) -> tuple[ torch.Tensor, torch.Tensor, @@ -233,11 +255,13 @@ def _load_examples( int, list[dict[str, Any]], torch.Tensor, + FramePersonFeatures | None, ]: all_states: list[torch.Tensor] = [] all_targets: list[torch.Tensor] = [] all_transition_targets: list[torch.Tensor] = [] all_local: list[torch.Tensor] = [] + all_people: list[FramePersonFeatures] = [] metadata: list[dict[str, Any]] = [] hidden_size: int | None = None num_input_layers: int | None = None @@ -251,6 +275,7 @@ def _load_examples( local_feature_root=local_feature_root, local_modalities=local_modalities, local_contract=local_contract, + local_fusion=local_fusion, ) local_contract = data.feature_metadata.get("local_contract") states, targets, future_indices = data.future_examples(NUM_HORIZONS) @@ -260,6 +285,11 @@ def _load_examples( if data.local_features is not None else torch.empty(len(states), 0) ) + people = ( + data.local_person_features.select_rows(current_indices) + if data.local_person_features is not None + else None + ) if compute_transition_targets: activity_sequence = torch.cat( (data.targets[current_indices, None, :], targets), @@ -286,6 +316,8 @@ def _load_examples( targets = targets[keep] transition_targets = transition_targets[keep] local = local[keep] + if people is not None: + people = people.select_rows(torch.nonzero(keep).flatten()) if not len(states): raise ValueError(f"Sparse anchor selection kept no examples for {name}") if hidden_size is None: @@ -306,6 +338,8 @@ def _load_examples( all_targets.append(targets) all_transition_targets.append(transition_targets) all_local.append(local) + if people is not None: + all_people.append(people) metadata.append( { "session": name, @@ -332,6 +366,8 @@ def _load_examples( ) if hidden_size is None or num_input_layers is None: raise ValueError("No sessions selected") + if all_people and len(all_people) != len(all_states): + raise ValueError("Person sets are missing for some sessions") return ( torch.cat(all_states, dim=0), torch.cat(all_targets, dim=0), @@ -340,6 +376,7 @@ def _load_examples( num_input_layers, metadata, torch.cat(all_local, dim=0), + FramePersonFeatures.concatenate(all_people) if all_people else None, ) @@ -407,6 +444,7 @@ def _evaluate_valid_horizons( batch_size: int, loss_name: str, local: torch.Tensor | None = None, + people: FramePersonFeatures | None = None, ) -> dict[str, Any]: head.eval() chunks: list[torch.Tensor] = [] @@ -420,7 +458,19 @@ def _evaluate_valid_horizons( local[start : start + batch_size].to(device=device, dtype=torch.float32) if local is not None else None ) - chunks.append(marginal_probabilities_from_logits(head(x, side), loss_name).cpu()) + person_batch = None + person_mask = None + if people is not None: + rows = torch.arange(start, min(start + batch_size, len(states))) + person_batch, person_mask = people.padded(rows) + person_batch = person_batch.to(device=device, dtype=torch.float32) + person_mask = person_mask.to(device=device) + chunks.append( + marginal_probabilities_from_logits( + head(x, side, person_batch, person_mask), + loss_name, + ).cpu() + ) probabilities = torch.cat(chunks, dim=0).numpy() target_array = targets.numpy() @@ -455,7 +505,7 @@ def _checkpoint( history: list[dict[str, Any]], ) -> dict[str, Any]: return { - "format_version": 4, + "format_version": 5, "head": head.state_dict(), "head_config": { "hidden_size": hidden_size, @@ -469,6 +519,9 @@ def _checkpoint( "local_visual_dim": head.local_visual_dim, "local_audio_dim": head.local_audio_dim, "local_projection_dim": head.local_projection_dim, + "local_fusion": head.local_fusion, + "person_attention_dim": head.person_attention_dim, + "person_attention_heads": head.person_attention_heads, }, "local_contract": feature_metadata[0].get("local_contract"), "alignment": { @@ -513,7 +566,10 @@ def _checkpoint( "layer_fusion": args.layer_fusion, "local_feature_root": str(args.local_feature_root) if args.local_feature_root else None, "local_modalities": args.local_modalities if args.local_feature_root else None, + "local_fusion": args.local_fusion, "local_projection_dim": args.local_projection_dim, + "person_attention_dim": args.person_attention_dim, + "person_attention_heads": args.person_attention_heads, "learnable_layer_mixture": args.layer_fusion == "scalar-mixture", "supervision_policy": ( "transition-sparse-manifest" @@ -605,6 +661,7 @@ def main() -> None: num_input_layers, feature_metadata, local_features, + person_features, ) = _load_examples( args.feature_root, args.target_root, @@ -615,6 +672,7 @@ def main() -> None: args.local_feature_root, args.local_modalities, compute_transition_targets=args.transition_objective != "none", + local_fusion=args.local_fusion, ) local_contract = feature_metadata[0].get("local_contract") if not len(states): @@ -645,6 +703,7 @@ def main() -> None: dev_states: torch.Tensor | None = None dev_targets: torch.Tensor | None = None dev_local: torch.Tensor | None = None + dev_people: FramePersonFeatures | None = None if dev_sessions: ( dev_states, @@ -654,6 +713,7 @@ def main() -> None: dev_num_input_layers, _, dev_local, + dev_people, ) = _load_examples( args.feature_root, args.target_root, @@ -665,6 +725,7 @@ def main() -> None: args.local_modalities, local_contract, compute_transition_targets=args.transition_objective != "none", + local_fusion=args.local_fusion, ) if dev_hidden_size != hidden_size: raise ValueError( @@ -678,7 +739,13 @@ def main() -> None: ) print(f"development examples: {len(dev_states)}") - dataset = TensorDataset(states, targets, transition_targets, local_features) + dataset = TensorDataset( + states, + targets, + transition_targets, + local_features, + torch.arange(len(states), dtype=torch.long), + ) generator = torch.Generator().manual_seed(args.seed) loader = DataLoader( dataset, @@ -702,13 +769,26 @@ def main() -> None: layer_fusion=args.layer_fusion, transition_auxiliary=args.transition_objective != "none", local_projection_dim=args.local_projection_dim, + local_fusion=args.local_fusion, + person_attention_dim=args.person_attention_dim, + person_attention_heads=args.person_attention_heads, **local_head_dimensions(local_contract), ).to(device) if local_contract is not None: print( - f"local fusion: {args.local_modalities}; raw V={head.local_visual_dim}, " - f"A={head.local_audio_dim}; projection={args.local_projection_dim or 'identity'}" + f"local fusion: {args.local_fusion}/{args.local_modalities}; " + f"raw V={head.local_visual_dim}, A={head.local_audio_dim}; " + f"projection={args.local_projection_dim or 'identity'}" ) + if person_features is not None: + print( + "frame sets: " + f"{len(person_features.frame_offsets) - 1} source frames, " + f"{len(person_features.features)} person observations, " + f"max_people={person_features.max_people}, " + f"attention={args.person_attention_dim}D/" + f"{args.person_attention_heads} heads" + ) trainable_parameters = sum( parameter.numel() for parameter in head.parameters() if parameter.requires_grad ) @@ -782,7 +862,13 @@ def main() -> None: unit="batch", disable=args.no_progress, ) - for state_batch, target_batch, transition_target_batch, local_batch in progress: + for ( + state_batch, + target_batch, + transition_target_batch, + local_batch, + example_index_batch, + ) in progress: local_batch = local_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) @@ -801,15 +887,37 @@ def main() -> None: dtype=torch.float32, non_blocking=True, ) + person_batch = None + person_mask = None + if person_features is not None: + person_batch, person_mask = person_features.padded( + example_index_batch + ) + person_batch = person_batch.to( + device=device, + dtype=torch.float32, + non_blocking=True, + ) + person_mask = person_mask.to(device=device, non_blocking=True) optimizer.zero_grad(set_to_none=True) if args.transition_objective == "none": - logits = head(state_batch, local_batch) + logits = head( + state_batch, + local_batch, + person_batch, + person_mask, + ) activity_loss = loss_function(logits, target_batch) loss = activity_loss transition_loss = activity_loss.detach().new_zeros(()) expected_reward = activity_loss.detach().new_zeros(()) else: - logits, transition_logits = head.forward_with_transition(state_batch, local_batch) + logits, transition_logits = head.forward_with_transition( + state_batch, + local_batch, + person_batch, + person_mask, + ) activity_loss = loss_function(logits, target_batch) components = transition_aware_objective( activity_loss, @@ -890,6 +998,7 @@ def main() -> None: args.batch_size, args.loss, dev_local, + dev_people, ) dev_map = float(dev_metrics["mean_map"]) record["dev_valid_mean_map"] = dev_map