Add no-box full-frame Qwen AV baselines
This commit is contained in:
Xin Wang
@@ -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()
|
||||
Reference in New Issue
Block a user