wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add causal frame set attention fusion

Browse Source
This commit is contained in:
Xin Wang
2026-09-02 18:48:26 +08:00
parent e769094b3d
commit 8e70281330
8 changed files with 837 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
PROJECT_CONTEXT.md
View File

@@ -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

21
docs/local_feature_baseline.md
View File

@@ -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 批量提取、

11
docs/progressive_feature_fusion_route_2026-09-02.md
View File

@@ -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 必跑对照

226
easycom_head.py
View File

@@ -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",
(

33
infer_turn_head.py
View File

@@ -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(

278
local_turn_features.py
View File

@@ -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}

174
test_local_turn_features.py
View File

@@ -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)],

125
train_turn_head.py
View File

@@ -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