Pure formatting pass — no behavior changes. Most edits are one of: * yapf splitting / re-flowing kwarg-heavy callsites (argparse setup, ``MetricResult(...)`` construction). * ruff auto-fixes around imports (collapsing ``Iterable`` to ``collections.abc``, removing the ``typing.Tuple`` shim where ``tuple[...]`` works, dropping unused imports). * ``zip(..., strict=False)`` added to every two-iterable zip call. * ``getattr(np, "trapz")`` → ``np.trapz`` where the fallback can read the attribute directly (the ``trapezoid`` lookup still uses ``getattr`` because the name itself is the moving target). Verified the test suite still passes after the pass.
19 KiB
Porting Eval Metrics into fastvideo.eval
This guide is for contributors adding new evaluation metrics to FastVideo's
eval suite. If you just want to run the existing metrics, see
fastvideo/eval/README.md.
When you need this guide
Use this guide when you are:
- Adding a new metric (native or wrapping a third-party library).
- Porting a benchmark (e.g. VBench, MIND, EvalCrafter) whose Python code needs to be importable from a pinned upstream.
- Adding a new metric group (audio, vlm, etc.).
TL;DR
Metrics are auto-discovered from
fastvideo/eval/metrics/<group>/<name>/metric.py. Each declares itself
with @register("<group>.<name>") and subclasses BaseMetric. Three
recipes:
- Native metric (pure-PyTorch, no submodule). Drop a file, declare
deps, implement
compute(sample). - Library-wrapped metric (CLIP, torch.hub, transformers, pyiqa).
Same as above, plus route the library's cache through
get_cache_dir()if it has adownload_root=/cache_dir=kwarg. - Upstream-submodule-wrapped metric (vbench-style). Pin upstream as
a git submodule under
fastvideo/third_party/eval/<bench>/. Adapter__init__.pydoes thesys.pathinsert + any runtime compat shims for modern dep versions. Nosetup.sh. No on-disk patches.
The full recipes are below.
0) Layout & auto-discovery
fastvideo/eval/metrics/
├── base.py # BaseMetric + lifecycle contract
├── common/ # group: SSIM, PSNR, LPIPS
├── optical_flow/ # group: gt_optical_flow, synthetic_optical_flow
├── audio/ # group: CLAP, WER, audiobox, …
├── vlm/ # group: VideoScore-2
├── physics_iq/ # group + sub-metrics
└── vbench/ # group: 16 sub-metrics
├── __init__.py # sys.path bootstrap + runtime compat shims
├── _grit_helper.py # shared upstream-touching helpers
└── <sub_metric>/metric.py
Auto-discovery (fastvideo/eval/metrics/__init__.py) walks each group
dir and imports every metric.py it finds, which fires the @register
decorators. Names starting with _ are skipped — use that prefix
for shared helpers or vendored code that shouldn't register itself.
1) The BaseMetric contract
Every metric subclasses fastvideo.eval.metrics.base.BaseMetric and
declares:
class YourMetric(BaseMetric):
name: str = "common.your_metric" # must match @register
requires_reference: bool = True # needs sample["reference"]
higher_is_better: bool = True # for ranking / aggregates
dependencies: list[str] = [] # importable module names;
# registry surfaces clean
# ImportError if missing
needs_gpu: bool = False
backbone: str | None = None # e.g. "clip_vit_l14"
You must implement:
def compute(self, sample: dict) -> list[MetricResult]:
"""sample['video'] is (1, T, C, H, W). Return a one-element list.
The leading 1 is preserved for forward-compat with batched eval;
today :class:`EvalWorker` always invokes metrics with B=1.
"""
You may override:
setup(self) -> None— eager model loading. Called once bycreate_evaluator. Idempotent (re-entrant). Use theif self._model is not None: returnpattern.to(self, device)— move the metric (and its submodels) todevice.
If a required input is missing (e.g. an fps-aware metric called without
fps), return self._skip(sample, reason) — do not raise.
2) Recipe A — Native metric (no external deps)
Smallest case. Pixel math, simple closed-form.
# fastvideo/eval/metrics/common/your_metric/metric.py
from __future__ import annotations
import torch
from fastvideo.eval.metrics.base import BaseMetric
from fastvideo.eval.registry import register
from fastvideo.eval.types import MetricResult
@register("common.your_metric")
class YourMetric(BaseMetric):
name = "common.your_metric"
requires_reference = True
higher_is_better = True
needs_gpu = False
dependencies: list[str] = [] # nothing extra
def compute(self, sample: dict) -> list[MetricResult]:
gen, ref = sample["video"], sample["reference"] # (B,T,C,H,W) each
per_video = ((gen - ref) ** 2).mean(dim=(1, 2, 3, 4)).sqrt()
return [
MetricResult(name=self.name, score=float(s), details={})
for s in per_video
]
That's it. Drop the file, the registry picks it up.
3) Recipe B — Library-wrapped metric (CLIP, torch.hub, transformers, pyiqa)
If your metric loads a backbone from a Python package, you are
responsible for pointing the library at the eval cache so users get
one knob (FASTVIDEO_EVAL_CACHE) to redirect everything.
Cache routing rules
| library | how to route | location after redirect |
|---|---|---|
clip.load("ViT-X") |
pass download_root=str(get_cache_dir() / "clip") |
${FASTVIDEO_EVAL_CACHE}/clip/ |
torch.hub.load(...) |
nothing — TORCH_HOME is redirected at fastvideo.eval import time |
${FASTVIDEO_EVAL_CACHE}/torch/hub/ |
transformers.from_pretrained(...) |
nothing — leave HF's default cache (~/.cache/huggingface/hub/) so users dedupe with other ML projects |
~/.cache/huggingface/hub/ |
huggingface_hub.snapshot_download / hf_hub_download |
use ensure_checkpoint(...) (it wraps these with filelock) |
same as above |
pyiqa.create_metric(...) |
no env var or kwarg honored — accept the spread, document in metric docstring | pyiqa-internal |
lpips, ptlflow |
torch.hub-based — auto-redirected | ${FASTVIDEO_EVAL_CACHE}/torch/hub/ |
| Raw URL (no HF Hub) | use ensure_checkpoint(name, source="https://...") |
${FASTVIDEO_EVAL_CACHE}/models/<name> |
Example — CLIP backbone + LAION head
# fastvideo/eval/metrics/your_group/your_metric/metric.py
from __future__ import annotations
import torch
import torch.nn as nn
from fastvideo.eval.metrics.base import BaseMetric
from fastvideo.eval.registry import register
from fastvideo.eval.types import MetricResult
@register("your_group.your_metric")
class YourMetric(BaseMetric):
name = "your_group.your_metric"
requires_reference = False
needs_gpu = True
dependencies = ["clip"] # "openai-clip" PyPI; importable as `clip`
def __init__(self) -> None:
super().__init__()
self._clip = None
self._head = None
def setup(self) -> None:
if self._clip is not None:
return
import clip
from fastvideo.eval.models import ensure_checkpoint, get_cache_dir
# Backbone: route CLIP's cache through our root.
self._clip, _ = clip.load(
"ViT-L/14",
device=self.device,
download_root=str(get_cache_dir() / "clip"),
)
self._clip.eval()
# URL-fetched head: ensure_checkpoint downloads to
# ${FASTVIDEO_EVAL_CACHE}/models/ with filelock + atomic rename.
ckpt = ensure_checkpoint(
"your_head.pth",
source="https://example.com/path/to/your_head.pth",
)
self._head = nn.Linear(768, 1)
self._head.load_state_dict(
torch.load(ckpt, map_location="cpu", weights_only=True)
)
self._head.to(self.device).eval()
def to(self, device):
super().to(device)
if self._clip is not None:
self._clip = self._clip.to(self.device)
if self._head is not None:
self._head = self._head.to(self.device)
return self
def compute(self, sample: dict) -> list[MetricResult]:
...
Don't reach into other ~/.cache/... dirs
If a third-party library hard-codes ~/.cache/<lib>/ and offers no
override, don't try to redirect it by setting os.environ or
patching os.path.expanduser. Just document the exception in the
metric's docstring. Forcing redirection is fragile and breaks user
expectations of where the library's cache lives.
4) Recipe C — Upstream-submodule-wrapped metric (vbench pattern)
Use this when the upstream benchmark ships Python code (vbench/,
MIND/, etc.) that's not pip-installable cleanly and we need to import
it. The pattern (see fastvideo/eval/metrics/vbench/__init__.py for the
full worked example):
4.1 Pin the upstream as a submodule
git submodule add <upstream-url> fastvideo/third_party/eval/<bench>
cd fastvideo/third_party/eval/<bench>
git checkout <pinned-sha>
cd -
git add .gitmodules fastvideo/third_party/eval/<bench>
The submodule pulls under the standard git submodule update --init --recursive flow that users run for kernel deps. No special opt-in
flag, no update = none, no setup script.
4.2 Bootstrap on sys.path
# fastvideo/eval/metrics/<bench>/__init__.py
from __future__ import annotations
import sys
from pathlib import Path
# fastvideo/eval/metrics/<bench>/__init__.py → ../../../../third_party/eval/<bench>
_UPSTREAM = Path(__file__).resolve().parents[3] / "third_party" / "eval" / "<bench>"
if _UPSTREAM.is_dir() and str(_UPSTREAM) not in sys.path:
sys.path.insert(0, str(_UPSTREAM))
We don't pip install the upstream — its egg-link/.pth would just
re-do this sys.path.insert. Skipping the install also skips the
upstream's setup.py (which often gates on a specific CUDA version).
4.3 Modern-dep compat: runtime shims, NOT on-disk patches
Upstream code pinned to e.g. transformers==4.33.2, numpy<2 typically
breaks against modern versions in 3-4 known places (API renames). Fix
those at import time, in the same __init__.py:
def _install_compat_shims() -> None:
# Example: transformers.modeling_utils API moved.
try:
import transformers.modeling_utils as _mu
import transformers.pytorch_utils as _pu
for _n in ("apply_chunking_to_forward",
"find_pruneable_heads_and_indices",
"prune_linear_layer"):
if not hasattr(_mu, _n) and hasattr(_pu, _n):
setattr(_mu, _n, getattr(_pu, _n))
except ImportError:
pass
# Example: numpy.lib.function_base.disp removed in numpy>=2.
try:
import types, numpy.lib as _nl
if not hasattr(_nl, "function_base"):
_stub = types.ModuleType("numpy.lib.function_base")
_stub.disp = lambda *a, **k: None
sys.modules["numpy.lib.function_base"] = _stub
_nl.function_base = _stub
except ImportError:
pass
_install_compat_shims()
For function-level patches that can't be expressed as attribute writes
(e.g. wrapping a model factory function), use a sys.meta_path finder
that wraps the loader. See the _install_modeling_finetune_hook()
helper in fastvideo/eval/metrics/vbench/__init__.py for the pattern —
about 30 lines.
Why shims and not git apply patches: patches go stale (new upstream
SHA → patches don't apply); shims are versioned Python code in our
repo, grep-able, and only ever run if the targeted module is being
imported.
4.4 Per-sub-metric files
Each sub-metric is a normal BaseMetric subclass that imports from the
upstream:
# fastvideo/eval/metrics/<bench>/<sub>/metric.py
from fastvideo.eval.metrics.base import BaseMetric
from fastvideo.eval.registry import register
from fastvideo.eval.types import MetricResult
@register("<bench>.<sub>")
class YourSubMetric(BaseMetric):
...
def setup(self) -> None:
if self._model is not None:
return
# The sys.path bootstrap fired when fastvideo.eval.metrics.<bench>
# was imported (which auto-discovery does before importing this
# sub-package). So upstream imports just work:
from <bench>.something import SomeModel
...
4.5 Conditional registration when upstream is missing
If a user installed fastvideo[eval] but didn't run git submodule update --init, <bench>.* metrics should not register. The
auto-discovery walker imports each sub-package's metric module; have
that import bail out cleanly:
# fastvideo/eval/metrics/<bench>/__init__.py — at the bottom
_AVAILABLE = (_UPSTREAM / "<bench>" / "__init__.py").is_file()
# fastvideo/eval/metrics/<bench>/<sub>/__init__.py
from fastvideo.eval.metrics.<bench> import _AVAILABLE
if _AVAILABLE:
from .metric import YourSubMetric # noqa
fastvideo eval list then reflects what the user actually has, not what
they could have.
4.6 Don't fix upstream bugs unless they block the metric
The upstream is pinned. If a metric works against the pinned SHA, leave it. If it actively breaks against modern deps (the import-drift cases above), shim it. Resist the urge to "improve" upstream code in a fastvideo-side fork — that's how patches go stale and parity drifts.
5) Model checkpoints — ensure_checkpoint
Use ensure_checkpoint(name, source, filename=None) for any non-package
weights. It resolves a local path, downloading on miss, with filelock
safety across processes and SLURM ranks.
source form |
what happens |
|---|---|
"/abs/path/to/file.pth" |
passthrough — returned unchanged |
"https://..." |
downloaded to ${FASTVIDEO_EVAL_CACHE}/models/<name> via huggingface_hub.http_get, atomic rename, filelock |
"org/repo" (no filename) |
snapshot_download(repo_id) → ~/.cache/huggingface/hub/ |
"org/repo" (with filename) |
hf_hub_download(repo_id, filename) → ~/.cache/huggingface/hub/ |
name is only used as the local filename for URL sources; HF sources
ignore it (HF manages its own cache key by content hash).
from fastvideo.eval.models import ensure_checkpoint
# URL — name matters
ckpt = ensure_checkpoint(
"amt-s.pth",
source="https://huggingface.co/lalala125/AMT/resolve/main/amt-s.pth",
)
# HF single file — name is decorative
ckpt = ensure_checkpoint(
"raft-things.pth", # ignored; HF cache uses repo+sha
source="OpenGVLab/VBench_Used_Models",
filename="raft-things.pth",
)
6) Declaring dependencies
Set dependencies = ["pkg1", "pkg2"] on your metric class (importable
module names, not PyPI distribution names). The registry checks each via
importlib.util.find_spec at instantiation time and raises a clean
ImportError pointing the user at pip install fastvideo[eval]:
class YourMetric(BaseMetric):
dependencies = ["clip", "timm"] # importable as `import clip`, `import timm`
If a dep is in [project.optional-dependencies.eval], you don't need to
do anything more. If it's a new dep, add it to that group in
pyproject.toml.
7) Common gotchas
- Do not write a
setup.shfor your benchmark. Submodule init via the standardgit submodule update --init --recursiveis enough; any modern-dep compat goes into your__init__.pyas runtime shims. - Do not modify upstream files on disk. The submodule should always
match its pinned SHA. Compat lives in our
__init__.py. - Do not pip-install the upstream. The egg-link is just a glorified
sys.path.insert— we do that directly in__init__.py. - Don't
import torch.hub.set_dir(...)from your metric. It's done globally infastvideo/eval/__init__.py. - Don't put cache-redirection env vars in your metric's
setup(). By the timesetup()runs, the library has likely already cached the default-location decision. Set env vars at package-init time only. - Skip rather than raise when an input is missing. Use
self._skip(sample, reason)for any expected-missing input — it returns a list ofMetricResult(score=None)so other metrics in the same evaluator continue. - Watch for upstream re-registration conflicts. If the upstream uses
a global registry (detectron2's
META_ARCH_REGISTRY, MMCV, etc.), loading the same model twice in the same process will throw. The evaluator already loads each metric once; if you write a custom setup-then-call pattern, mirror that single-load discipline.
8) Training-time eval — keep evaluators hot, free caches between calls
When wiring eval into a training loop, the working pattern is:
-
Construct the
Evaluatoronce, attach it to the pipeline (self._eval = create_evaluator(...)). Don't recreate it per validation round — that re-pays the model load cost. -
Save validation videos to disk (the diffusion path already does this). Pass paths to
evaluator.evaluate, not in-memory tensors that share GPU memory with the training model. -
Run validation only on rank 0 of each sequence-parallel group; gather paths from other ranks and let rank 0 score everything.
-
After every
evaluate(...)call, callevaluator.release_cuda_memory()in afinallyblock —gc.collect()+torch.cuda.empty_cache()+torch.cuda.ipc_collect(). The eval model stays loaded; only transient activation buffers from the just-finished call get freed::for video_path in batch: try: scores = self._eval.evaluate(video=load_video(video_path)) finally: self._eval.release_cuda_memory() -
If memory pressure spikes (rare on H200), call
evaluator.unload()to drop every metric reference and let the GPU memory be GC'd.unloadis reversible:evaluator.reload()rebuilds the same metrics with the original config (re-paying the model load cost). Callingevaluatebetweenunloadandreloadraises a clearRuntimeError.
This is the pattern in
fastvideo.training.{ptlflow,vbench}_validation (mhuo's branch).
For most metrics (sub-1 GB backbones, e.g. CLIP/DINO/RAFT/AMT) the
eval model can stay co-resident with the training model in
transformer.eval() mode without any swap. For larger ones
(VideoScore2 at 14 GB), measure first; if it fits on the rank-0 GPU
during validation (training model in eval mode = no grads/optimizer
updates), keep it hot. If not, unload between rounds.
9) Local verification
Native and library-wrapped metrics — single-GPU smoke is enough:
import torch
from fastvideo.eval import create_evaluator
ev = create_evaluator(metrics=["<group>.<your_metric>"], device="cuda")
video = torch.randn(1, 49, 3, 256, 256, device="cuda").clamp(0, 1)
print(ev.evaluate(video=video))
Submodule-wrapped metrics — also do a parity check against the upstream
once. The pattern from research/ (see commit history) gives a clean
recipe: clone upstream into a separate venv, run the same video through
both, expect EXACT match on bit-deterministic metrics and ≤1% drift on
backbone-heavy ones (driven by transformers/torch version differences).
For quick-and-dirty parity in CI: pin a tiny test video, record expected
scores ± tolerance, and add a calibration test under
fastvideo/tests/eval/.
10) When not to add a metric
- Set-vs-set distribution metrics (FVD, FID-style) don't fit
BaseMetric.compute(sample)cleanly — they need a population. Adding them requires a stateful accumulator interface that doesn't exist yet. Open an issue first. - Metrics requiring a single-GPU model > available memory. Eval is not the place for tensor-parallel sharding; metrics are expected to fit on one GPU.
- Metrics that need
mmcvwith conflicting CUDA ABI. Document the affected sub-metrics as unsupported and skip them — do not build isolation infrastructure (subprocess engine, per-metric venv).