217 lines
11 KiB
Python
217 lines
11 KiB
Python
"""EasyCom encoder-only split, packing and real CPU train/infer checks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
import copy
|
|
import io
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from easycom_head import load_session_data
|
|
from local_turn_features import MODALITIES
|
|
import prepare_encoder_only_features as prepare
|
|
import run_easycom_encoder_only as runner
|
|
from test_local_turn_features import fixture
|
|
|
|
|
|
class EasyComEncoderOnlyTests(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"Session_{i}" for i in range(1, 13)]
|
|
self.train = self.names[3:]
|
|
self.test = ["Session_3", "Session_2", "Session_1"]
|
|
self.source = self.root / "local"
|
|
self.targets = self.root / "easycom_target_frame"
|
|
self.source.mkdir()
|
|
self.targets.mkdir()
|
|
self.payloads = {}
|
|
for name in self.names:
|
|
# No Thinker features, manifest, media, annotations or model exists.
|
|
_, side, target = fixture(name=name, dataset="easycom")
|
|
side["reference_features"] = "/must-not-be-read/thinker.pt"
|
|
side["states"] = "POISON: not an encoder input"
|
|
torch.save(side, self.source / f"{name}.pt")
|
|
np.save(self.targets / f"{name}.npy", target)
|
|
self.payloads[name] = side
|
|
|
|
def output_root(self, modality="av"):
|
|
return self.root / f"encoder_{modality}"
|
|
|
|
def prepare_args(self, modality="av", extra=()):
|
|
return ["--dataset", "easycom", "--data-root", str(self.root),
|
|
"--local-feature-root", str(self.source), "--modalities", modality,
|
|
"--output-root", str(self.output_root(modality)), "--no-progress", *extra]
|
|
|
|
def runner_args(self, modality="av", extra=()):
|
|
return ["--data-root", str(self.root), "--feature-root", str(self.output_root(modality)),
|
|
"--modalities", modality, "--checkpoint-root", str(self.root / "checkpoints"),
|
|
"--result-root", str(self.root / "results"), "--seeds", "0", "--device", "cpu",
|
|
"--no-progress", *extra]
|
|
|
|
def export(self, modality="av"):
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
prepare.main(self.prepare_args(modality))
|
|
|
|
def test_dataset_defaults_and_ascending_b1_training_order(self):
|
|
args = prepare.parse_args(["--dataset", "easycom"])
|
|
root = Path("/data1/yfzx/wangx/Datasets/easycom")
|
|
self.assertEqual(args.data_root, root)
|
|
self.assertEqual(args.target_root, root / "easycom_target_frame")
|
|
self.assertIsNone(args.manifest)
|
|
run = runner.parse_args(["--feature-root", str(self.output_root())])
|
|
self.assertEqual(run.dataset, "easycom")
|
|
self.assertEqual(run.data_root, root)
|
|
self.assertEqual(run.target_root, root / "easycom_target_frame")
|
|
self.assertIsNone(run.manifest)
|
|
self.assertEqual(run.result_root.name, "easycom_encoder_only_1s")
|
|
self.assertEqual(run.checkpoint_root.name, "easycom_encoder_only_1s")
|
|
self.assertEqual(prepare.encoder_sample_split("easycom"), (self.names, self.train, self.test))
|
|
# Preserve the original export command's Ego4D default.
|
|
self.assertEqual(prepare.parse_args([]).dataset, "ego4d")
|
|
|
|
def test_all_modalities_export_twelve_real_files_and_exact_native_inputs(self):
|
|
for modality, parts in MODALITIES.items():
|
|
with self.subTest(modality=modality):
|
|
self.export(modality)
|
|
output = self.output_root(modality)
|
|
self.assertEqual((output / "complete.txt").read_text(), "12 sessions\n")
|
|
manifest = json.loads((output / "export_manifest.json").read_text())
|
|
self.assertEqual(manifest["dataset"], "easycom")
|
|
self.assertIsNone(manifest["manifest"])
|
|
self.assertEqual(manifest["train_samples"], self.train)
|
|
self.assertEqual(manifest["test_samples"], self.test)
|
|
self.assertEqual(manifest["sessions"], dict.fromkeys(self.names, 20))
|
|
self.assertEqual(manifest["num_states"], 240)
|
|
for name in self.names:
|
|
path = output / f"{name}.pt"
|
|
self.assertTrue(path.is_file())
|
|
self.assertFalse(path.is_symlink())
|
|
packed = torch.load(path, weights_only=False)
|
|
original = self.payloads[name]
|
|
expected = torch.cat([tensor for part in parts for tensor in (
|
|
original[part].float(), original[f"{part}_valid"][:, None].float(),
|
|
)], dim=-1).to(torch.bfloat16)
|
|
self.assertTrue(torch.equal(packed["states"], expected))
|
|
self.assertTrue(torch.equal(packed["times"], original["times"]))
|
|
self.assertTrue(torch.equal(packed["label_indices"], original["label_indices"]))
|
|
self.assertEqual(packed["dataset"], "easycom")
|
|
self.assertIs(packed["config"]["uses_thinker_states"], False)
|
|
data = load_session_data(output, self.targets, name,
|
|
expected_mode=f"encoder-only-{modality}")
|
|
self.assertIsNone(data.local_features)
|
|
self.assertEqual(tuple(data.states.shape), tuple(expected.shape))
|
|
with self.assertRaises(FileExistsError):
|
|
prepare.main(self.prepare_args(modality))
|
|
|
|
def test_partial_wrong_dataset_and_missing_heldout_session_fail_before_write(self):
|
|
name = "Session_3"
|
|
path = self.source / f"{name}.pt"
|
|
for key, value in (("partial", True), ("dataset", "ego4d")):
|
|
side = copy.deepcopy(self.payloads[name])
|
|
side[key] = value
|
|
torch.save(side, path)
|
|
with self.subTest(key=key), self.assertRaises(ValueError):
|
|
prepare.main(self.prepare_args())
|
|
self.assertFalse(self.output_root().exists())
|
|
path.unlink()
|
|
with self.assertRaises(FileNotFoundError):
|
|
prepare.main(self.prepare_args())
|
|
self.assertFalse(self.output_root().exists())
|
|
|
|
def test_easycom_rejects_ego4d_manifest(self):
|
|
with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
|
|
prepare.parse_args(self.prepare_args(extra=["--manifest", "/ego4d.tsv"]))
|
|
with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
|
|
runner.parse_args(self.runner_args(extra=["--manifest", "/ego4d.tsv"]))
|
|
|
|
def test_dry_runs_write_nothing_and_use_fixed_sessions(self):
|
|
with contextlib.redirect_stdout(io.StringIO()):
|
|
prepare.main(self.prepare_args(extra=["--dry-run"]))
|
|
self.assertFalse(self.output_root().exists())
|
|
self.export()
|
|
args = runner.parse_args(self.runner_args(extra=["--seeds", "0", "1", "2"]))
|
|
commands = list(runner.experiment_commands(args))
|
|
self.assertEqual(len(commands), 3)
|
|
for seed, (train, infer) in enumerate(commands):
|
|
self.assertEqual(train[train.index("--epochs") + 1], "4")
|
|
self.assertEqual(train[train.index("--warmup-epochs") + 1], "1.6")
|
|
self.assertEqual(train[train.index("--seed") + 1], str(seed))
|
|
self.assertEqual(train[train.index("--sessions") + 1:], self.train)
|
|
self.assertEqual(infer[infer.index("--sessions") + 1:], self.test)
|
|
for command in (train, infer):
|
|
self.assertEqual(command[command.index("--dataset") + 1], "easycom")
|
|
self.assertEqual(command[command.index("--expected-mode") + 1], "encoder-only-av")
|
|
for forbidden in ("--manifest", "--local-feature-root", "--layer-fusion", "--dev-sessions"):
|
|
self.assertNotIn(forbidden, command)
|
|
with patch.object(runner.workflow.subprocess, "run") as run, contextlib.redirect_stdout(io.StringIO()):
|
|
runner.main(self.runner_args(extra=["--dry-run"]))
|
|
run.assert_not_called()
|
|
self.assertFalse((self.root / "checkpoints").exists())
|
|
|
|
def test_runner_rejects_wrong_marker_missing_test_input_and_existing_results(self):
|
|
self.export()
|
|
args = runner.parse_args(self.runner_args())
|
|
marker = self.output_root() / "complete.txt"
|
|
marker.write_text("12 clips\n")
|
|
with self.assertRaisesRegex(ValueError, "incomplete"):
|
|
runner.preflight(args)
|
|
marker.write_text("12 sessions\n")
|
|
_, output = runner.run_paths(args, 0)
|
|
output.mkdir(parents=True)
|
|
with self.assertRaises(FileExistsError):
|
|
runner.preflight(args)
|
|
output.rmdir()
|
|
(self.output_root() / "Session_1.pt").unlink()
|
|
with self.assertRaises(FileNotFoundError):
|
|
runner.preflight(args)
|
|
|
|
def test_cpu_round_trip_three_modalities_four_epochs_no_dev(self):
|
|
env = os.environ.copy()
|
|
env.update(OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1")
|
|
for modality, parts in MODALITIES.items():
|
|
with self.subTest(modality=modality):
|
|
self.export(modality)
|
|
result = subprocess.run(
|
|
[sys.executable, str(Path(runner.__file__).resolve()), *self.runner_args(modality)],
|
|
env=env, capture_output=True, text=True, timeout=120,
|
|
)
|
|
self.assertEqual(result.returncode, 0, result.stdout + result.stderr)
|
|
args = runner.parse_args(self.runner_args(modality))
|
|
checkpoint, output = runner.run_paths(args, 0)
|
|
saved = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
|
self.assertEqual(saved["dataset"], "easycom")
|
|
self.assertEqual(saved["epoch"], 4)
|
|
self.assertEqual(saved["train_sessions"], self.train)
|
|
self.assertEqual(saved["dev_sessions"], [])
|
|
self.assertIsNone(saved["local_contract"])
|
|
self.assertIsNone(saved["training_config"]["local_feature_root"])
|
|
width = sum(self.payloads["Session_1"]["config"][f"{part}_dim"] + 1 for part in parts)
|
|
self.assertEqual(saved["head_config"]["hidden_size"], width)
|
|
self.assertEqual(saved["head_config"]["local_visual_dim"], 0)
|
|
self.assertEqual(saved["head_config"]["local_audio_dim"], 0)
|
|
self.assertEqual(saved["head_config"]["layer_fusion"], "single")
|
|
self.assertFalse(any("local_projections" in key for key in saved["head"]))
|
|
metrics = json.loads((output / "metrics.json").read_text())
|
|
self.assertEqual(metrics["sessions"], self.test)
|
|
self.assertEqual(metrics["checkpoint_epoch"], 4)
|
|
for name in self.test:
|
|
self.assertEqual(np.load(output / "pred" / f"{name}.npy").shape, (20, 3, 10))
|
|
self.assertEqual({p.stem for p in (output / "pred").glob("*.npy")}, set(self.test))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|