180 lines
6.3 KiB
Python
180 lines
6.3 KiB
Python
"""User-facing scorer.
|
||
|
||
Layering (mirrors FastVideo's VideoGenerator → Worker pattern, but
|
||
in-process)::
|
||
|
||
Evaluator ← user-facing; round-robins samples across workers
|
||
└── EvalWorker × N ← single-GPU; owns metric replicas
|
||
|
||
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 — same method, return type follows the
|
||
input shape.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from collections.abc import Iterable
|
||
|
||
from fastvideo.eval.registry import (list_metrics, missing_dependencies, resolve_group)
|
||
from fastvideo.eval.types import 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``.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
metrics: list[str] | str = "all",
|
||
device: str = "cuda:0",
|
||
num_gpus: int = 1,
|
||
compile: bool = False,
|
||
) -> None:
|
||
names = _resolve_metric_names(metrics)
|
||
if num_gpus > 1:
|
||
self._workers = [EvalWorker(names, f"cuda:{i}", compile=compile) for i in range(num_gpus)]
|
||
else:
|
||
self._workers = [EvalWorker(names, device, compile=compile)]
|
||
self._pool = (ThreadPoolExecutor(max_workers=num_gpus) if num_gpus > 1 else None)
|
||
|
||
@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,
|
||
**kwargs,
|
||
) -> dict[str, MetricResult] | list[dict[str, MetricResult]]:
|
||
"""Score one sample (kwargs form) or many samples (list form).
|
||
|
||
``video`` and ``reference`` may be either a pre-loaded
|
||
``(T, C, H, W)`` tensor or a path-like (``str`` / ``Path``).
|
||
Paths are decoded inside the worker thread that picks up the
|
||
sample, so memory stays bounded by ``num_gpus`` even when
|
||
thousands of paths are queued — see ``score_folder.py`` for
|
||
the canonical pattern.
|
||
|
||
One sample::
|
||
|
||
ev.evaluate(video=tensor, text_prompt="...", fps=24.0)
|
||
ev.evaluate(video="path/to/clip.mp4", fps=24.0)
|
||
|
||
Many samples — fan out across GPU replicas, results in input order::
|
||
|
||
ev.evaluate(samples=[
|
||
{"video": "a.mp4", "reference": "ref_a.mp4"},
|
||
{"video": "b.mp4", "reference": "ref_b.mp4"},
|
||
...
|
||
])
|
||
|
||
Multi-GPU dispatch fires automatically iff ``num_gpus > 1`` *and*
|
||
the list form is used. The kwargs form always runs on worker 0;
|
||
if you have a single sample but multiple GPUs, wrap it in a
|
||
one-element list to use the pool, or just accept that a single
|
||
call uses one GPU — that's fine.
|
||
"""
|
||
if samples is None:
|
||
return self._workers[0].evaluate(**kwargs)
|
||
|
||
samples = list(samples)
|
||
if self._pool is None or len(samples) <= 1:
|
||
return [self._workers[0].evaluate(**s) for s in samples]
|
||
|
||
n = len(self._workers)
|
||
# Round-robin: worker i handles samples i, i+n, i+2n, ...
|
||
futures = [self._pool.submit(self._workers[idx % n].evaluate, **sample) for idx, sample in enumerate(samples)]
|
||
return [f.result() for f in futures]
|
||
|
||
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:
|
||
"""Tear down the worker thread pool. Idempotent."""
|
||
if self._pool is not None:
|
||
self._pool.shutdown(wait=True)
|
||
self._pool = None
|
||
|
||
|
||
def create_evaluator(
|
||
metrics: list[str] | str = "all",
|
||
device: str = "cuda:0",
|
||
num_gpus: int = 1,
|
||
compile: bool = False,
|
||
) -> Evaluator:
|
||
return Evaluator(metrics=metrics, device=device, num_gpus=num_gpus, compile=compile)
|
||
|
||
|
||
def _resolve_metric_names(metrics: list[str] | str) -> list[str]:
|
||
"""Resolve metric names, supporting groups (``"vbench"``) and ``"all"``.
|
||
|
||
Group / ``"all"`` selectors 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"``)
|
||
always pass through unchanged — the missing dep then surfaces as
|
||
:class:`ImportError` at construction time, which is what the user
|
||
asked for.
|
||
"""
|
||
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)
|
||
candidates = _filter_satisfied(group, context=m) if group is not None else [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 from a group expansion."""
|
||
keep: list[str] = []
|
||
for n in names:
|
||
missing = missing_dependencies(n)
|
||
if missing:
|
||
logger.warning(
|
||
"eval: skipping %s in group '%s'; missing dependency: %s. "
|
||
"Install instructions: pass the metric name explicitly to see them.", n, context, ", ".join(missing))
|
||
continue
|
||
keep.append(n)
|
||
return keep
|