diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index 0ebfb51..d70f603 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -86,6 +86,12 @@ All are pending in that log; Omni LoRA is deferred and no semantic gate is compl per cap. CPU tests match separate single-window replay and both head/local loaders; real-GPU replay and media smoke remain required. See [`docs/easycom_multiwindow_features.md`](docs/easycom_multiwindow_features.md). +- Two no-box full-frame Qwen AV screens are implemented on top of the matched + global-Omni + recent-Qwen-audio inputs. F2 queries the latest strictly + pre-cutoff post-merger spatial-token grid; F3 first applies a shared causal + GRU to the latest same-grid-position histories. They use no face boxes or + person IDs and therefore test full-frame visual context, not speaker + grounding. See [`docs/easycom_fullframe_qwen_av.md`](docs/easycom_fullframe_qwen_av.md). - Qwen2.5-Omni now has two matched chunk-handling controls around the existing 5 Hz partial-tail path: native completed 2-second blocks held between boundaries, and non-native 200-ms payload commits. Both preserve dense causal diff --git a/docs/README.md b/docs/README.md index 354cd29..e555c45 100644 --- a/docs/README.md +++ b/docs/README.md @@ -26,6 +26,7 @@ tests. | [`easycom_multiwindow_features.md`](easycom_multiwindow_features.md) | Shared-AV extraction of independent Qwen 2/4/6/8/10/12-second context banks, GPU smoke and matched head training | | [`local_feature_baseline.md`](local_feature_baseline.md) | Extracting/validating EasyCom pooled local Qwen visual and audio sidecars | | [`easycom_encoder_only.md`](easycom_encoder_only.md) | Comparing EasyCom Qwen encoder-only A/V/AV representations | +| [`easycom_fullframe_qwen_av.md`](easycom_fullframe_qwen_av.md) | No-box F2/F3 full-frame Qwen spatial and causal spatiotemporal fusion | | [`ego4d_encoder_only.md`](ego4d_encoder_only.md) | Preparing and running the Ego4D encoder-only comparison | | [`ego4d_local_baselines.md`](ego4d_local_baselines.md) | Previewing Ego4D boxes, extracting local features and running 8-GPU jobs | | [`easycom_track_gru_residual.md`](easycom_track_gru_residual.md) | Matched pooled-to-track GRU experiment, causal identity contract and L20 commands | diff --git a/docs/easycom_fullframe_qwen_av.md b/docs/easycom_fullframe_qwen_av.md new file mode 100644 index 0000000..d734e4b --- /dev/null +++ b/docs/easycom_fullframe_qwen_av.md @@ -0,0 +1,143 @@ +# EasyCom no-box full-frame Qwen AV: F2/F3 + +## Definition + +This experiment adds full-frame Qwen vision-encoder evidence to the existing +strictly causal global Omni plus recent Qwen-audio baseline. It never reads +face/head boxes, track IDs, active-speaker labels or speech labels. + +```text +F2: G(t) + A(t) query latest full-frame Qwen spatial tokens before t + +F3: same-grid-position tokens from the latest causal frames + -> shared chronological GRU + -> G(t) + A(t) queries the temporally encoded spatial grid +``` + +At the default 2 FPS, F3 keeps two observations per grid position: at most the +latest one second of visual history. `G` already contains global audio/video +context, so this tests whether an explicit recent encoder readout adds timely +evidence. It is not an audio-free Omni baseline. + +F2/F3 are no-box full-frame visual-context models. Improvement does not by +itself establish speaker grounding, identity or addressee reasoning. + +## Causal contract + +`extract_fullframe_qwen_av_features.py` creates +`qwen-fullframe-av-tokens-v1` sidecars: + +- a cutoff `t` sees only fixed-grid frames with timestamp strictly less than + `t`; a frame exactly at `t` is excluded; +- images are complete EasyCom camera frames after the standard session-banner + crop; +- Qwen post-merger spatial tokens are stored without mean pooling; +- F3 links only the same merged-grid row/column within one source clip, and + resets on clip boundaries, grid changes or excessive time gaps; +- audio is copied from the existing audited `[t-1,t)` Qwen mixed-audio + sidecar; its person-crop tensors are never consumed; +- visual/audio model identity, config digest and Transformers version must + match exactly; +- partial smoke sidecars are rejected by training. + +The targets remain the unchanged 5 Hz `background/wearer/others` labels and +the ten 0.2--2.0-second horizons. + +## 1. Four-second GPU smoke extraction + +```bash +cd /data1/yfzx/wangx/Code/icassp2027-experiment + +DATA_ROOT=/data1/yfzx/wangx/Datasets/easycom/extracted_easycom +MODEL=/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B +FEATURE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack +TARGET_ROOT=/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame +AUDIO_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576 +SMOKE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/fullframe_qwen_av_smoke + +CUDA_VISIBLE_DEVICES=0 python extract_fullframe_qwen_av_features.py \ + --data-root "$DATA_ROOT" --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" --audio-feature-root "$AUDIO_ROOT" \ + --output-root "$SMOKE_ROOT" --samples 4 --max-steps 20 \ + --model-id "$MODEL" --device cuda:0 --dtype bfloat16 \ + --flash-attention-2 --video-fps 2 --max-pixels 50176 \ + --image-batch-size 8 +``` + +Expected: 20 states spanning 0.2--4.0 s; eight visual frames at +`0.0,...,3.5`; finite nonempty spatial tokens; matched 2048D audio; and +`partial=True`. Never train on this smoke file. + +If the audio sidecars were made with an explicit `--encoder-id`, pass that +exact value here. Otherwise omit it, as above, so both branches use the same +recorded local checkpoint path. + +Run the regression tests in the remote environment: + +```bash +python -m unittest -v test_fullframe_qwen_av test_local_turn_features +``` + +They verify strict boundary exclusion, same-grid F3 history, malformed and +partial-contract rejection, and F2/F3 forward/backward behavior. + +## 2. Extract all sessions + +Use a fresh root. Existing files are not overwritten unless `--overwrite` is +explicitly provided. + +```bash +FULLFRAME_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/fullframe_qwen_av_2fps_224 + +CUDA_VISIBLE_DEVICES=0 python extract_fullframe_qwen_av_features.py \ + --data-root "$DATA_ROOT" --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" --audio-feature-root "$AUDIO_ROOT" \ + --output-root "$FULLFRAME_ROOT" \ + --samples 1 2 3 4 5 6 7 8 9 10 11 12 \ + --model-id "$MODEL" --device cuda:0 --dtype bfloat16 \ + --flash-attention-2 --video-fps 2 --max-pixels 50176 \ + --image-batch-size 8 +``` + +On interruption, rerun only missing session numbers. Do not mix checkpoints, +Transformers versions, FPS settings or pixel budgets in one root. + +## 3. Train and evaluate + +The launcher uses Sessions 4--12 for four fixed epochs and evaluates Sessions +3/2/1 with the existing LSTR-MCE ten-horizon protocol. + +```bash +export FEATURE_ROOT TARGET_ROOT FULLFRAME_ROOT +export CHECKPOINT_ROOT=checkpoints/easycom_fullframe_qwen_av +export RESULT_ROOT=results/easycom_fullframe_qwen_av + +CUDA_VISIBLE_DEVICES=0 bash run_easycom_fullframe_qwen_av.sh f2 0 +CUDA_VISIBLE_DEVICES=1 bash run_easycom_fullframe_qwen_av.sh f3 0 +``` + +After seed 0 passes shape/loss/inference checks: + +```bash +CUDA_VISIBLE_DEVICES=0 bash run_easycom_fullframe_qwen_av.sh f2 1 2 +CUDA_VISIBLE_DEVICES=1 bash run_easycom_fullframe_qwen_av.sh f3 1 2 +``` + +Defaults are train batch 64, inference batch 128, F3 history 2 frames and +maximum inter-frame link gap 0.75 s. Lower `TRAIN_BATCH_SIZE` or +`INFER_BATCH_SIZE` if needed; do not change the feature contract. + +The checkpoint-compatible option remains named `--person-history-frames`; in +F3 it means observations per spatial grid position, not person history. + +## Required comparison and claim boundary + +Compare matched seeds for `G+A`, F2 and F3. Add a past time-shifted full-frame +control before attributing gains to AV synchronization. F3 contains a GRU that +F2 does not, so F3–F2 combines extra history and temporal capacity; a +parameter-matched audio/query temporal control is required before claiming a +visual-motion mechanism. + +Dense AP/mAP is an architecture screen. Floor-entry/yield interpretations +remain timing-derived proxies until the competing-risk real-data audit is +frozen. diff --git a/easycom_head.py b/easycom_head.py index 0683e1e..8297941 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -67,6 +67,13 @@ CURRENT_FRAME_PERSON_SET_FUSION_MODES = ( "audio-person-frame-set", "gated-person-state-adapter", "highrate-person-av-attention", + "fullframe-spatial-attention", +) +FULLFRAME_SPATIAL_FUSION_MODES = ("fullframe-spatial-attention",) +FULLFRAME_TEMPORAL_FUSION_MODES = ("fullframe-spatiotemporal-attention",) +FULLFRAME_FUSION_MODES = ( + *FULLFRAME_SPATIAL_FUSION_MODES, + *FULLFRAME_TEMPORAL_FUSION_MODES, ) HIGHRATE_PERSON_AV_FUSION_MODES = ("highrate-person-av-attention",) TRACK_PERSON_FUSION_MODES = ( @@ -74,12 +81,17 @@ TRACK_PERSON_FUSION_MODES = ( "causal-person-adapter-v2", "causal-person-adapter-egospeak-gru", ) +HISTORY_VISUAL_SET_FUSION_MODES = ( + *TRACK_PERSON_FUSION_MODES, + *FULLFRAME_TEMPORAL_FUSION_MODES, +) PERSON_ADAPTER_V2_MODES = ( "causal-person-adapter-v2", "causal-person-adapter-egospeak-gru", ) PERSON_SET_FUSION_MODES = ( *CURRENT_FRAME_PERSON_SET_FUSION_MODES, + *FULLFRAME_TEMPORAL_FUSION_MODES, *TRACK_PERSON_FUSION_MODES, ) LOCAL_FUSION_MODES = ("pooled", *PERSON_SET_FUSION_MODES) @@ -92,6 +104,12 @@ LOCAL_FUSION_SCOPES = { "highrate-person-av-attention": ( "global-omni-traditional-audio-base-plus-zero-init-25hz-person-residual" ), + "fullframe-spatial-attention": ( + "f2-global-omni-qwen-audio-latest-fullframe-spatial-tokens" + ), + "fullframe-spatiotemporal-attention": ( + "f3-global-omni-qwen-audio-causal-fullframe-token-history" + ), "track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual", "causal-person-adapter-v2": ( "causal-anonymous-track-gru-dynamic-gated-logit-residual" @@ -206,7 +224,10 @@ class TurnPredictionHead(nn.Module): fused_size = hidden_size projected_modalities = ( (("audio", local_audio_dim),) - if local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES + if local_fusion in ( + *CURRENT_FRAME_PERSON_SET_FUSION_MODES, + *FULLFRAME_TEMPORAL_FUSION_MODES, + ) else (("visual", local_visual_dim), ("audio", local_audio_dim)) ) for name, dimension in projected_modalities: @@ -291,6 +312,13 @@ class TurnPredictionHead(nn.Module): batch_first=True, ) self.audio_person_gate = nn.Parameter(torch.zeros(())) + if local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES: + self.person_track_gru = nn.GRU( + person_attention_dim, + person_attention_dim, + num_layers=1, + batch_first=True, + ) if local_fusion == "gated-person-state-adapter": # Screening version of F6 for precomputed frozen states. It # injects selected person evidence after the Thinker readout, @@ -303,6 +331,8 @@ class TurnPredictionHead(nn.Module): nn.Linear(person_attention_dim, hidden_size), ) self.person_state_adapter_gate = nn.Parameter(torch.zeros(())) + elif local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES: + fused_size += person_attention_dim + 1 elif local_fusion in TRACK_PERSON_FUSION_MODES: # Apply one shared causal GRU independently to every anonymous # face track. The regular pooled G+A+V head remains the exact @@ -419,7 +449,10 @@ class TurnPredictionHead(nn.Module): @property def local_input_size(self) -> int: - if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES: + if self.local_fusion in ( + *CURRENT_FRAME_PERSON_SET_FUSION_MODES, + *FULLFRAME_TEMPORAL_FUSION_MODES, + ): return self.local_audio_dim + 1 return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d) @@ -567,6 +600,65 @@ class TurnPredictionHead(nn.Module): audio_valid, ) + def _fullframe_temporal_context( + self, + states: torch.Tensor, + local: torch.Tensor, + token_histories: torch.Tensor | None, + history_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Encode same-grid-position visual histories, then query the grid.""" + + if states.ndim != 2: + raise ValueError("F3 expects batched [B,D] global states") + if local.shape != (len(states), self.local_audio_dim + 1): + raise ValueError("F3 requires aligned [Qwen audio,valid] features") + if token_histories is None or token_histories.ndim != 4: + raise ValueError("F3 requires token histories [B,K,L,visual_dim]") + if ( + token_histories.shape[0] != len(states) + or token_histories.shape[2] != self.person_history_frames + or token_histories.shape[3] != self.local_visual_dim + ): + raise ValueError( + "F3 history shape does not match its configured length/dimension" + ) + 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 + token_states, token_valid = self._encode_person_track_histories( + token_histories, + history_mask, + ) + audio_values = local[:, : self.local_audio_dim] + audio_valid = local[:, self.local_audio_dim :] + audio_projected = self.local_projections["audio"](audio_values) * audio_valid + query = self.person_state_query(states) + query = query + self.person_audio_query(audio_values) * audio_valid + memory = torch.cat( + (self.null_person.expand(len(states), -1, -1), token_states), + dim=1, + ) + memory_valid = torch.cat( + ( + torch.ones( + (len(states), 1), dtype=torch.bool, device=states.device + ), + token_valid, + ), + dim=1, + ) + context, _ = self.person_attention( + query[:, None, :], + memory, + memory, + key_padding_mask=~memory_valid, + need_weights=False, + ) + visual_valid = token_valid.any(dim=1, keepdim=True).to(states.dtype) + return audio_projected, audio_valid, context[:, 0, :], visual_valid + def _encode_person_track_histories( self, person_histories: torch.Tensor, @@ -685,6 +777,21 @@ class TurnPredictionHead(nn.Module): (states, audio, audio_valid, person_context, person_valid), dim=-1, ) + elif self.local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES: + if local is None: + raise ValueError("F3 requires local Qwen audio features") + audio, audio_valid, visual_context, visual_valid = ( + self._fullframe_temporal_context( + states, + local, + person_features, + person_mask, + ) + ) + states = torch.cat( + (states, audio, audio_valid, visual_context, visual_valid), + dim=-1, + ) elif self.local_input_size: if local is None or local.shape != (*states.shape[:-1], self.local_input_size): raise ValueError( @@ -1554,7 +1661,26 @@ def load_session_data( name, dataset=dataset, ) - if local_fusion in TRACK_PERSON_FUSION_MODES: + if local_fusion in FULLFRAME_FUSION_MODES: + if dataset != "easycom" or local_modalities != "av": + raise ValueError( + "Full-frame F2/F3 require EasyCom and --local-modalities av" + ) + from fullframe_qwen_av import load_fullframe_av_features + + local, local_people, contract = load_fullframe_av_features( + local_path, + payload, + sample_id=name, + dataset=dataset, + expected_contract=local_contract, + history_max_gap_seconds=( + person_track_max_gap_seconds + if local_fusion in FULLFRAME_TEMPORAL_FUSION_MODES + else None + ), + ) + elif local_fusion in TRACK_PERSON_FUSION_MODES: if local_modalities != "av": raise ValueError( "Track person fusion requires --local-modalities av" @@ -1610,7 +1736,7 @@ def load_session_data( elif local_contract is not None: raise ValueError("--local-feature-root is required for this checkpoint") elif local_fusion != "pooled": - raise ValueError("Person-set fusion requires --local-feature-root") + raise ValueError("Set/token fusion requires --local-feature-root") return SessionData( name, states, diff --git a/extract_fullframe_qwen_av_features.py b/extract_fullframe_qwen_av_features.py new file mode 100644 index 0000000..21b17ef --- /dev/null +++ b/extract_fullframe_qwen_av_features.py @@ -0,0 +1,402 @@ +"""Extract causal full-frame Qwen vision tokens plus matched Qwen audio. + +This EasyCom-only extractor never reads face/head boxes or speaker labels. It +reuses the audited recent mixed-audio vectors from a local-Qwen sidecar and +encodes complete camera frames on a fixed causal grid. The output supports F2 +latest-frame spatial attention and F3 same-grid-position temporal modeling. +""" + +from __future__ import annotations + +import argparse +import os +import tempfile +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from tqdm.auto import tqdm + +from easycom_head import normalize_sample_name, resolve_feature_path +from extract_turn_local_features import ( + model_config_digest, + validate_local_transformers_version, + validate_source_timeline, +) +from fullframe_qwen_av import FULLFRAME_AV_FORMAT +from local_turn_features import causal_frame_ticks, load_local_features, reference_alignment + + +class FrozenQwenFullFrameEncoder: + """Keep only the frozen Qwen vision tower and merger on the GPU.""" + + def __init__(self, args: argparse.Namespace) -> None: + import transformers + from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration + + self.transformers_version = validate_local_transformers_version(transformers.__version__) + self.device = torch.device(args.device) + if self.device.type == "cuda" and not torch.cuda.is_available(): + raise RuntimeError(f"CUDA requested but unavailable: {args.device}") + dtype_key = "dtype" if self.transformers_version == "4.57.6" else "torch_dtype" + kwargs = { + dtype_key: getattr(torch, args.dtype), + "device_map": {"": "cpu"}, + "output_loading_info": True, + } + if args.flash_attention_2: + kwargs["attn_implementation"] = "flash_attention_2" + model, info = Qwen2_5OmniThinkerForConditionalGeneration.from_pretrained( + args.model_id, **kwargs + ) + bad = [key for key in info.get("missing_keys", []) if key.startswith("visual.")] + if bad or info.get("mismatched_keys") or info.get("error_msgs"): + raise RuntimeError(f"Incomplete Qwen vision weights: {info}") + model.requires_grad_(False).eval() + self.visual_dim = int(model.config.vision_config.out_hidden_size) + self.model_config_sha256 = model_config_digest(model.config.to_dict()) + model.model = None + model.lm_head = None + model.audio_tower = None + model.visual.to(self.device) + self.model = model + self.processor = Qwen2_5OmniProcessor.from_pretrained( + args.model_id, use_fast=False + ) + self.max_pixels = args.max_pixels + self.image_batch_size = args.image_batch_size + + @torch.inference_mode() + def image_tokens( + self, images: list[Any] + ) -> tuple[list[torch.Tensor], list[tuple[int, int]]]: + pieces: list[torch.Tensor] = [] + grids: list[tuple[int, int]] = [] + for start in range(0, len(images), self.image_batch_size): + batch = self.processor.image_processor( + images=images[start : start + self.image_batch_size], + min_pixels=4 * 28 * 28, + max_pixels=self.max_pixels, + return_tensors="pt", + ) + grid = batch["image_grid_thw"].to(self.device) + pixels = batch["pixel_values"].to( + self.device, dtype=self.model.visual.dtype + ) + tokens = self.model.get_image_features(pixels, grid) + merge = int(self.model.config.vision_config.spatial_merge_size) + counts = (grid.prod(-1) // (merge * merge)).tolist() + if not isinstance(tokens, torch.Tensor) or tokens.shape != ( + sum(counts), self.visual_dim + ): + raise RuntimeError(f"Unexpected Qwen image output: {getattr(tokens, 'shape', None)}") + pieces.extend(part.float().cpu() for part in tokens.split(counts)) + for row in grid.tolist(): + temporal, height, width = (int(value) for value in row) + if temporal != 1 or height % merge or width % merge: + raise RuntimeError(f"Unexpected image grid before merger: {row}") + grids.append((height // merge, width // merge)) + if len(pieces) != len(images) or len(grids) != len(images): + raise AssertionError("Qwen did not return one token set per image") + return pieces, grids + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--data-root", type=Path, required=True) + parser.add_argument("--feature-root", type=Path, required=True) + parser.add_argument("--target-root", type=Path, required=True) + parser.add_argument("--audio-feature-root", type=Path, required=True) + parser.add_argument("--output-root", type=Path, required=True) + parser.add_argument("--samples", nargs="+", required=True) + parser.add_argument("--model-id", default="Qwen/Qwen2.5-Omni-3B") + parser.add_argument("--encoder-id") + parser.add_argument("--device", default="cuda:0") + parser.add_argument( + "--dtype", choices=("float32", "float16", "bfloat16"), default="bfloat16" + ) + parser.add_argument("--flash-attention-2", action="store_true") + parser.add_argument("--video-fps", type=float, default=2.0) + parser.add_argument("--max-pixels", type=int, default=224 * 224) + parser.add_argument("--image-batch-size", type=int, default=8) + parser.add_argument( + "--decoder-backend", + choices=("decord", "torchcodec", "pynvcodec"), + default="decord", + ) + parser.add_argument("--decoder-threads", type=int, default=1) + parser.add_argument("--decoder-gpu-id", type=int, default=0) + parser.add_argument("--max-steps", type=int) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--overwrite", action="store_true") + parser.add_argument("--no-progress", action="store_true") + args = parser.parse_args(argv) + if not np.isfinite(args.video_fps) or args.video_fps <= 0: + parser.error("--video-fps must be finite and positive") + if args.max_pixels < 4 * 28 * 28 or args.image_batch_size <= 0: + parser.error("Require max-pixels >= 3136 and positive image-batch-size") + if args.max_steps is not None and args.max_steps <= 0: + parser.error("--max-steps must be positive") + protected = { + args.feature_root.resolve(), + args.audio_feature_root.resolve(), + args.target_root.resolve(), + } + if args.output_root.resolve() in protected: + parser.error("--output-root must not overwrite an input root") + return args + + +def _build_media(args: argparse.Namespace, sample: str, end: float): + from streaming_turn_model_directly import CausalSourceAV, _easycom_directory_specs + + specs = _easycom_directory_specs( + args.data_root / "Main" / "Video_Compressed" / sample + ) + return CausalSourceAV( + specs, + output_fps=args.video_fps, + required_duration=end, + decoder_backend=args.decoder_backend, + decoder_threads=args.decoder_threads, + decoder_gpu_id=args.decoder_gpu_id, + ) + + +def _requested_config( + args: argparse.Namespace, + encoder: FrozenQwenFullFrameEncoder, + audio_contract: dict[str, Any], +) -> dict[str, Any]: + source = audio_contract["config"] + if source.get("causal_policy") != "half-open-trailing-window": + raise ValueError("Source Qwen audio is not a half-open causal window") + if source.get("uses_speech_labels") is not False: + raise ValueError("Source Qwen audio must not use speech labels") + audio_dim = source.get("audio_dim") + if not isinstance(audio_dim, int) or audio_dim <= 0: + raise ValueError("Source Qwen audio contract lacks audio_dim") + encoder_id = args.encoder_id or args.model_id + if source.get("encoder_id") != encoder_id: + raise ValueError( + f"Vision/audio encoder identities differ: {encoder_id!r} vs " + f"{source.get('encoder_id')!r}" + ) + if source.get("model_config_sha256") != encoder.model_config_sha256: + raise ValueError("Vision/audio Qwen model configs differ") + if source.get("transformers") != encoder.transformers_version: + raise ValueError("Vision/audio Transformers versions differ") + return { + "encoder_id": encoder_id, + "model_config_sha256": encoder.model_config_sha256, + "transformers": encoder.transformers_version, + "dtype": args.dtype, + "attention": "flash_attention_2" if args.flash_attention_2 else "default", + "visual_dim": encoder.visual_dim, + "audio_dim": audio_dim, + "visual_fps": float(args.video_fps), + "max_pixels": int(args.max_pixels), + "visual_feature_kind": "qwen-vision-post-merger-spatial-token", + "visual_source": "easycom-full-frame-after-session-banner-crop", + "visual_pooling": "none-ragged-spatial-tokens", + "token_grid_identity": "merged-grid-row-column-within-source-clip", + "audio_source": "reused-qwen-local-mixed-audio", + "audio_window_seconds": source.get("audio_window_seconds"), + "source_audio_contract": audio_contract, + "uses_boxes": False, + "uses_speech_labels": False, + "causal_policy": "strict-half-open-fullframe-grid", + } + + +def _extract_visual_tokens( + media, + encoder: FrozenQwenFullFrameEncoder, + *, + end: float, + fps: float, + batch_size: int, + no_progress: bool, +) -> dict[str, Any]: + ticks = list(causal_frame_ticks(end, fps)) + frame_times = torch.tensor(ticks, dtype=torch.float64) / fps + feature_parts: list[torch.Tensor] = [] + grid_shapes: list[tuple[int, int]] = [] + source_clip_indices: list[int] = [] + offsets = [0] + progress = tqdm( + range(0, len(ticks), batch_size), + desc="full-frame Qwen vision", + unit="batch", + disable=no_progress, + ) + for start in progress: + batch_ticks = ticks[start : start + batch_size] + images, sources = [], [] + for tick in batch_ticks: + time = tick / fps + source = media._source_frame_for_time(time) + images.append(media.frame_at_or_before(time)) + sources.append(source) + token_sets, grids = encoder.image_tokens(images) + for tick, tokens, grid, source in zip(batch_ticks, token_sets, grids, sources): + if tokens.shape != (grid[0] * grid[1], encoder.visual_dim): + raise RuntimeError( + f"Qwen token/grid mismatch at tick {tick}: " + f"{tuple(tokens.shape)} vs {grid}" + ) + if not torch.isfinite(tokens).all(): + raise RuntimeError(f"Non-finite full-frame tokens at tick {tick}") + stored = tokens.to(torch.bfloat16) + feature_parts.append(stored) + grid_shapes.append(grid) + source_clip_indices.append(int(source[0])) + offsets.append(offsets[-1] + len(stored)) + features = ( + torch.cat(feature_parts, dim=0) + if feature_parts + else torch.empty(0, encoder.visual_dim, dtype=torch.bfloat16) + ) + return { + "features": features, + "frame_times": frame_times, + "frame_offsets": torch.tensor(offsets, dtype=torch.long), + "grid_shapes": torch.tensor(grid_shapes, dtype=torch.long), + "frame_source_clip_indices": torch.tensor( + source_clip_indices, dtype=torch.long + ), + } + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + samples = [normalize_sample_name(value, "easycom") for value in args.samples] + if len(samples) != len(set(samples)): + raise ValueError("Duplicate sessions") + references, audio_paths = {}, {} + for sample in samples: + references[sample] = resolve_feature_path(args.feature_root, sample, dataset="easycom") + audio_paths[sample] = resolve_feature_path( + args.audio_feature_root, sample, dataset="easycom" + ) + for path in ( + references[sample], + audio_paths[sample], + args.target_root / f"{sample}.npy", + ): + if not path.is_file(): + raise FileNotFoundError(path) + destination = args.output_root / f"{sample}.pt" + if destination.exists() and not args.overwrite: + raise FileExistsError( + f"{destination} exists; select a new root or pass --overwrite" + ) + + audio_contract: dict[str, Any] | None = None + session_steps = {} + for sample in samples: + reference = torch.load(references[sample], map_location="cpu", weights_only=False) + times, _ = reference_alignment(reference) + target = np.load( + args.target_root / f"{sample}.npy", allow_pickle=False, mmap_mode="r" + ) + if target.shape != (len(times), 3): + raise ValueError(f"Reference/target mismatch for {sample}") + packed_audio, audio_contract = load_local_features( + audio_paths[sample], + reference, + sample_id=sample, + dataset="easycom", + modalities="audio", + expected_contract=audio_contract, + ) + session_steps[sample] = len(times) + steps = min(len(times), args.max_steps or len(times)) + print( + f"{sample}: {steps} states, {steps / 5:g}s; " + f"matched audio={packed_audio.shape[1] - 1}D", + flush=True, + ) + del reference, target, packed_audio + if args.dry_run: + print("Dry run only: source video decode and Qwen vision were not checked.") + return + + args.output_root.mkdir(parents=True, exist_ok=True) + encoder = FrozenQwenFullFrameEncoder(args) + assert audio_contract is not None + config = _requested_config(args, encoder, audio_contract) + for position, sample in enumerate(samples, 1): + print(f"[{position}/{len(samples)}] {sample}", flush=True) + reference = torch.load(references[sample], map_location="cpu", weights_only=False) + full_times, full_indices = reference_alignment(reference) + packed_audio, current_contract = load_local_features( + audio_paths[sample], + reference, + sample_id=sample, + dataset="easycom", + modalities="audio", + expected_contract=audio_contract, + ) + if current_contract != audio_contract: + raise AssertionError("Audio contract changed after preflight") + full_steps = session_steps[sample] + steps = min(full_steps, args.max_steps or full_steps) + times, indices, packed_audio = ( + full_times[:steps], + full_indices[:steps], + packed_audio[:steps], + ) + end = steps / 5 + media = _build_media(args, sample, end) + validate_source_timeline(reference, media) + visual_tokens = _extract_visual_tokens( + media, + encoder, + end=end, + fps=args.video_fps, + batch_size=args.image_batch_size, + no_progress=args.no_progress, + ) + audio = packed_audio[:, :-1].to(torch.bfloat16) + audio_valid = packed_audio[:, -1].bool() + if torch.count_nonzero(audio[~audio_valid]): + raise ValueError("Invalid inherited audio rows are not zero") + output = { + "format": FULLFRAME_AV_FORMAT, + "dataset": "easycom", + "sample_id": sample, + "times": times, + "label_indices": indices, + "partial": steps != full_steps, + "audio": audio, + "audio_valid": audio_valid, + "visual_tokens": visual_tokens, + "config": config, + "reference_features": str(references[sample].resolve()), + "source_audio_features": str(audio_paths[sample].resolve()), + "direct_media": media.metadata(), + } + destination = args.output_root / f"{sample}.pt" + with tempfile.NamedTemporaryFile( + dir=args.output_root, + prefix=f".{sample}.", + suffix=".tmp", + delete=False, + ) as handle: + temporary = Path(handle.name) + try: + torch.save(output, temporary) + os.replace(temporary, destination) + finally: + temporary.unlink(missing_ok=True) + print( + f"Saved {destination}: frames={len(visual_tokens['frame_times'])}, " + f"tokens={len(visual_tokens['features'])}, partial={output['partial']}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/fullframe_qwen_av.py b/fullframe_qwen_av.py new file mode 100644 index 0000000..faa2968 --- /dev/null +++ b/fullframe_qwen_av.py @@ -0,0 +1,192 @@ +"""Strict no-box full-frame Qwen AV sidecar contract and loader.""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +from local_turn_features import FramePersonFeatures, reference_alignment + + +FULLFRAME_AV_FORMAT = "qwen-fullframe-av-tokens-v1" + + +def load_fullframe_av_features( + path: Path, + reference: dict[str, Any], + *, + sample_id: str, + dataset: str, + expected_contract: dict[str, Any] | None = None, + history_max_gap_seconds: float | None = None, +) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]: + """Load aligned Qwen audio and ragged full-frame spatial token sets.""" + + times, indices = reference_alignment(reference) + payload = torch.load(path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict) or payload.get("format") != FULLFRAME_AV_FORMAT: + raise ValueError(f"Invalid full-frame feature format: {path}") + if payload.get("sample_id") != sample_id or payload.get("dataset") != dataset: + raise ValueError(f"Full-frame sample/dataset mismatch: {path}") + if dataset != "easycom": + raise ValueError("Full-frame F2/F3 currently support EasyCom only") + if payload.get("partial", False): + raise ValueError(f"Partial smoke-test features cannot be used for training: {path}") + local_indices = torch.as_tensor(payload.get("label_indices", [])) + local_times = torch.as_tensor(payload.get("times", []), dtype=torch.float64) + if local_indices.dtype != torch.long or not torch.equal(indices, local_indices): + raise ValueError(f"Full-frame label_indices mismatch: {path}") + if local_times.shape != times.shape or not torch.allclose( + local_times, times, atol=1e-7, rtol=0 + ): + raise ValueError(f"Full-frame times mismatch: {path}") + config = payload.get("config", {}) + if ( + config.get("causal_policy") != "strict-half-open-fullframe-grid" + or config.get("uses_speech_labels") is not False + or config.get("uses_boxes") is not False + or config.get("visual_feature_kind") + != "qwen-vision-post-merger-spatial-token" + or config.get("visual_source") + != "easycom-full-frame-after-session-banner-crop" + or config.get("audio_source") != "reused-qwen-local-mixed-audio" + ): + raise ValueError(f"Full-frame AV features lack the causal/no-box contract: {path}") + contract = {"format": FULLFRAME_AV_FORMAT, "modalities": "av", "config": config} + if expected_contract is not None and contract != expected_contract: + raise ValueError(f"Full-frame contract differs across sessions/checkpoint: {path}") + + 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) + or not audio.is_floating_point() + or not torch.isfinite(audio).all() + ): + raise ValueError(f"Invalid full-frame sidecar audio: {path}") + if ( + not isinstance(audio_valid, torch.Tensor) + or audio_valid.dtype != torch.bool + or audio_valid.shape != (len(times),) + ): + raise ValueError(f"Invalid full-frame 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") + visual_fps = config.get("visual_fps") + tokens = payload.get("visual_tokens") + if not isinstance(visual_dim, int) or visual_dim <= 0: + raise ValueError(f"Missing visual_dim: {path}") + if not isinstance(visual_fps, (int, float)) or not math.isfinite(visual_fps) or visual_fps <= 0: + raise ValueError(f"Invalid visual_fps: {path}") + if not isinstance(tokens, dict): + raise ValueError(f"Missing full-frame visual tokens: {path}") + features = tokens.get("features") + frame_times = tokens.get("frame_times") + frame_offsets = tokens.get("frame_offsets") + grid_shapes = tokens.get("grid_shapes") + source_clip_indices = tokens.get("frame_source_clip_indices") + 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 full-frame token features: {path}") + if ( + not isinstance(frame_times, torch.Tensor) + or frame_times.ndim != 1 + or not frame_times.is_floating_point() + ): + raise ValueError(f"Invalid full-frame frame_times: {path}") + frame_times = frame_times.to(dtype=torch.float64) + if ( + not len(frame_times) + or not torch.isfinite(frame_times).all() + or bool((frame_times < 0).any()) + or bool((frame_times[1:] <= frame_times[:-1]).any()) + ): + raise ValueError(f"Full-frame times must be finite/increasing: {path}") + expected_frame_count = max( + 0, + math.ceil(float(times[-1]) * float(visual_fps) - 1e-8), + ) + expected_frame_times = ( + torch.arange(expected_frame_count, dtype=torch.float64) + / float(visual_fps) + ) + if frame_times.shape != expected_frame_times.shape or not torch.allclose( + frame_times, expected_frame_times, atol=1e-8, rtol=0 + ): + raise ValueError(f"Full-frame times are not the fixed causal grid: {path}") + if ( + not isinstance(frame_offsets, torch.Tensor) + or frame_offsets.dtype != torch.long + or frame_offsets.shape != (len(frame_times) + 1,) + 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 full-frame frame_offsets: {path}") + if ( + not isinstance(grid_shapes, torch.Tensor) + or grid_shapes.dtype != torch.long + or grid_shapes.shape != (len(frame_times), 2) + or bool((grid_shapes <= 0).any()) + ): + raise ValueError(f"Invalid full-frame grid_shapes: {path}") + if not torch.equal(frame_offsets[1:] - frame_offsets[:-1], grid_shapes.prod(-1)): + raise ValueError(f"Token counts do not match full-frame grids: {path}") + if ( + not isinstance(source_clip_indices, torch.Tensor) + or source_clip_indices.dtype != torch.long + or source_clip_indices.shape != (len(frame_times),) + or bool((source_clip_indices < 0).any()) + or bool((source_clip_indices[1:] < source_clip_indices[:-1]).any()) + ): + raise ValueError(f"Invalid full-frame source-clip indices: {path}") + if float(frame_times[-1]) >= float(times[-1]) - 1e-8: + raise ValueError(f"Full-frame timeline extends past the reference cutoff: {path}") + + row_to_frame = torch.searchsorted(frame_times, times, right=False) - 1 + chosen = torch.full(times.shape, float("nan"), dtype=torch.float64) + valid_rows = row_to_frame >= 0 + chosen[valid_rows] = frame_times[row_to_frame[valid_rows]] + if bool((chosen[valid_rows] >= times[valid_rows]).any()): + raise ValueError(f"Full-frame mapping selected boundary/future evidence: {path}") + + previous = None + if history_max_gap_seconds is not None: + if not math.isfinite(history_max_gap_seconds) or history_max_gap_seconds <= 0: + raise ValueError("history_max_gap_seconds must be finite and positive") + previous = torch.full((len(features),), -1, dtype=torch.long) + for frame in range(1, len(frame_times)): + gap = float(frame_times[frame] - frame_times[frame - 1]) + same_clip = bool(source_clip_indices[frame] == source_clip_indices[frame - 1]) + same_grid = torch.equal(grid_shapes[frame], grid_shapes[frame - 1]) + if not same_clip or not same_grid or gap > history_max_gap_seconds + 1e-8: + continue + current_start, current_end = int(frame_offsets[frame]), int(frame_offsets[frame + 1]) + old_start, old_end = int(frame_offsets[frame - 1]), int(frame_offsets[frame]) + if current_end - current_start != old_end - old_start: + raise AssertionError("Equal grids have unequal token counts") + previous[current_start:current_end] = torch.arange(old_start, old_end) + + frames = FramePersonFeatures( + features, + frame_offsets, + row_to_frame.to(torch.long), + chosen, + previous, + ) + packed_audio = torch.cat((audio.float(), audio_valid[:, None].float()), dim=-1) + return packed_audio, frames, contract diff --git a/infer_turn_head.py b/infer_turn_head.py index 068e92e..4ce053d 100644 --- a/infer_turn_head.py +++ b/infer_turn_head.py @@ -14,12 +14,12 @@ from tqdm.auto import tqdm from easycom_head import ( CLASS_NAMES, DATASET_NAMES, + HISTORY_VISUAL_SET_FUSION_MODES, LOCAL_FUSION_SCOPES, NUM_CLASSES, NUM_HORIZONS, OFFICIAL_TEST_SESSIONS, STEP_SECONDS, - TRACK_PERSON_FUSION_MODES, TurnPredictionHead, average_precision_metrics, checkpoint_head_config, @@ -150,7 +150,7 @@ def _predict_states( start, min(start + batch_size, len(states)), ) - if head.local_fusion in TRACK_PERSON_FUSION_MODES: + if head.local_fusion in HISTORY_VISUAL_SET_FUSION_MODES: person_batch, person_mask = local_person_features.padded_history( rows, max_observations=head.person_history_frames, diff --git a/run_easycom_fullframe_qwen_av.sh b/run_easycom_fullframe_qwen_av.sh new file mode 100755 index 0000000..7746967 --- /dev/null +++ b/run_easycom_fullframe_qwen_av.sh @@ -0,0 +1,65 @@ +#!/usr/bin/env bash +# Foreground only. Usage: bash run_easycom_fullframe_qwen_av.sh [all|f2|f3] [seed ...] +set -euo pipefail +cd "$(dirname "$0")" + +FEATURE_ROOT=${FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack} +TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame} +FULLFRAME_ROOT=${FULLFRAME_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/fullframe_qwen_av_2fps_224} +CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_fullframe_qwen_av} +RESULT_ROOT=${RESULT_ROOT:-results/easycom_fullframe_qwen_av} +LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256} +ATTENTION_DIM=${ATTENTION_DIM:-256} +ATTENTION_HEADS=${ATTENTION_HEADS:-4} +HISTORY_FRAMES=${HISTORY_FRAMES:-2} +HISTORY_MAX_GAP_SECONDS=${HISTORY_MAX_GAP_SECONDS:-0.75} +TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-64} +INFER_BATCH_SIZE=${INFER_BATCH_SIZE:-128} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} +export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4} +export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4} + +variant=${1:-all} +if (( $# )); then shift; fi +if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi +case "$variant" in + all) variants=(f2 f3) ;; + f2|f3) variants=("$variant") ;; + *) echo "Expected all, f2, or f3" >&2; exit 2 ;; +esac + +for variant in "${variants[@]}"; do + case "$variant" in + f2) fusion=fullframe-spatial-attention ;; + f3) fusion=fullframe-spatiotemporal-attention ;; + esac + for seed in "${seeds[@]}"; do + checkpoint="$CHECKPOINT_ROOT/${variant}_seed${seed}.pth" + result="$RESULT_ROOT/${variant}_seed${seed}" + if [[ -e "$checkpoint" || -e "$result" ]]; then + echo "Refusing to overwrite $checkpoint or $result; select new output roots." >&2 + exit 1 + fi + python train_turn_head.py \ + --feature-root "$FEATURE_ROOT" --target-root "$TARGET_ROOT" \ + --local-feature-root "$FULLFRAME_ROOT" --local-modalities av \ + --local-fusion "$fusion" \ + --local-projection-dim "$LOCAL_PROJECTION_DIM" \ + --person-attention-dim "$ATTENTION_DIM" \ + --person-attention-heads "$ATTENTION_HEADS" \ + --person-history-frames "$HISTORY_FRAMES" \ + --person-track-max-gap-seconds "$HISTORY_MAX_GAP_SECONDS" \ + --output "$checkpoint" --sessions 4 5 6 7 8 9 10 11 12 \ + --epochs 4 --batch-size "$TRAIN_BATCH_SIZE" \ + --learning-rate 7e-5 --weight-decay 5e-5 \ + --warmup-epochs 1.6 --warmup-factor 0.3 --dropout 0.2 \ + --loss lstr-mce --transition-objective none --layer-fusion concat-mlp \ + --expected-mode open-stream-replay --seed "$seed" --device cuda:0 + python infer_turn_head.py \ + --feature-root "$FEATURE_ROOT" --target-root "$TARGET_ROOT" \ + --local-feature-root "$FULLFRAME_ROOT" \ + --checkpoint "$checkpoint" --output-dir "$result" --sessions 3 2 1 \ + --batch-size "$INFER_BATCH_SIZE" \ + --expected-mode open-stream-replay --device cuda:0 + done +done diff --git a/test_fullframe_qwen_av.py b/test_fullframe_qwen_av.py new file mode 100644 index 0000000..91d28a3 --- /dev/null +++ b/test_fullframe_qwen_av.py @@ -0,0 +1,196 @@ +"""CPU contract and head tests for no-box full-frame F2/F3.""" + +from __future__ import annotations + +import copy +from pathlib import Path +import tempfile +import unittest + +import numpy as np +import torch + +from easycom_head import TurnPredictionHead, checkpoint_head_config, load_session_data +from fullframe_qwen_av import FULLFRAME_AV_FORMAT, load_fullframe_av_features + + +def _fixture(*, partial: bool = False): + count = 6 + times = (torch.arange(count, dtype=torch.float64) + 1) / 5 + reference = { + "states": torch.randn(count, 2, 8), + "times": times, + "label_indices": torch.arange(count), + "start": 0, + "mode": "open-stream-replay", + } + sidecar = { + "format": FULLFRAME_AV_FORMAT, + "dataset": "easycom", + "sample_id": "Session_4", + "times": times.clone(), + "label_indices": torch.arange(count), + "partial": partial, + "audio": torch.randn(count, 5), + "audio_valid": torch.ones(count, dtype=torch.bool), + "visual_tokens": { + "features": torch.arange(12 * 6, dtype=torch.float32).reshape(12, 6), + "frame_times": torch.tensor([0.0, 0.5, 1.0], dtype=torch.float64), + "frame_offsets": torch.tensor([0, 4, 8, 12], dtype=torch.long), + "grid_shapes": torch.tensor([[2, 2], [2, 2], [2, 2]], dtype=torch.long), + "frame_source_clip_indices": torch.zeros(3, dtype=torch.long), + }, + "config": { + "visual_dim": 6, + "audio_dim": 5, + "visual_fps": 2.0, + "causal_policy": "strict-half-open-fullframe-grid", + "uses_speech_labels": False, + "uses_boxes": False, + "visual_feature_kind": "qwen-vision-post-merger-spatial-token", + "visual_source": "easycom-full-frame-after-session-banner-crop", + "audio_source": "reused-qwen-local-mixed-audio", + }, + } + return reference, sidecar + + +class FullFrameQwenAVTests(unittest.TestCase): + def setUp(self): + torch.set_num_threads(1) + + def _load(self, sidecar, reference, *, history=False): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "Session_4.pt" + torch.save(sidecar, path) + return load_fullframe_av_features( + path, + reference, + sample_id="Session_4", + dataset="easycom", + history_max_gap_seconds=0.75 if history else None, + ) + + def test_strict_boundary_mapping_and_same_grid_history(self): + reference, sidecar = _fixture() + audio, frames, contract = self._load(sidecar, reference, history=True) + self.assertEqual(audio.shape, (6, 6)) + self.assertEqual(frames.row_to_frame.tolist(), [0, 0, 1, 1, 1, 2]) + self.assertEqual( + frames.chosen_frame_times.tolist(), [0.0, 0.0, 0.5, 0.5, 0.5, 1.0] + ) + self.assertEqual(contract["format"], FULLFRAME_AV_FORMAT) + histories, mask = frames.padded_history( + torch.tensor([5]), max_observations=2 + ) + self.assertEqual(histories.shape, (1, 4, 2, 6)) + self.assertTrue(mask.all()) + tokens = sidecar["visual_tokens"]["features"] + self.assertTrue(torch.equal(histories[0, 0, 0], tokens[4])) + self.assertTrue(torch.equal(histories[0, 0, 1], tokens[8])) + + def test_rejects_partial_and_malformed_contracts(self): + reference, sidecar = _fixture(partial=True) + with self.assertRaisesRegex(ValueError, "Partial smoke-test"): + self._load(sidecar, reference) + + _, sidecar = _fixture() + wrong = copy.deepcopy(sidecar) + wrong["config"]["uses_boxes"] = True + with self.assertRaisesRegex(ValueError, "causal/no-box"): + self._load(wrong, reference) + + future = copy.deepcopy(sidecar) + future["visual_tokens"]["frame_times"][-1] = reference["times"][-1] + 0.01 + with self.assertRaisesRegex(ValueError, "fixed causal grid|extends past"): + self._load(future, reference) + + def test_f2_f3_forward_backward_and_checkpoint_config(self): + reference, sidecar = _fixture() + local, frames, _ = self._load(sidecar, reference, history=True) + rows = torch.tensor([0, 2, 3]) + states = torch.randn(3, 8) + + tokens, token_mask = frames.padded(rows) + f2 = TurnPredictionHead( + hidden_size=8, + local_visual_dim=6, + local_audio_dim=5, + local_projection_dim=4, + local_fusion="fullframe-spatial-attention", + person_attention_dim=4, + person_attention_heads=2, + person_history_frames=2, + ) + logits = f2(states, local[rows], tokens, token_mask) + self.assertEqual(logits.shape, (3, 10, 3)) + logits.sum().backward() + self.assertIsNotNone(f2.person_attention.in_proj_weight.grad) + + histories, history_mask = frames.padded_history(rows, max_observations=2) + f3 = TurnPredictionHead( + hidden_size=8, + local_visual_dim=6, + local_audio_dim=5, + local_projection_dim=4, + local_fusion="fullframe-spatiotemporal-attention", + person_attention_dim=4, + person_attention_heads=2, + person_history_frames=2, + ) + logits = f3(states, local[rows], histories, history_mask) + self.assertEqual(logits.shape, (3, 10, 3)) + logits.sum().backward() + self.assertIsNotNone(f3.person_track_gru.weight_ih_l0.grad) + config = checkpoint_head_config( + { + "head_config": { + "hidden_size": 8, + "local_visual_dim": 6, + "local_audio_dim": 5, + "local_fusion": "fullframe-spatiotemporal-attention", + "person_history_frames": 2, + } + } + ) + self.assertEqual(config["local_fusion"], "fullframe-spatiotemporal-attention") + self.assertEqual(config["person_history_frames"], 2) + + def test_session_loader_routes_fullframe_contract(self): + reference, sidecar = _fixture() + labels = np.tile( + np.asarray( + [[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float32 + ), + (2, 1), + ) + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + feature_root = root / "global" + target_root = root / "target" + local_root = root / "fullframe" + for folder in (feature_root, target_root, local_root): + folder.mkdir() + torch.save(reference, feature_root / "Session_4.pt") + torch.save(sidecar, local_root / "Session_4.pt") + np.save(target_root / "Session_4.npy", labels) + data = load_session_data( + feature_root, + target_root, + "Session_4", + expected_mode="open-stream-replay", + local_feature_root=local_root, + local_modalities="av", + local_fusion="fullframe-spatiotemporal-attention", + person_track_max_gap_seconds=0.75, + ) + self.assertEqual(data.local_features.shape, (6, 6)) + self.assertIsNotNone(data.local_person_features.previous_observation) + self.assertEqual( + data.feature_metadata["local_contract"]["format"], + FULLFRAME_AV_FORMAT, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_turn_head.py b/train_turn_head.py index 542c473..1aa8236 100644 --- a/train_turn_head.py +++ b/train_turn_head.py @@ -24,6 +24,8 @@ from tqdm.auto import tqdm from easycom_head import ( CLASS_NAMES, DATASET_NAMES, + FULLFRAME_FUSION_MODES, + HISTORY_VISUAL_SET_FUSION_MODES, HIGHRATE_PERSON_AV_FUSION_MODES, LAYER_FUSION_MODES, LOCAL_FUSION_MODES, @@ -87,6 +89,9 @@ def parse_args() -> argparse.Namespace: "highrate-person-av-attention lets global Omni plus one shared " "traditional CPC audio stream query causal 25 Hz person visual/ASD " "states sampled at the 5 Hz cutoffs; " + "fullframe-spatial-attention is F2 over the latest strictly " + "pre-cutoff full-frame grid; fullframe-spatiotemporal-attention " + "is F3 with a shared causal grid-position GRU; " "audio-person-frame-set and gated-person-state-adapter are " "current-sidecar F5/F6 screening variants, not full F5/F6." ), @@ -237,7 +242,7 @@ def _validate_args(args: argparse.Namespace) -> None: args.local_feature_root is None or args.local_modalities != "av" ): raise ValueError( - "Person-set --local-fusion modes require --local-feature-root and " + "Set/token --local-fusion modes require --local-feature-root and " "--local-modalities av" ) if args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES: @@ -528,7 +533,7 @@ def _evaluate_valid_horizons( person_mask = None if people is not None: rows = torch.arange(start, min(start + batch_size, len(states))) - if head.local_fusion in TRACK_PERSON_FUSION_MODES: + if head.local_fusion in HISTORY_VISUAL_SET_FUSION_MODES: person_batch, person_mask = people.padded_history( rows, max_observations=head.person_history_frames, @@ -1004,14 +1009,25 @@ def main() -> None: 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" - ) + if args.local_fusion in FULLFRAME_FUSION_MODES: + print( + "full-frame token sets: " + f"{len(person_features.frame_offsets) - 1} source frames, " + f"{len(person_features.features)} spatial tokens, " + f"max_tokens_per_frame={person_features.max_people}, " + f"history_frames={args.person_history_frames}, " + f"attention={args.person_attention_dim}D/" + f"{args.person_attention_heads} heads" + ) + else: + 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 ) @@ -1132,7 +1148,7 @@ def main() -> None: person_batch = None person_mask = None if person_features is not None: - if head.local_fusion in TRACK_PERSON_FUSION_MODES: + if head.local_fusion in HISTORY_VISUAL_SET_FUSION_MODES: person_batch, person_mask = person_features.padded_history( example_index_batch, max_observations=head.person_history_frames,