342 lines
13 KiB
Python
342 lines
13 KiB
Python
"""User-facing scorer.
|
||
|
||
Layering (mirrors FastVideo's VideoGenerator → Worker pattern, but
|
||
in-process)::
|
||
|
||
Evaluator ← user-facing
|
||
└── EvalWorker × N ← single-GPU; owns metric replicas
|
||
└── VideoPool ← async path-→-tensor prefetch (per evaluate call)
|
||
|
||
The constructor builds one :class:`EvalWorker` per GPU and loads every
|
||
metric on every worker eagerly. :meth:`evaluate` is the single entry
|
||
point: pass kwargs for one sample, or pass a list of sample dicts to
|
||
fan-out across GPU replicas with pipelined decoding — same method,
|
||
return type follows the input shape.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import threading
|
||
from collections.abc import Iterable
|
||
from typing import Any
|
||
|
||
from fastvideo.eval.registry import (list_metrics, missing_dependencies, resolve_group)
|
||
from fastvideo.eval.types import EvalResults, MetricResult
|
||
from fastvideo.eval.worker import EvalWorker
|
||
from fastvideo.logger import init_logger
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
class Evaluator:
|
||
"""Pre-initialized scorer for repeated evaluation.
|
||
|
||
Parameters
|
||
----------
|
||
metrics : list[str] | str
|
||
Metric names, group prefixes (``"vbench"``), or ``"all"``.
|
||
device : str
|
||
Single-GPU device (e.g. ``"cuda:0"``). Ignored when *num_gpus* > 1.
|
||
num_gpus : int
|
||
Number of GPU replicas. Each gets its own :class:`EvalWorker`.
|
||
compile : bool
|
||
Apply :func:`torch.compile` to each metric's ``_model``.
|
||
loader_threads : int
|
||
Background decode threads in the :class:`VideoPool`. Default 1
|
||
(hide decode behind compute). Bump for I/O-heavy benchmark sets
|
||
where one loader can't keep up with the workers.
|
||
prefetch_factor : int
|
||
``pool max_size = prefetch_factor * num_workers``. Default 2 —
|
||
one sample being consumed, one prefetched per worker.
|
||
pre_upload : bool
|
||
When ``True`` (default), the worker performs a single
|
||
host→device upload of ``video`` / ``reference`` per sample
|
||
before the metric loop, and every metric reads from that
|
||
shared GPU-resident tensor. Without it, each metric pays its
|
||
own ``.to(self.device)`` — N transfers of the same clip for N
|
||
metrics, which dominates at high resolution. Set ``False`` for
|
||
training-time eval, where keeping a clip resident on GPU
|
||
across the metric loop would fight the training step for VRAM.
|
||
skip_missing_deps : bool
|
||
When ``True``, silently drop explicit metric names whose
|
||
optional deps aren't importable (with a one-line warning per
|
||
skipped metric). Default ``False`` — an explicit name with a
|
||
missing dep raises :class:`ImportError` at construction time.
|
||
Group selectors (``"vbench"``, ``"all"``) always silent-skip
|
||
regardless of this flag.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
metrics: list[str] | str = "all",
|
||
device: str = "cuda:0",
|
||
num_gpus: int = 1,
|
||
compile: bool = False,
|
||
*,
|
||
loader_threads: int = 1,
|
||
prefetch_factor: int = 2,
|
||
pre_upload: bool = True,
|
||
skip_missing_deps: bool = False,
|
||
) -> None:
|
||
names = _resolve_metric_names(metrics, skip_missing_deps=skip_missing_deps)
|
||
if num_gpus > 1:
|
||
self._workers = [
|
||
EvalWorker(names,
|
||
f"cuda:{i}",
|
||
compile=compile,
|
||
pre_upload=pre_upload,
|
||
skip_missing_deps=skip_missing_deps) for i in range(num_gpus)
|
||
]
|
||
else:
|
||
self._workers = [
|
||
EvalWorker(names, device, compile=compile, pre_upload=pre_upload, skip_missing_deps=skip_missing_deps)
|
||
]
|
||
self._loader_threads = max(1, loader_threads)
|
||
self._prefetch_factor = max(1, prefetch_factor)
|
||
|
||
@property
|
||
def num_gpus(self) -> int:
|
||
return len(self._workers)
|
||
|
||
@property
|
||
def metric_names(self) -> list[str]:
|
||
return self._workers[0].metric_names
|
||
|
||
def evaluate(
|
||
self,
|
||
samples: Iterable[dict] | None = None,
|
||
*,
|
||
metrics: list[str] | None = None,
|
||
**kwargs,
|
||
) -> dict[str, MetricResult] | EvalResults:
|
||
"""Score one sample (kwargs form) or many samples (list form).
|
||
|
||
Both forms go through the same :class:`VideoPool` pipeline;
|
||
``video`` / ``reference`` paths are decoded asynchronously.
|
||
|
||
Parameters
|
||
----------
|
||
samples :
|
||
Iterable of sample dicts. Omit and pass kwargs for a
|
||
single-sample call.
|
||
metrics :
|
||
Subset of this Evaluator's registered metrics to actually
|
||
run on this batch. ``None`` (default) runs all registered.
|
||
Lets a single long-lived Evaluator score different (gen,
|
||
ref) corpora with different metric subsets across multiple
|
||
``evaluate()`` calls — e.g. LPIPS on a paired corpus, FVD
|
||
on an unequal-cardinality corpus — without burning model
|
||
loads. Set-metric accumulators are reset only for the
|
||
metrics included in *metrics*, so state for other set
|
||
metrics is preserved across calls.
|
||
|
||
Single sample::
|
||
|
||
ev.evaluate(video=tensor, text_prompt="...", fps=24.0)
|
||
|
||
Many samples::
|
||
|
||
ev.evaluate(samples=[{"video": ..., "reference": ...}, ...])
|
||
|
||
Many samples with a metric filter::
|
||
|
||
ev.evaluate(samples=lpips_samples, metrics=["common.lpips"])
|
||
ev.evaluate(samples=fvd_samples, metrics=["common.fvd"])
|
||
|
||
Returns
|
||
-------
|
||
dict[str, MetricResult] for the single-sample form;
|
||
:class:`EvalResults` (list-of-dict subclass with ``.corpus``) for
|
||
the list form.
|
||
"""
|
||
if metrics is not None:
|
||
unknown = [m for m in metrics if m not in self.metric_names]
|
||
if unknown:
|
||
raise ValueError(f"metrics filter contains names not registered on this Evaluator: "
|
||
f"{unknown}; registered: {self.metric_names}")
|
||
|
||
single = samples is None
|
||
sample_list: list[dict] = [kwargs] if samples is None else list(samples)
|
||
if not sample_list:
|
||
return EvalResults(samples=[], corpus={})
|
||
|
||
if single:
|
||
set_names = self._workers[0].set_metrics().keys()
|
||
active_set = (set_names if metrics is None else (set_names & set(metrics)))
|
||
if active_set:
|
||
# Set metrics need a population. A single sample can't produce a
|
||
# meaningful corpus result, and silently discarding it (return
|
||
# ``per_sample[0]`` only) hides the no-op. Force the list form.
|
||
raise ValueError("Set-vs-set metrics require samples=[...] with >=2 entries; "
|
||
"the kwargs form (single sample) cannot produce a corpus "
|
||
f"result. Active set metrics: {sorted(active_set)}")
|
||
|
||
per_sample, corpus = self._run(sample_list, metric_filter=metrics)
|
||
|
||
if single:
|
||
return per_sample[0]
|
||
return EvalResults(samples=per_sample, corpus=corpus)
|
||
|
||
def _run(
|
||
self,
|
||
samples: list[dict],
|
||
*,
|
||
metric_filter: list[str] | None = None,
|
||
) -> tuple[list[dict[str, MetricResult]], dict[str, MetricResult]]:
|
||
"""Pool-driven sample pipeline + set-metric finalize.
|
||
|
||
``metric_filter`` (when not ``None``) restricts both the worker
|
||
dispatch and the set-metric finalize to that subset of names.
|
||
Set metrics outside the filter keep their previous state.
|
||
"""
|
||
from fastvideo.eval.pool import VideoPool
|
||
|
||
filter_set: set[str] | None = set(metric_filter) if metric_filter is not None else None
|
||
|
||
# Reset only the set metrics that will actually run this batch —
|
||
# other set metrics keep state across calls (e.g. FVD reference
|
||
# features built by an earlier call still apply to this call).
|
||
for w in self._workers:
|
||
for name, m in w.set_metrics().items():
|
||
if filter_set is None or name in filter_set:
|
||
m.reset()
|
||
|
||
n_workers = len(self._workers)
|
||
max_size = self._prefetch_factor * n_workers
|
||
per_sample: list[Any] = [None] * len(samples)
|
||
|
||
with VideoPool(samples, loader_threads=self._loader_threads, max_size=max_size) as pool:
|
||
if n_workers == 1:
|
||
while True:
|
||
item = pool.get()
|
||
if item is None:
|
||
break
|
||
idx, decoded = item
|
||
per_sample[idx] = self._workers[0].evaluate(metrics=metric_filter, **decoded)
|
||
else:
|
||
# Multi-GPU: every worker drains the shared pool (work-stealing).
|
||
errors: list[BaseException] = []
|
||
threads: list[threading.Thread] = []
|
||
for w in self._workers:
|
||
t = threading.Thread(target=self._consumer_loop,
|
||
args=(w, pool, per_sample, errors, metric_filter),
|
||
daemon=True)
|
||
t.start()
|
||
threads.append(t)
|
||
for t in threads:
|
||
t.join()
|
||
if errors:
|
||
raise errors[0]
|
||
|
||
# Finalize only the set metrics in the filter. Merge per-worker
|
||
# state into worker 0, then finalize once.
|
||
corpus: dict[str, MetricResult] = {}
|
||
base_set = self._workers[0].set_metrics()
|
||
for name, m in base_set.items():
|
||
if filter_set is not None and name not in filter_set:
|
||
continue
|
||
for w in self._workers[1:]:
|
||
base_set[name].merge_from(w.set_metrics()[name])
|
||
corpus[name] = m.finalize()
|
||
|
||
return per_sample, corpus
|
||
|
||
@staticmethod
|
||
def _consumer_loop(
|
||
worker: EvalWorker,
|
||
pool: Any,
|
||
results: list,
|
||
errors: list,
|
||
metric_filter: list[str] | None,
|
||
) -> None:
|
||
try:
|
||
while True:
|
||
item = pool.get()
|
||
if item is None:
|
||
return
|
||
idx, decoded = item
|
||
results[idx] = worker.evaluate(metrics=metric_filter, **decoded)
|
||
except BaseException as e: # noqa: BLE001 — surface to parent thread via shared list
|
||
errors.append(e)
|
||
|
||
def release_cuda_memory(self) -> None:
|
||
"""Free CUDA caches on every replica without dropping models."""
|
||
for w in self._workers:
|
||
w.release_cuda_memory()
|
||
|
||
def unload(self) -> None:
|
||
"""Drop metric refs on every replica. Reverse with :meth:`reload`."""
|
||
for w in self._workers:
|
||
w.unload()
|
||
|
||
def reload(self) -> None:
|
||
"""Rebuild metrics dropped by :meth:`unload`."""
|
||
for w in self._workers:
|
||
w.reload()
|
||
|
||
def shutdown(self) -> None:
|
||
"""No-op; kept for API compatibility with older callers."""
|
||
|
||
|
||
def create_evaluator(
|
||
metrics: list[str] | str = "all",
|
||
device: str = "cuda:0",
|
||
num_gpus: int = 1,
|
||
compile: bool = False,
|
||
*,
|
||
skip_missing_deps: bool = False,
|
||
) -> Evaluator:
|
||
return Evaluator(metrics=metrics,
|
||
device=device,
|
||
num_gpus=num_gpus,
|
||
compile=compile,
|
||
skip_missing_deps=skip_missing_deps)
|
||
|
||
|
||
def _resolve_metric_names(metrics: list[str] | str, *, skip_missing_deps: bool = False) -> list[str]:
|
||
"""Resolve metric names, supporting groups (``"vbench"``) and ``"all"``.
|
||
|
||
Group / ``"all"`` selectors always silently skip metrics whose declared
|
||
dependencies aren't importable in this environment, with a single
|
||
warning per skipped metric.
|
||
|
||
Explicit names (e.g. ``"vbench.color"``) raise :class:`ImportError` at
|
||
construction time by default — the user asked for this metric specifically,
|
||
so a missing dep is a real error. Pass ``skip_missing_deps=True`` to opt
|
||
into the same silent-skip behavior for explicit names too (useful when
|
||
you're enumerating "every metric that works in this venv" rather than
|
||
requiring each one).
|
||
"""
|
||
if metrics == "all":
|
||
return _filter_satisfied(list_metrics(), context="all")
|
||
if isinstance(metrics, str):
|
||
metrics = [metrics]
|
||
|
||
seen: set[str] = set()
|
||
names: list[str] = []
|
||
for m in metrics:
|
||
group = resolve_group(m)
|
||
if group is not None:
|
||
candidates = _filter_satisfied(group, context=m)
|
||
elif skip_missing_deps:
|
||
candidates = _filter_satisfied([m], context=m)
|
||
else:
|
||
candidates = [m]
|
||
for n in candidates:
|
||
if n not in seen:
|
||
seen.add(n)
|
||
names.append(n)
|
||
return names
|
||
|
||
|
||
def _filter_satisfied(names: list[str], *, context: str) -> list[str]:
|
||
"""Drop metrics with missing deps, warning once per skipped metric."""
|
||
keep: list[str] = []
|
||
for n in names:
|
||
missing = missing_dependencies(n)
|
||
if missing:
|
||
ctx = f"in group '{context}'" if context != n else "(skip_missing_deps=True)"
|
||
logger.warning("eval: skipping %s %s; missing dependency: %s", n, ctx, ", ".join(missing))
|
||
continue
|
||
keep.append(n)
|
||
return keep
|