Add Ego4D local AV ablation workflow
This commit is contained in:
Xin Wang
@@ -0,0 +1,206 @@
|
||||
"""CPU checks for manifest extraction, safe reuse and the four-run Ego4D family."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import copy
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from easycom_head import load_ego4d_manifest_split
|
||||
import extract_turn_local_features as extract
|
||||
import run_ego4d_local_baselines as runner
|
||||
from test_local_turn_features import fixture
|
||||
|
||||
|
||||
class Ego4DLocalWorkflowTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
torch.set_num_threads(1)
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
self.root = Path(self.temp.name)
|
||||
self.names = [f"00000000-0000-0000-0000-{i:012d}" for i in (4, 1, 5)]
|
||||
self.splits = ["train", "test", "train"]
|
||||
for folder in ("features", "targets", "local", "v1/clips", "v1/annotations"):
|
||||
(self.root / folder).mkdir(parents=True)
|
||||
self.manifest = self.root / "ego4d_egospeak.tsv"
|
||||
self.manifest.write_text("split\tuid\n" + "".join(
|
||||
f"{split}\t{name}\n" for split, name in zip(self.splits, self.names)))
|
||||
for name in self.names:
|
||||
reference, side, labels = fixture(name=name, dataset="ego4d")
|
||||
torch.save(reference, self.root / "features" / f"{name}.pt")
|
||||
torch.save(side, self.root / "local" / f"{name}.pt")
|
||||
np.save(self.root / "targets" / f"{name}.npy", labels)
|
||||
(self.root / "v1/clips" / f"{name}.mp4").write_bytes(b"fake-media-not-decoded")
|
||||
for annotation_split, split in (("train", "train"), ("val", "test")):
|
||||
payload = {"videos": [{"clips": [{"clip_uid": name, "persons": []}
|
||||
for name, part in zip(self.names, self.splits) if part == split]}]}
|
||||
(self.root / "v1/annotations" / f"av_{annotation_split}.json").write_text(json.dumps(payload))
|
||||
|
||||
def extraction_args(self, extra=()):
|
||||
return ["--dataset", "ego4d", "--data-root", str(self.root),
|
||||
"--feature-root", str(self.root / "features"), "--target-root", str(self.root / "targets"),
|
||||
"--manifest", str(self.manifest), "--output-root", str(self.root / "extracted"), *extra]
|
||||
|
||||
def runner_args(self, extra=()):
|
||||
return ["--data-root", str(self.root), "--feature-root", str(self.root / "features"),
|
||||
"--target-root", str(self.root / "targets"), "--local-feature-root", str(self.root / "local"),
|
||||
"--checkpoint-root", str(self.root / "checkpoints"), "--result-root", str(self.root / "results"),
|
||||
"--seeds", "0", "--device", "cpu", "--no-progress", *extra]
|
||||
|
||||
def test_manifest_order_split_and_selection(self):
|
||||
self.assertEqual(load_ego4d_manifest_split(self.manifest, "all"), self.names)
|
||||
args = extract.parse_args(self.extraction_args(["--split", "train", "--start-index", "2", "--limit", "1"]))
|
||||
self.assertEqual(extract.select_samples(args), [self.names[2]])
|
||||
args = extract.parse_args(self.extraction_args(["--start-index", "2", "--end-index", "2"]))
|
||||
self.assertEqual(extract.select_samples(args), [self.names[1]])
|
||||
args = extract.parse_args(self.extraction_args(["--start-index", "100"]))
|
||||
with self.assertRaisesRegex(ValueError, "No samples"):
|
||||
extract.select_samples(args)
|
||||
with self.manifest.open("a") as handle:
|
||||
handle.write(f"test\t{self.names[0]}\n")
|
||||
with self.assertRaisesRegex(ValueError, "Duplicate"):
|
||||
load_ego4d_manifest_split(self.manifest, "train")
|
||||
|
||||
def test_invalid_cli_combinations(self):
|
||||
for extra in (["--samples", self.names[0]], ["--dataset", "easycom"], ["--start-index", "0"],
|
||||
["--limit", "0"], ["--max-steps", "20", "--reuse-existing"],
|
||||
["--overwrite", "--reuse-existing"]):
|
||||
with self.subTest(extra=extra), contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
|
||||
extract.parse_args(self.extraction_args(extra))
|
||||
|
||||
def test_dry_run_no_model_decode_tracks_or_writes(self):
|
||||
with patch("extract_turn_local_features.FrozenQwenLocalEncoder") as encoder, \
|
||||
patch("extract_turn_local_features.build_media_and_boxes") as media, \
|
||||
patch("extract_turn_local_features.load_ego4d_boxes") as boxes:
|
||||
extract.main(self.extraction_args(["--dry-run"]))
|
||||
encoder.assert_not_called()
|
||||
media.assert_not_called()
|
||||
boxes.assert_not_called()
|
||||
self.assertFalse((self.root / "extracted").exists())
|
||||
np.save(self.root / "targets" / f"{self.names[1]}.npy", np.zeros((19, 3)))
|
||||
with self.assertRaisesRegex(ValueError, "Reference/target mismatch"):
|
||||
extract.main(self.extraction_args(["--dry-run"]))
|
||||
|
||||
def test_one_encoder_per_batch_and_validated_reuse(self):
|
||||
fake_encoder = SimpleNamespace(transformers_version="4.57.6", visual_dim=6, audio_dim=5,
|
||||
model_config_sha256="same-checkpoint")
|
||||
_, side, _ = fixture(dataset="ego4d")
|
||||
media = SimpleNamespace(clips=[], metadata=lambda: {})
|
||||
with patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="4.57.6")}), \
|
||||
patch("extract_turn_local_features.FrozenQwenLocalEncoder", return_value=fake_encoder) as encoder, \
|
||||
patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])) as build, \
|
||||
patch("extract_turn_local_features.extract_local", side_effect=lambda *a: copy.deepcopy(side)):
|
||||
extract.main(self.extraction_args())
|
||||
encoder.assert_called_once()
|
||||
self.assertEqual(build.call_count, 3)
|
||||
encoder.reset_mock()
|
||||
build.reset_mock()
|
||||
extract.main(self.extraction_args(["--reuse-existing"]))
|
||||
encoder.assert_not_called()
|
||||
build.assert_not_called()
|
||||
with self.assertRaisesRegex(ValueError, "audio_window_seconds"):
|
||||
extract.main(self.extraction_args(["--reuse-existing", "--audio-window-seconds", "2"]))
|
||||
first = self.root / "extracted" / f"{self.names[0]}.pt"
|
||||
payload = torch.load(first, weights_only=False)
|
||||
payload["partial"] = True
|
||||
torch.save(payload, first)
|
||||
with self.assertRaisesRegex(ValueError, "Partial smoke-test"):
|
||||
extract.main(self.extraction_args(["--reuse-existing"]))
|
||||
|
||||
def test_four_commands_keep_splits_epochs_and_feature_ablation(self):
|
||||
args = runner.parse_args(self.runner_args())
|
||||
fusion = runner.preflight(args)
|
||||
self.assertEqual(fusion, "concat-mlp")
|
||||
commands = list(runner.experiment_commands(args, fusion))
|
||||
self.assertEqual([name for name, _, _ in commands], [v + "_seed0" for v in runner.VARIANTS])
|
||||
for name, train, infer in commands:
|
||||
self.assertEqual(train[train.index("--epochs") + 1], "4")
|
||||
for command in (train, infer):
|
||||
self.assertEqual(command[command.index("--dataset") + 1], "ego4d")
|
||||
self.assertEqual(command[command.index("--manifest") + 1], str(self.manifest.resolve()))
|
||||
self.assertNotIn("--sessions", command)
|
||||
self.assertNotIn("--dev-sessions", command)
|
||||
self.assertEqual("--local-feature-root" in command, name != "baseline_seed0")
|
||||
if name != "baseline_seed0":
|
||||
self.assertEqual(train[train.index("--local-modalities") + 1], name.split("_")[0])
|
||||
|
||||
def test_family_preflight_catches_missing_test_sidecar_before_training(self):
|
||||
(self.root / "local" / f"{self.names[1]}.pt").unlink()
|
||||
with patch("run_ego4d_local_baselines.subprocess.run") as run:
|
||||
with self.assertRaisesRegex(FileNotFoundError, self.names[1]):
|
||||
runner.main(self.runner_args())
|
||||
run.assert_not_called()
|
||||
self.assertEqual(runner.preflight(runner.parse_args(self.runner_args(["--variants", "baseline"]))),
|
||||
"concat-mlp")
|
||||
|
||||
def test_preflight_rejects_mixed_contracts_and_existing_outputs(self):
|
||||
path = self.root / "local" / f"{self.names[1]}.pt"
|
||||
side = torch.load(path, weights_only=False)
|
||||
side["config"]["encoder_id"] = "different-test-encoder"
|
||||
torch.save(side, path)
|
||||
with self.assertRaisesRegex(ValueError, "contract differs"):
|
||||
runner.preflight(runner.parse_args(self.runner_args()))
|
||||
(self.root / "results" / "baseline_seed0").mkdir(parents=True)
|
||||
with self.assertRaisesRegex(FileExistsError, "Refusing"):
|
||||
runner.preflight(runner.parse_args(self.runner_args(["--variants", "baseline"])))
|
||||
|
||||
def test_single_layer_auto_and_dry_run(self):
|
||||
for name in self.names:
|
||||
path = self.root / "features" / f"{name}.pt"
|
||||
reference = torch.load(path, weights_only=False)
|
||||
reference["states"] = reference["states"].mean(1)
|
||||
torch.save(reference, path)
|
||||
args = runner.parse_args(self.runner_args())
|
||||
self.assertEqual(runner.preflight(args), "single")
|
||||
for _, train, _ in runner.experiment_commands(args, "single"):
|
||||
self.assertNotIn("--layer-fusion", train)
|
||||
with self.assertRaisesRegex(ValueError, "requires stacked"):
|
||||
runner.preflight(runner.parse_args(self.runner_args(["--layer-fusion", "concat-mlp"])))
|
||||
with patch("run_ego4d_local_baselines.subprocess.run") as run:
|
||||
runner.main(self.runner_args(["--dry-run"]))
|
||||
run.assert_not_called()
|
||||
self.assertFalse((self.root / "checkpoints").exists())
|
||||
|
||||
def test_full_cpu_family_round_trip_and_summary(self):
|
||||
env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1")
|
||||
completed = subprocess.run([sys.executable, str(runner.REPO / "run_ego4d_local_baselines.py"),
|
||||
*self.runner_args()], env=env, capture_output=True, text=True, timeout=90)
|
||||
self.assertEqual(completed.returncode, 0, completed.stdout + completed.stderr)
|
||||
self.assertIn("av - audio (seeds=[0])", completed.stdout)
|
||||
scores = {}
|
||||
for variant in runner.VARIANTS:
|
||||
path = self.root / "results" / f"{variant}_seed0" / "metrics.json"
|
||||
metrics = json.loads(path.read_text())
|
||||
self.assertEqual(metrics["checkpoint_epoch"], 4)
|
||||
self.assertEqual(metrics["dataset"], "ego4d")
|
||||
self.assertEqual(metrics["checkpoint_dev_sessions"], [])
|
||||
self.assertEqual(metrics["sessions"], [self.names[1]])
|
||||
self.assertEqual(metrics["checkpoint_train_sessions"], [self.names[0], self.names[2]])
|
||||
self.assertEqual(bool(metrics["local_contract"]), variant != "baseline")
|
||||
scores[variant] = metrics["causal_valid_only"]["horizons"][0]["map"]
|
||||
expected_delta = 100 * (scores["av"] - scores["audio"])
|
||||
self.assertIn(f"av - audio (seeds=[0]): {expected_delta:+.5f} ± n/a", completed.stdout)
|
||||
# Do not silently pool a mismatched epoch as a paired comparison.
|
||||
metrics["checkpoint_epoch"] = 25
|
||||
path.write_text(json.dumps(metrics))
|
||||
completed = subprocess.run([sys.executable, str(runner.REPO / "summarize_local_baselines.py"),
|
||||
"--root", str(self.root / "results"), "--dataset", "ego4d", "--paired-deltas"],
|
||||
env=env, capture_output=True, text=True, timeout=20)
|
||||
self.assertNotEqual(completed.returncode, 0)
|
||||
self.assertIn("different splits/epochs/loss/fusion", completed.stderr)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user