327 lines
18 KiB
Python
327 lines
18 KiB
Python
"""CPU tests for B1 alignment, causal pooling, fusion, and train/infer round trips."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import importlib.util
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
from easycom_head import TurnPredictionHead, checkpoint_head_config
|
|
from extract_turn_local_features import (
|
|
FrozenQwenLocalEncoder, ego4d_clip_boxes, extract_local,
|
|
validate_source_timeline,
|
|
model_config_digest,
|
|
)
|
|
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,
|
|
)
|
|
from train_turn_head import _load_examples
|
|
|
|
|
|
def fixture(n=20, name="Session_4", dataset="easycom"):
|
|
times = (torch.arange(n, dtype=torch.float64) + 1) / 5
|
|
reference = {"states": torch.randn(n, 4, 8), "label_indices": torch.arange(n),
|
|
"times": times, "mode": "open-stream-replay", "start": 0}
|
|
side = {"format": LOCAL_FORMAT, "dataset": dataset, "sample_id": name,
|
|
"times": times.clone(), "label_indices": torch.arange(n), "partial": False,
|
|
"visual": torch.randn(n, 6), "audio": torch.randn(n, 5),
|
|
"visual_valid": torch.ones(n, dtype=torch.bool),
|
|
"audio_valid": torch.ones(n, dtype=torch.bool),
|
|
"config": {"visual_dim": 6, "audio_dim": 5, "encoder_id": "synthetic-test",
|
|
"causal_policy": "half-open-trailing-window", "uses_speech_labels": False}}
|
|
side["visual_valid"][0] = False
|
|
side["visual"][0] = 0
|
|
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
|
|
|
|
|
|
class LocalTests(unittest.TestCase):
|
|
def setUp(self):
|
|
torch.set_num_threads(1)
|
|
|
|
def test_half_open_grid_and_pooling(self):
|
|
self.assertEqual(list(causal_frame_ticks(.2, 2)), [0])
|
|
self.assertEqual(list(causal_frame_ticks(1., 2)), [0, 1])
|
|
ts = torch.tensor([0., .5, 1., 1.5], dtype=torch.float64)
|
|
vectors = torch.tensor([[2.], [4.], [90.], [100.]])
|
|
valid = torch.ones(4, dtype=torch.bool)
|
|
value, present = pooled_recent_visual(ts, vectors, valid, 1., 1.)
|
|
self.assertTrue(present)
|
|
self.assertEqual(value.item(), 3.) # no frame at t=1
|
|
value, _ = pooled_recent_visual(ts, vectors, valid, 1.2, 1.)
|
|
self.assertEqual(value.item(), 47.) # .5 and 1.0, not 0.0
|
|
value, present = pooled_recent_visual(ts, vectors, valid * False, 1., 1.)
|
|
self.assertFalse(present)
|
|
self.assertEqual(value.item(), 0.)
|
|
|
|
def test_box_formats_and_crop(self):
|
|
person = {"Participant_ID": -1, "x1": 2, "y1": 3, "x2": 7, "y2": 9}
|
|
faces = easycom_boxes([{"Frame_Number": 1, "Participants": [person]}], "face")
|
|
heads = easycom_boxes([{"Frame_Number": 0, "Participants": person}], "head")
|
|
self.assertEqual(faces, heads)
|
|
self.assertEqual(crop_box([202., 3., 207., 9.], 1920, 1080, offset_x=200), (2, 3, 7, 9))
|
|
self.assertIsNone(crop_box([4., 4., 3., 5.], 10, 10))
|
|
self.assertIsNone(crop_box([float("nan"), 4., 9., 5.], 10, 10))
|
|
|
|
def test_ego4d_clip_frames(self):
|
|
boxes = [{"frame": 3, "clip_frame": None, "x": 1, "y": 2, "width": 3, "height": 4}]
|
|
person = {"person_id": "x", "camera_wearer": False, "tracking_paths": [{"track": boxes}, {"track": boxes}]}
|
|
wearer = {**person, "camera_wearer": True}
|
|
actual = ego4d_clip_boxes({"persons": [wearer, person]})
|
|
self.assertEqual(actual, {3: [{"id": "x", "xyxy": [1., 2., 4., 6.]}]})
|
|
|
|
def test_sidecar_alignment_contract_and_modalities(self):
|
|
reference, side, _ = fixture()
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
path = Path(directory) / "Session_4.pt"
|
|
for mode, width in (("visual", 7), ("audio", 6), ("av", 13)):
|
|
torch.save(side, path)
|
|
features, contract = load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities=mode)
|
|
self.assertEqual(features.shape, (20, width))
|
|
dims = local_head_dimensions(contract)
|
|
self.assertEqual(dims["local_visual_dim"], 0 if mode == "audio" else 6)
|
|
for field, value in (("times", side["times"] + .2), ("label_indices", torch.arange(20) + 1),
|
|
("sample_id", "Session_5"), ("partial", True),
|
|
("visual", torch.full((20, 6), float("nan"))),
|
|
("visual_valid", torch.ones(20))):
|
|
broken = {**side, field: value}
|
|
torch.save(broken, path)
|
|
with self.subTest(field=field), self.assertRaises(ValueError):
|
|
load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities="av")
|
|
torch.save(side, path)
|
|
wrong_contract = copy.deepcopy(contract)
|
|
wrong_contract["config"]["encoder_id"] = "other-weights"
|
|
with self.assertRaisesRegex(ValueError, "contract"):
|
|
load_local_features(path, reference, sample_id="Session_4", dataset="easycom", modalities="av", expected_contract=wrong_contract)
|
|
with self.assertRaisesRegex(ValueError, "times"):
|
|
reference_alignment({**reference, "times": reference["times"] + .2})
|
|
|
|
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)):
|
|
for projection in (0, 3):
|
|
head = TurnPredictionHead(hidden_size=8, num_input_layers=4, layer_fusion="concat-mlp",
|
|
local_visual_dim=v, local_audio_dim=a, local_projection_dim=projection,
|
|
transition_auxiliary=True, dropout=0)
|
|
width = sum(d + 1 for d in (v, a) if d)
|
|
local = torch.randn(3, width)
|
|
offset = 0
|
|
for d in (v, a):
|
|
if d:
|
|
local[:, offset + d] = torch.tensor([0., 1., 1.])
|
|
offset += d + 1
|
|
logits, transition = head.forward_with_transition(states, local)
|
|
self.assertEqual(logits.shape, (3, 10, 3))
|
|
self.assertEqual(transition.shape, (3, 10))
|
|
(logits.square().mean() + transition.square().mean()).backward()
|
|
self.assertIsNotNone(head.net[1].weight.grad)
|
|
if width:
|
|
# Changing invalid raw features cannot change the readout.
|
|
changed = local.clone()
|
|
offset = 0
|
|
for d in (v, a):
|
|
if d:
|
|
changed[0, offset:offset + d] = 1000
|
|
offset += d + 1
|
|
self.assertTrue(torch.equal(head(states, local)[0], head(states, changed)[0]))
|
|
with self.assertRaisesRegex(ValueError, "requires aligned"):
|
|
head(states)
|
|
config = {"hidden_size": 8, "num_input_layers": 4, "layer_fusion": "concat-mlp",
|
|
"local_visual_dim": v, "local_audio_dim": a, "local_projection_dim": projection,
|
|
"transition_auxiliary": True, "dropout": 0}
|
|
restored = TurnPredictionHead(**checkpoint_head_config({"head_config": config}))
|
|
restored.load_state_dict(head.state_dict())
|
|
self.assertTrue(torch.equal(restored(states, local), head(states, local)))
|
|
|
|
def test_legacy_head_parameter_keys_and_values(self):
|
|
torch.manual_seed(42)
|
|
old = torch.nn.Sequential(torch.nn.LayerNorm(8), torch.nn.Linear(8, 4), torch.nn.GELU(),
|
|
torch.nn.Dropout(.1), torch.nn.Linear(4, 30))
|
|
torch.manual_seed(42)
|
|
head = TurnPredictionHead(hidden_size=8)
|
|
self.assertEqual(set(head.state_dict()), {f"net.{key}" for key in old.state_dict()})
|
|
for key, value in old.state_dict().items():
|
|
self.assertTrue(torch.equal(value, head.state_dict()[f"net.{key}"]))
|
|
|
|
def test_sparse_selection_keeps_local_rows(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
for folder in ("features", "local", "targets"):
|
|
(root / folder).mkdir()
|
|
reference, side, targets = fixture()
|
|
torch.save(reference, root / "features/Session_4.pt")
|
|
torch.save(side, root / "local/Session_4.pt")
|
|
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[0], reference["states"][[1, 4]]))
|
|
self.assertTrue(torch.equal(data[1][:, 0], torch.from_numpy(targets[[2, 5]])))
|
|
|
|
def test_source_inventory_rejects_different_clip_order(self):
|
|
source = {"path": "/old/00-00-000.mp4", "source_frames": 1200, "dense_start": 0., "duration": 60.}
|
|
media = SimpleNamespace(clips=[object()], metadata=lambda: {"source_clips": [source]})
|
|
validate_source_timeline({"direct_media": {"source_clips": [source]}}, media)
|
|
bad = {**source, "path": "/old/01-00-000.mp4"}
|
|
with self.assertRaisesRegex(ValueError, "timeline differs"):
|
|
validate_source_timeline({"direct_media": {"source_clips": [bad]}}, media)
|
|
|
|
def test_config_digest_ignores_only_checkpoint_location(self):
|
|
a = {"_name_or_path": "/old/model", "audio_config": {"_name_or_path": "/old/model", "output_dim": 8}}
|
|
b = {"_name_or_path": "/new/model", "audio_config": {"_name_or_path": "/new/model", "output_dim": 8}}
|
|
self.assertEqual(model_config_digest(a), model_config_digest(b))
|
|
b["audio_config"]["output_dim"] = 16
|
|
self.assertNotEqual(model_config_digest(a), model_config_digest(b))
|
|
|
|
def test_encoder_image_pool_splits_people_before_mean(self):
|
|
encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder)
|
|
encoder.visual_dim, encoder.device, encoder.image_batch_size, encoder.max_pixels = 2, torch.device("cpu"), 8, 50176
|
|
encoder.processor = SimpleNamespace(image_processor=lambda **kwargs: {
|
|
"pixel_values": torch.ones(12, 2), "image_grid_thw": torch.tensor([[1, 2, 2], [1, 2, 4]])})
|
|
encoder.model = SimpleNamespace(visual=SimpleNamespace(dtype=torch.float32),
|
|
config=SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)),
|
|
get_image_features=lambda *_: torch.tensor([[1., 1.], [3., 3.], [5., 5.]]))
|
|
actual = encoder.images([object(), object()])
|
|
self.assertTrue(torch.equal(actual, torch.tensor([[1., 1.], [4., 4.]])))
|
|
|
|
def test_extraction_never_encodes_future_audio_and_saves_ragged_people(self):
|
|
class Media:
|
|
clips = [SimpleNamespace(crop=None, dense_start=0.)]
|
|
def _source_frame_for_time(self, t):
|
|
return 0, round(t * 20), t
|
|
def frame_at_or_before(self, t):
|
|
return Image.new("RGB", (16, 16))
|
|
def audio_slice(self, start, end):
|
|
return np.arange(round(start * 16000), round(end * 16000), dtype=np.float32)
|
|
class Encoder:
|
|
visual_dim = audio_dim = 2
|
|
def __init__(self):
|
|
self.ends = []
|
|
def images(self, crops):
|
|
return torch.ones(len(crops), 2)
|
|
def audio(self, waveform):
|
|
self.ends.append(float(waveform[-1]))
|
|
return torch.tensor([float(waveform[0]), float(waveform[-1])])
|
|
encoder = Encoder()
|
|
args = SimpleNamespace(video_fps=2., visual_window_seconds=1., audio_window_seconds=1.,
|
|
box_coordinate_space="camera", crop_expansion=1., no_progress=True)
|
|
annotations = [{0: [{"id": "1", "xyxy": [0., 0., 8., 8.]}]}]
|
|
result = extract_local(Media(), annotations, encoder, torch.arange(1, 11, dtype=torch.float64) / 5, args)
|
|
self.assertEqual(encoder.ends, [3200 * (i + 1) - 1 for i in range(10)])
|
|
self.assertEqual(result["persons"]["frame_offsets"].tolist(), [0, 1, 1, 1, 1])
|
|
self.assertTrue(result["visual_valid"][:5].all())
|
|
self.assertFalse(result["visual_valid"][5:].any())
|
|
self.assertTrue(result["audio_valid"].all())
|
|
|
|
|
|
@unittest.skipUnless(importlib.util.find_spec("transformers"), "optional: install transformers==4.52.3")
|
|
class TinyQwenEncoderTests(unittest.TestCase):
|
|
def test_real_encoder_apis_without_downloading_weights(self):
|
|
import transformers
|
|
if transformers.__version__ != "4.52.3":
|
|
self.skipTest("Pinned encoder API test requires transformers==4.52.3")
|
|
from transformers import (
|
|
Qwen2_5OmniThinkerConfig, Qwen2_5OmniThinkerForConditionalGeneration,
|
|
Qwen2VLImageProcessor, WhisperFeatureExtractor,
|
|
)
|
|
torch.set_num_threads(1)
|
|
config = Qwen2_5OmniThinkerConfig(
|
|
audio_config=dict(d_model=16, encoder_layers=1, encoder_attention_heads=2,
|
|
encoder_ffn_dim=32, output_dim=8),
|
|
vision_config=dict(depth=1, hidden_size=16, intermediate_size=32, num_heads=2,
|
|
out_hidden_size=8, fullatt_block_indexes=[0]),
|
|
text_config=dict(vocab_size=32, hidden_size=8, intermediate_size=16,
|
|
num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2),
|
|
)
|
|
encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder)
|
|
encoder.model = Qwen2_5OmniThinkerForConditionalGeneration(config).eval().requires_grad_(False)
|
|
encoder.model.model = None
|
|
encoder.model.lm_head = None
|
|
encoder.visual_dim = encoder.audio_dim = 8
|
|
encoder.device, encoder.max_pixels, encoder.image_batch_size = torch.device("cpu"), 224 * 224, 8
|
|
encoder.processor = SimpleNamespace(image_processor=Qwen2VLImageProcessor(),
|
|
feature_extractor=WhisperFeatureExtractor(feature_size=128))
|
|
crops = [Image.new("RGB", (56, 84), "red"), Image.new("RGB", (112, 56), "blue")]
|
|
batch = encoder.images(crops)
|
|
single = encoder.images(crops[:1])
|
|
self.assertEqual(batch.shape, (2, 8))
|
|
self.assertTrue(torch.allclose(batch[0], single[0], atol=1e-5))
|
|
for n in (3200, 6400, 16000):
|
|
audio = encoder.audio(np.sin(np.arange(n, dtype=np.float32) / 10))
|
|
self.assertEqual(audio.shape, (8,))
|
|
self.assertTrue(torch.isfinite(audio).all())
|
|
self.assertFalse(audio.requires_grad)
|
|
|
|
|
|
class RoundTripTests(unittest.TestCase):
|
|
def test_train_infer_easycom_and_ego4d(self):
|
|
repo = Path(__file__).resolve().parent
|
|
env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1")
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
for folder in ("features", "local", "targets"):
|
|
(root / folder).mkdir()
|
|
train_uid = "00000000-0000-0000-0000-000000000004"
|
|
test_uid = "00000000-0000-0000-0000-000000000001"
|
|
manifest = root / "ego4d.tsv"
|
|
manifest.write_text(f"split\tuid\ntrain\t{train_uid}\ntest\t{test_uid}\n")
|
|
for name in ("Session_1", "Session_4", "Session_5", train_uid, test_uid):
|
|
dataset = "easycom" if name.startswith("Session") else "ego4d"
|
|
reference, side, targets = fixture(name=name, dataset=dataset)
|
|
torch.save(reference, root / "features" / f"{name}.pt")
|
|
torch.save(side, root / "local" / f"{name}.pt")
|
|
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"):
|
|
checkpoint = root / f"{mode}.pth"
|
|
train = [sys.executable, str(repo / "train_turn_head.py"), *common,
|
|
"--epochs", "1", "--warmup-epochs", "0", "--batch-size", "4",
|
|
"--layer-fusion", "concat-mlp", "--output", str(checkpoint)]
|
|
infer = [sys.executable, str(repo / "infer_turn_head.py"), *common,
|
|
"--checkpoint", str(checkpoint), "--output-dir", str(root / f"result_{mode}")]
|
|
if mode == "ego4d":
|
|
train += ["--dataset", "ego4d", "--manifest", str(manifest)]
|
|
infer += ["--dataset", "ego4d", "--manifest", str(manifest)]
|
|
else:
|
|
train += ["--sessions", "4", "5"]
|
|
infer += ["--sessions", "1"]
|
|
if mode == "visual":
|
|
train += ["--dev-sessions", "5", "--local-projection-dim", "0"]
|
|
if mode == "av":
|
|
train += ["--transition-objective", "r1-supervised"]
|
|
if mode != "baseline":
|
|
train += ["--local-feature-root", str(root / "local"), "--local-modalities", "av" if mode == "ego4d" else mode]
|
|
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)
|
|
self.assertEqual(completed.returncode, 0, f"{mode}: {completed.stdout}\n{completed.stderr}")
|
|
metrics = json.loads((root / f"result_{mode}/metrics.json").read_text())
|
|
self.assertEqual(metrics["checkpoint_epoch"], 1)
|
|
self.assertEqual(bool(metrics["local_contract"]), mode != "baseline")
|
|
if mode == "av":
|
|
self.assertTrue(metrics["transition_auxiliary"])
|
|
# The reporting command uses the same metric keys as inference.
|
|
completed = subprocess.run(
|
|
[sys.executable, str(repo / "summarize_local_baselines.py"), "--root", str(root)],
|
|
cwd=repo, capture_output=True, text=True, timeout=20,
|
|
)
|
|
self.assertEqual(completed.returncode, 0, completed.stderr)
|
|
self.assertIn("Mean ± sample standard deviation", completed.stdout)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|