188 lines
7.8 KiB
Python
188 lines
7.8 KiB
Python
"""EvalWorker: a single-GPU bag of metric replicas.
|
|
|
|
One ``EvalWorker`` per GPU. Holds an instance of every requested metric on
|
|
its own device, scores one sample at a time. Stateless across calls — the
|
|
worker never sees more than one ``evaluate(...)`` invocation at a time
|
|
(the parent :class:`Evaluator` ensures this by handing the worker to one
|
|
thread at a time).
|
|
|
|
This mirrors FastVideo's ``Worker`` layer under :class:`VideoGenerator`,
|
|
but in-process (threads, not processes) — eval metrics are independent
|
|
and need no NCCL / TP / SP / distributed init, so process isolation is
|
|
unnecessary overhead.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import contextlib
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from fastvideo.eval.memory import clear_cache
|
|
from fastvideo.eval.metrics.base import BaseMetric
|
|
from fastvideo.eval.registry import get_metric
|
|
from fastvideo.eval.types import MetricResult, Video
|
|
|
|
|
|
class EvalWorker:
|
|
"""Owns metric replicas on one device. Single-GPU, single-sample.
|
|
|
|
Per-sample metrics (``is_set_metric=False``) return one
|
|
:class:`MetricResult` per ``evaluate`` call. Set metrics
|
|
(``is_set_metric=True``) accumulate state on the worker's own
|
|
instance and contribute nothing to the per-sample return — the
|
|
Evaluator finalizes them after the pool drains.
|
|
"""
|
|
|
|
def __init__(self,
|
|
metric_names: list[str],
|
|
device: str,
|
|
*,
|
|
compile: bool = False,
|
|
pre_upload: bool = True,
|
|
skip_missing_deps: bool = False) -> None:
|
|
self._names = list(metric_names)
|
|
self._device = device
|
|
self._compile = compile
|
|
self._pre_upload = pre_upload
|
|
self._skip_missing_deps = skip_missing_deps
|
|
self._metrics: dict[str, BaseMetric] = {}
|
|
self._unloaded = False
|
|
self._load()
|
|
|
|
@property
|
|
def device(self) -> str:
|
|
return self._device
|
|
|
|
@property
|
|
def metric_names(self) -> list[str]:
|
|
"""Names of the metrics this worker owns, in load order."""
|
|
return list(self._metrics.keys())
|
|
|
|
def _load(self) -> None:
|
|
for name in self._names:
|
|
try:
|
|
m = get_metric(name)
|
|
m.to(self._device)
|
|
m.setup()
|
|
except (ImportError, ModuleNotFoundError) as e:
|
|
# Declared deps are checked up-front in get_metric; this
|
|
# catches transitive imports done lazily in setup() (e.g.
|
|
# vbench upstream pulling decord from its utils.py).
|
|
if self._skip_missing_deps:
|
|
import logging
|
|
logging.getLogger(__name__).warning("eval: skipping %s; setup-time import failed: %s", name, e)
|
|
continue
|
|
raise
|
|
if self._compile and getattr(m, "_model", None) is not None:
|
|
m._model = torch.compile(m._model)
|
|
self._metrics[name] = m
|
|
# Skip pre_upload when no loaded metric actually runs on GPU
|
|
# (e.g. an all-physics_iq run) — it would just pay a wasted
|
|
# host→device→host round-trip per sample.
|
|
self._any_needs_gpu = any(m.needs_gpu for m in self._metrics.values())
|
|
self._unloaded = False
|
|
|
|
def evaluate(self, *, metrics: list[str] | None = None, **kwargs) -> dict[str, MetricResult]:
|
|
"""Score one already-decoded sample.
|
|
|
|
Inputs (``video`` / ``reference`` paths) are decoded and
|
|
normalized upstream by the :class:`VideoPool`. The worker
|
|
handles the optional pre-upload to its device and dispatches to
|
|
each metric's ``compute`` or ``accumulate``.
|
|
|
|
Samples tagged ``role="reference"`` are corpus context for set
|
|
metrics only — per-sample metrics skip them so a mixed
|
|
Evaluator (e.g. FVD + LPIPS + vbench) doesn't produce spurious
|
|
per-sample scores on the reference half of the corpus.
|
|
|
|
``metrics`` (when not ``None``) further restricts dispatch to
|
|
that subset of registered names — used by the per-call
|
|
``Evaluator.evaluate(metrics=...)`` filter.
|
|
"""
|
|
if self._unloaded:
|
|
raise RuntimeError("EvalWorker was unloaded; call reload() before evaluating.")
|
|
|
|
sample = dict(kwargs)
|
|
# Unwrap ``Video`` → its ``.frames`` tensor for per-sample metric
|
|
# consumers (PSNR/SSIM/LPIPS/optical_flow). The pool decodes Video
|
|
# instances but leaves the wrapper in place so source-aware metrics
|
|
# (audio.imagebind_score) can still read ``Video.source``; those
|
|
# metrics should additionally pass ``sample["video_path"]`` if they
|
|
# need the path. Per-sample tensor consumers get a tensor either way.
|
|
for _k in ("video", "reference"):
|
|
_v = sample.get(_k)
|
|
if isinstance(_v, Video):
|
|
sample[_k] = _v.frames
|
|
if self._pre_upload and self._any_needs_gpu:
|
|
sample["video"] = _to_device(sample.get("video"), self._device)
|
|
if "reference" in sample:
|
|
sample["reference"] = _to_device(sample["reference"], self._device)
|
|
|
|
is_ref = sample.get("role") == "reference"
|
|
filter_set: set[str] | None = set(metrics) if metrics is not None else None
|
|
results: dict[str, MetricResult] = {}
|
|
broken: list[str] = []
|
|
for name, m in self._metrics.items():
|
|
if filter_set is not None and name not in filter_set:
|
|
continue
|
|
try:
|
|
if m.is_set_metric:
|
|
m.accumulate(sample)
|
|
elif not is_ref:
|
|
results[name] = m.compute(sample)
|
|
except (ImportError, ModuleNotFoundError, FileNotFoundError) as e:
|
|
# In skip_missing_deps mode, compute/accumulate-time failures
|
|
# whose root cause is a missing dependency (lazy import of a
|
|
# lib like torchcodec) or a missing resource (model checkpoint
|
|
# not on disk) drop the metric for the remainder of this
|
|
# Evaluator instead of bringing the whole run down. Strict
|
|
# mode re-raises. Programmer bugs, OOM, and other runtime
|
|
# failures are intentionally NOT caught here — they surface.
|
|
if not self._skip_missing_deps:
|
|
raise
|
|
import logging
|
|
logging.getLogger(__name__).exception("eval: dropping %s after %s: %s", name, type(e).__name__, e)
|
|
broken.append(name)
|
|
for n in broken:
|
|
self._metrics.pop(n, None)
|
|
return results
|
|
|
|
def set_metrics(self) -> dict[str, BaseMetric]:
|
|
"""Return ``{name: instance}`` for set metrics on this worker."""
|
|
return {n: m for n, m in self._metrics.items() if m.is_set_metric}
|
|
|
|
def reset_set_metrics(self) -> None:
|
|
"""Clear accumulator state on every set metric."""
|
|
for m in self._metrics.values():
|
|
if m.is_set_metric:
|
|
m.reset()
|
|
|
|
def release_cuda_memory(self) -> None:
|
|
"""Free CUDA caches without dropping models."""
|
|
clear_cache()
|
|
if torch.cuda.is_available():
|
|
with contextlib.suppress(Exception):
|
|
torch.cuda.ipc_collect()
|
|
|
|
def unload(self) -> None:
|
|
"""Drop metric refs so models become GC-able. Reverse with reload()."""
|
|
self._metrics = {}
|
|
self._unloaded = True
|
|
self.release_cuda_memory()
|
|
|
|
def reload(self) -> None:
|
|
"""Rebuild metrics dropped by :meth:`unload`."""
|
|
if self._unloaded:
|
|
self._load()
|
|
|
|
|
|
def _to_device(value: Any, device: str | torch.device) -> Any:
|
|
"""Move *value* to *device* if it's a tensor not already there."""
|
|
if value is None or not isinstance(value, torch.Tensor):
|
|
return value
|
|
target = torch.device(device)
|
|
if value.device == target:
|
|
return value
|
|
return value.to(target, non_blocking=True)
|