From c5e2e044a8362761fd06653ae0bae3b5091a3cca Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Mon, 24 Aug 2026 00:27:42 +0800 Subject: [PATCH] Add dense-cache evaluation for Omni LoRA --- README.md | 42 +++++ eval_omni_turn_lora.py | 365 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 407 insertions(+) create mode 100644 eval_omni_turn_lora.py diff --git a/README.md b/README.md index 63afab6..6b60b3e 100644 --- a/README.md +++ b/README.md @@ -36,6 +36,7 @@ classification head. | `TurnPredictionHead` | Trained and evaluated | Map `[2048]` to `[10, 3]` logits | | `train_turn_head.py` | Verified on EasyCom | Train the frozen-Omni anticipation head with the official EasyCom split | | `infer_turn_head.py` | Verified on EasyCom | Save EgoSpeak-format predictions and reproduce per-horizon AP/mAP | +| `eval_omni_turn_lora.py` | Implemented, remote validation pending | Evaluate Thinker LoRA from dense AV caches with the same AP/mAP protocol | The full Thinker input path, embedded audio, direct video decoding, FlashAttention-2, 5 Hz streaming extraction, cache replay, target alignment, @@ -1030,6 +1031,47 @@ for example, change 8 GPUs with batch 2 and accumulation 1 to 4 GPUs with batch post-hoc continuation because a one-epoch cosine schedule has already reached zero. +### Evaluate a Thinker-LoRA checkpoint + +Use `eval_omni_turn_lora.py` for the formal EasyCom Table 2 evaluation. It +loads the trained Thinker LoRA and `10 x 3` turn head, reconstructs every causal +state from dense AV caches, and reports both the EgoSpeak full-timeline and +strict causal-valid protocols. Sparse training caches are rejected because +they do not contain predictions at every 200 ms timestamp. + +The following command evaluates the official EasyCom test sessions on three +GPUs, with one complete session assigned to each rank: + +```bash +CHECKPOINT=/data1/yfzx/wangx/Code/icassp2027-experiment/checkpoints/omni_turn_lora_easycom_sparse_8s_ddp8_epoch1 + +CUDA_VISIBLE_DEVICES=0,1,2 \ +OMP_NUM_THREADS=1 \ +MKL_NUM_THREADS=1 \ +OPENBLAS_NUM_THREADS=1 \ +PYTORCH_ALLOC_CONF=expandable_segments:True \ +TORCH_NCCL_ASYNC_ERROR_HANDLING=1 \ +torchrun --standalone --nproc_per_node=3 eval_omni_turn_lora.py \ + --model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \ + --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_3 \ + --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_2 \ + --cache-dir /data1/yfzx/wangx/Datasets/easycom/omni_lora_cache_dense_8s/Session_1 \ + --adapter-checkpoint "${CHECKPOINT}/adapter" \ + --head-checkpoint "${CHECKPOINT}/turn_head.pt" \ + --output-dir results/omni_turn_lora_easycom_sparse_8s_ddp8_epoch1 \ + --batch-size 2 \ + --block-cache-size 32 \ + --dtype bfloat16 \ + --flash-attention-2 \ + --ddp-timeout-minutes 60 +``` + +The evaluator saves per-session probabilities and valid masks under `pred/` +and `valid/`, and writes the aggregate metrics to `metrics.json`. Add +`--overwrite` only when intentionally replacing an existing result directory. +Treat Sessions 1, 2, and 3 as the official test set: inspect their metrics only +after the training recipe and hyperparameters have been fixed. + For a speed-only alternative, `train_omni_turn_lora.py` supports a two-GPU producer-consumer benchmark. The producer GPU runs the frozen AV towers and places current-tail embeddings in a bounded CPU queue while the consumer GPU diff --git a/eval_omni_turn_lora.py b/eval_omni_turn_lora.py new file mode 100644 index 0000000..9306d44 --- /dev/null +++ b/eval_omni_turn_lora.py @@ -0,0 +1,365 @@ +#!/usr/bin/env python3 +"""Evaluate Thinker LoRA + turn head 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, required=True) + 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) + 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() + 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, + "adapter_checkpoint": 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()