388 lines
14 KiB
Python
388 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Evaluate a frozen or LoRA-adapted Thinker from dense Omni AV caches.
|
|
|
|
The script accepts one dense cache directory per EasyCom session. With
|
|
``torchrun``, complete sessions are assigned round-robin to ranks; rank zero
|
|
then loads the saved per-session predictions and reports the exact
|
|
EgoSpeak-compatible and strict-causal metrics used by ``infer_turn_head.py``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import gc
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any, Sequence
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.distributed as dist
|
|
from torch.utils.data import DataLoader, Dataset
|
|
from tqdm.auto import tqdm
|
|
|
|
from easycom_head import CLASS_NAMES, NUM_CLASSES, NUM_HORIZONS
|
|
from infer_turn_head import (
|
|
_align_to_egospeak_timeline,
|
|
_evaluate,
|
|
_load_head,
|
|
_print_metrics,
|
|
)
|
|
from omni_cached_av import (
|
|
SCHEMA_VERSION,
|
|
CachedBlockStore,
|
|
collate_cached_sequences,
|
|
forward_last_state,
|
|
load_json,
|
|
reconstruct_open_stream_sequence,
|
|
torch_load,
|
|
)
|
|
from train_omni_turn_lora import (
|
|
DistributedRun,
|
|
_distributed_barrier,
|
|
_initialize_distributed,
|
|
_load_base_model,
|
|
)
|
|
|
|
|
|
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--model-id", required=True)
|
|
parser.add_argument(
|
|
"--cache-dir",
|
|
type=Path,
|
|
action="append",
|
|
required=True,
|
|
help="Dense cache directory from cache_omni_av_pieces.py; repeat per session",
|
|
)
|
|
parser.add_argument(
|
|
"--adapter-checkpoint",
|
|
type=Path,
|
|
help=(
|
|
"PEFT adapter directory. Omit it to evaluate the frozen base "
|
|
"Thinker as a cache-path control."
|
|
),
|
|
)
|
|
parser.add_argument("--head-checkpoint", type=Path, required=True)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--batch-size", type=int, default=2)
|
|
parser.add_argument("--block-cache-size", type=int, default=32)
|
|
parser.add_argument(
|
|
"--dtype",
|
|
choices=("bfloat16", "float16"),
|
|
default="bfloat16",
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
default="cuda:0",
|
|
help="Single-GPU device; torchrun overrides this with LOCAL_RANK",
|
|
)
|
|
parser.add_argument("--flash-attention-2", action="store_true")
|
|
parser.add_argument("--ddp-timeout-minutes", type=float, default=60.0)
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
parser.add_argument("--no-progress", action="store_true")
|
|
return parser.parse_args(argv)
|
|
|
|
|
|
class DenseCacheSession(Dataset[dict[str, Any]]):
|
|
"""Reconstruct every 5-Hz causal anchor for one cached session."""
|
|
|
|
def __init__(self, cache_dir: Path, *, block_cache_size: int) -> None:
|
|
self.cache_dir = cache_dir.expanduser().resolve()
|
|
self.manifest = load_json(self.cache_dir / "manifest.json")
|
|
if int(self.manifest.get("schema_version", -1)) != SCHEMA_VERSION:
|
|
raise ValueError(
|
|
f"Unsupported cache schema in {self.cache_dir}: "
|
|
f"{self.manifest.get('schema_version')}"
|
|
)
|
|
policy = self.manifest.get("anchor_selection", {}).get("policy")
|
|
if policy != "dense":
|
|
raise ValueError(
|
|
f"Evaluation requires --anchor-policy dense, got {policy!r}: "
|
|
f"{self.cache_dir}"
|
|
)
|
|
self.session = str(self.manifest["session"])
|
|
self.static = torch_load(self.cache_dir / self.manifest["static_path"])
|
|
self.store = CachedBlockStore(
|
|
self.cache_dir,
|
|
max_blocks=block_cache_size,
|
|
)
|
|
self.anchors = list(self.manifest["anchors"])
|
|
target_path = Path(self.manifest["target_path"])
|
|
targets = np.load(target_path, allow_pickle=False)
|
|
if targets.ndim != 2 or targets.shape[1] != NUM_CLASSES:
|
|
raise ValueError(
|
|
f"Expected target [T,{NUM_CLASSES}], got {targets.shape}: "
|
|
f"{target_path}"
|
|
)
|
|
if len(targets) != int(self.manifest["target_steps"]):
|
|
raise ValueError(
|
|
f"Target length changed for {self.session}: {len(targets)} vs "
|
|
f"{self.manifest['target_steps']}"
|
|
)
|
|
self.targets = np.asarray(targets, dtype=np.float32)
|
|
indices = np.asarray(
|
|
[int(anchor["target_index"]) for anchor in self.anchors],
|
|
dtype=np.int64,
|
|
)
|
|
expected = np.arange(len(self.targets), dtype=np.int64)
|
|
if not np.array_equal(indices, expected):
|
|
raise ValueError(
|
|
f"{self.session} dense cache does not contain exactly one ordered "
|
|
"anchor for every target timestep"
|
|
)
|
|
self.label_indices = torch.from_numpy(indices.copy()).long()
|
|
|
|
def __len__(self) -> int:
|
|
return len(self.anchors)
|
|
|
|
def __getitem__(self, index: int) -> dict[str, Any]:
|
|
anchor = self.anchors[index]
|
|
embeddings, positions = reconstruct_open_stream_sequence(
|
|
store=self.store,
|
|
static=self.static,
|
|
anchor=anchor,
|
|
)
|
|
return {
|
|
"inputs_embeds": embeddings,
|
|
"position_ids": positions,
|
|
# Reuse the verified padding/readout collator. Evaluation ignores
|
|
# this zero-length placeholder after stacking it.
|
|
"targets": torch.empty(0, dtype=torch.float32),
|
|
"step_index": int(anchor["target_index"]),
|
|
}
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _predict_session(
|
|
*,
|
|
dataset: DenseCacheSession,
|
|
language_model: torch.nn.Module,
|
|
head: torch.nn.Module,
|
|
device: torch.device,
|
|
model_dtype: torch.dtype,
|
|
batch_size: int,
|
|
show_progress: bool,
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
|
loader = DataLoader(
|
|
dataset,
|
|
batch_size=batch_size,
|
|
shuffle=False,
|
|
num_workers=0,
|
|
pin_memory=True,
|
|
collate_fn=collate_cached_sequences,
|
|
)
|
|
current_probabilities = np.zeros(
|
|
(len(dataset.targets), NUM_HORIZONS, NUM_CLASSES),
|
|
dtype=np.float32,
|
|
)
|
|
seen = np.zeros(len(dataset.targets), dtype=bool)
|
|
progress = tqdm(
|
|
loader,
|
|
total=len(loader),
|
|
desc=f"evaluate {dataset.session}",
|
|
unit="batch",
|
|
disable=not show_progress,
|
|
)
|
|
for batch in progress:
|
|
inputs_embeds = batch["inputs_embeds"].to(
|
|
device=device,
|
|
dtype=model_dtype,
|
|
non_blocking=True,
|
|
)
|
|
position_ids = batch["position_ids"].to(device=device, non_blocking=True)
|
|
attention_mask = batch["attention_mask"].to(device=device, non_blocking=True)
|
|
readout_indices = batch["readout_indices"].to(device=device, non_blocking=True)
|
|
states = forward_last_state(
|
|
language_model,
|
|
inputs_embeds=inputs_embeds,
|
|
position_ids=position_ids,
|
|
attention_mask=attention_mask,
|
|
readout_indices=readout_indices,
|
|
)
|
|
probabilities = head(states.float()).softmax(dim=-1).cpu().numpy()
|
|
step_indices = batch["step_indices"].numpy()
|
|
current_probabilities[step_indices] = probabilities
|
|
seen[step_indices] = True
|
|
if not bool(seen.all()):
|
|
missing = np.flatnonzero(~seen)
|
|
raise RuntimeError(
|
|
f"{dataset.session} missed {len(missing)} dense anchors; "
|
|
f"first={missing[:10].tolist()}"
|
|
)
|
|
return _align_to_egospeak_timeline(
|
|
current_probabilities,
|
|
dataset.label_indices,
|
|
len(dataset.targets),
|
|
)
|
|
|
|
|
|
def _load_session_specs(
|
|
cache_dirs: Sequence[Path], block_cache_size: int
|
|
) -> list[DenseCacheSession]:
|
|
sessions = [
|
|
DenseCacheSession(path, block_cache_size=block_cache_size)
|
|
for path in cache_dirs
|
|
]
|
|
names = [session.session for session in sessions]
|
|
if len(names) != len(set(names)):
|
|
raise ValueError(f"Duplicate cache sessions: {names}")
|
|
contracts = {
|
|
(
|
|
item.manifest["model_id"],
|
|
item.manifest["memory_seconds"],
|
|
item.manifest["video_fps"],
|
|
item.manifest["positions_per_chunk"],
|
|
item.manifest["hidden_size"],
|
|
)
|
|
for item in sessions
|
|
}
|
|
if len(contracts) != 1:
|
|
raise ValueError("Evaluation caches do not share one model/timing contract")
|
|
return sessions
|
|
|
|
|
|
def _run(args: argparse.Namespace, run: DistributedRun) -> None:
|
|
if args.batch_size <= 0 or args.block_cache_size <= 0:
|
|
raise ValueError("Batch and block-cache sizes must be positive")
|
|
if run.world_size > len(args.cache_dir):
|
|
raise ValueError(
|
|
f"Requested {run.world_size} ranks for only {len(args.cache_dir)} sessions"
|
|
)
|
|
output_dir = args.output_dir.expanduser().resolve()
|
|
metrics_path = output_dir / "metrics.json"
|
|
if metrics_path.exists() and not args.overwrite:
|
|
raise FileExistsError(
|
|
f"Evaluation output exists: {metrics_path}; pass --overwrite"
|
|
)
|
|
if run.is_main:
|
|
(output_dir / "pred").mkdir(parents=True, exist_ok=True)
|
|
(output_dir / "valid").mkdir(parents=True, exist_ok=True)
|
|
_distributed_barrier(run)
|
|
|
|
sessions = _load_session_specs(args.cache_dir, args.block_cache_size)
|
|
assigned = [
|
|
session
|
|
for index, session in enumerate(sessions)
|
|
if index % run.world_size == run.rank
|
|
]
|
|
print(
|
|
f"[rank {run.rank}/{run.world_size}] sessions="
|
|
+ ",".join(item.session for item in assigned),
|
|
flush=True,
|
|
)
|
|
|
|
base_model, processor = _load_base_model(args, run)
|
|
base_model.requires_grad_(False)
|
|
adapter_path = None
|
|
if args.adapter_checkpoint is None:
|
|
language_model = base_model.model.eval()
|
|
model_variant = "frozen_base_thinker"
|
|
print(
|
|
f"[rank {run.rank}] frozen control: no LoRA adapter loaded",
|
|
flush=True,
|
|
)
|
|
else:
|
|
try:
|
|
from peft import PeftModel
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"PEFT is required to evaluate the LoRA adapter"
|
|
) from exc
|
|
adapter_path = args.adapter_checkpoint.expanduser().resolve()
|
|
if not (adapter_path / "adapter_config.json").is_file():
|
|
raise FileNotFoundError(f"Missing adapter config: {adapter_path}")
|
|
language_model = PeftModel.from_pretrained(
|
|
base_model.model,
|
|
str(adapter_path),
|
|
is_trainable=False,
|
|
).eval()
|
|
model_variant = "thinker_lora"
|
|
language_model.config.use_cache = False
|
|
model_dtype = next(language_model.parameters()).dtype
|
|
head, head_payload = _load_head(args.head_checkpoint, run.device)
|
|
hidden_size = int(head.hidden_size)
|
|
for session in sessions:
|
|
if int(session.manifest["hidden_size"]) != hidden_size:
|
|
raise ValueError(
|
|
f"{session.session} hidden size {session.manifest['hidden_size']} "
|
|
f"!= head {hidden_size}"
|
|
)
|
|
del base_model
|
|
del processor
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
|
|
for session in assigned:
|
|
predictions, valid = _predict_session(
|
|
dataset=session,
|
|
language_model=language_model,
|
|
head=head,
|
|
device=run.device,
|
|
model_dtype=model_dtype,
|
|
batch_size=args.batch_size,
|
|
show_progress=not args.no_progress,
|
|
)
|
|
np.save(output_dir / "pred" / f"{session.session}.npy", predictions)
|
|
np.save(output_dir / "valid" / f"{session.session}.npy", valid)
|
|
print(
|
|
f"[rank {run.rank}] {session.session}: pred={predictions.shape}, "
|
|
f"valid={valid.sum(axis=0).tolist()}",
|
|
flush=True,
|
|
)
|
|
|
|
_distributed_barrier(run)
|
|
if run.is_main:
|
|
session_predictions = {
|
|
item.session: np.load(
|
|
output_dir / "pred" / f"{item.session}.npy",
|
|
allow_pickle=False,
|
|
)
|
|
for item in sessions
|
|
}
|
|
session_valid = {
|
|
item.session: np.load(
|
|
output_dir / "valid" / f"{item.session}.npy",
|
|
allow_pickle=False,
|
|
)
|
|
for item in sessions
|
|
}
|
|
session_targets = {item.session: item.targets for item in sessions}
|
|
metrics = _evaluate(session_predictions, session_valid, session_targets)
|
|
metrics.update(
|
|
{
|
|
"sessions": [item.session for item in sessions],
|
|
"classes": list(CLASS_NAMES),
|
|
"prediction_layout": "[target_time, class, future_horizon]",
|
|
"model_id": args.model_id,
|
|
"model_variant": model_variant,
|
|
"adapter_checkpoint": (
|
|
None if adapter_path is None else str(adapter_path)
|
|
),
|
|
"head_checkpoint": str(args.head_checkpoint.expanduser().resolve()),
|
|
"head_training_step": head_payload.get("training_step"),
|
|
"cache_dirs": [str(item.cache_dir) for item in sessions],
|
|
"world_size": run.world_size,
|
|
"batch_size_per_gpu": args.batch_size,
|
|
}
|
|
)
|
|
metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
|
|
_print_metrics(metrics)
|
|
print(f"\nmetrics: {metrics_path}")
|
|
print(f"predictions: {output_dir / 'pred'}")
|
|
_distributed_barrier(run)
|
|
|
|
|
|
def main(argv: Sequence[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
run = _initialize_distributed(args)
|
|
try:
|
|
_run(args, run)
|
|
finally:
|
|
if run.enabled and dist.is_initialized():
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|