coarse-grained-forecasting-benchmark's picture
Add anonymous reproduction code
41dd7e0 verified
Raw History Blame Contribute Delete
11.6 kB
"""The Clean-FID pipeline, pinned and checked.
Everything after the dataset adapter is fixed here for every dataset:
* images arrive as uint8 (N, H, W, 3) in [0, 255] (one fixed image range);
* a single resize to 299x299 with clean-fid's "clean" resizer (PIL bicubic, float output, antialiasing is
PIL's default for its resampling filters); the TorchScript Inception is built with resize_inside=False
and asserts 299x299 input, so a second resize cannot happen silently;
* the pinned Inception-v3 weights (inception-2015-12-05.pt, sha256 checked), pool3 features (2048-D);
* FID = clean-fid's frechet_distance formula; KID = clean-fid's kernel_distance (unbiased polynomial-kernel
MMD^2, num_subsets / max_subset_size from the evaluation settings).
``pins()`` returns the record that every FID manifest carries; ``check_pins(expected)`` raises on any drift.
"""
from __future__ import annotations
import hashlib
import importlib.metadata as _md
import os
from pathlib import Path
import numpy as np
INCEPTION_DIR = Path(os.environ.get("CLEAN_FID_CACHE", Path.home() / ".cache" / "clean-fid"))
INCEPTION_FILE = "inception-2015-12-05.pt"
IMAGE_SIZE = 299
# ── the spatial Inception feature of csFID (Nash et al., ICML 2021) ─────────────────────
# The Inception-v3 activation "mixed_6/conv" -- the FIRST branch of the mixed_6 Inception block, BEFORE its
# concatenation with tower/tower_1/tower_2 (openai/guided-diffusion evaluator.py FID_SPATIAL_NAME) -- is a
# 192 x 17 x 17 tensor; its first 7 channels are kept and flattened NHWC to 17*17*7 = 2023 components per image,
# without spatial average pooling (that pooling is what makes the 2048-D pool3 cFID feature layout-tolerant).
# In this TorchScript port the tensor is the submodule layers.mixed_6.conv; the module names come from the TF op
# names. The stepwise traversal used to read it reproduces layers.forward exactly on CPU (verify_spatial_tap).
SPATIAL_LAYER = "mixed_6/conv"
SPATIAL_CHANNELS = 7
SPATIAL_GRID = 17
SPATIAL_DIM = SPATIAL_GRID * SPATIAL_GRID * SPATIAL_CHANNELS # 2023
POOL3_DIM = 2048
_MIXED5_SHAPE = (768, SPATIAL_GRID, SPATIAL_GRID) # canonical input of the mixed_6 block
def _sha256(path: Path) -> str:
h = hashlib.sha256()
with open(path, "rb") as f:
for chunk in iter(lambda: f.read(1 << 22), b""):
h.update(chunk)
return h.hexdigest()
def pins() -> dict:
import PIL
import torch
p = INCEPTION_DIR / INCEPTION_FILE
return {"backend": "clean-fid", "clean_fid_version": _md.version("clean-fid"), "mode": "clean",
"resize": {"backend": "PIL", "filter": "bicubic", "size": IMAGE_SIZE, "once": True,
"implementation": "cleanfid.resize.build_resizer('clean') = make_resizer('PIL', False, 'bicubic', (299, 299))"},
"inception": {"file": INCEPTION_FILE, "path": str(p), "sha256": _sha256(p) if p.is_file() else None,
"feature_layer": "pool3", "feature_dim": POOL3_DIM, "resize_inside": False},
"spatial": {"feature_layer": SPATIAL_LAYER, "submodule": "layers.mixed_6.conv",
"tap": "first branch of the mixed_6 block, before concatenation with tower/tower_1/tower_2",
"channels_kept": SPATIAL_CHANNELS, "grid": [SPATIAL_GRID, SPATIAL_GRID], "feature_dim": SPATIAL_DIM,
"flatten": "NHWC", "spatial_pooling": "none",
"reference": "Nash et al., ICML 2021; openai/guided-diffusion evaluator.py FID_SPATIAL_NAME='mixed_6/conv:0'"},
"fid": "cleanfid.fid.frechet_distance (scipy sqrtm, eps 1e-6 fallback)",
"kid": "cleanfid.fid.kernel_distance (unbiased polynomial-kernel MMD^2)",
"torch": torch.__version__, "pillow": PIL.__version__, "numpy": np.__version__}
def check_pins(expected: dict) -> dict:
"""``expected`` = the fid settings (clean_fid_version, inception.sha256, inception.file). Raises on drift."""
got = pins()
probs = []
if str(expected.get("clean_fid_version")) != got["clean_fid_version"]:
probs.append(f"clean-fid {got['clean_fid_version']} != pinned {expected.get('clean_fid_version')}")
exp_inc = expected.get("inception") or {}
if exp_inc.get("file") and exp_inc["file"] != got["inception"]["file"]:
probs.append("inception file name differs")
if got["inception"]["sha256"] is None:
probs.append(f"Inception weights not found at {got['inception']['path']} (download once with cleanfid on a node with network)")
elif exp_inc.get("sha256") and exp_inc["sha256"] != got["inception"]["sha256"]:
probs.append(f"Inception weights sha256 {got['inception']['sha256'][:12]} != pinned {exp_inc['sha256'][:12]}")
if probs:
raise RuntimeError("FID pipeline pin violated: " + "; ".join(probs))
return got
def build_extractor(device):
"""clean-fid 'clean' extractor: TorchScript TF-Inception, weights from INCEPTION_DIR, no internal resize."""
from cleanfid.inception_torchscript import InceptionV3W
INCEPTION_DIR.mkdir(parents=True, exist_ok=True)
model = InceptionV3W(str(INCEPTION_DIR), download=not (INCEPTION_DIR / INCEPTION_FILE).is_file(),
resize_inside=False).to(device).eval()
return model
def inception_features(imgs_u8: np.ndarray, extractor, device, batch: int = 200) -> np.ndarray:
"""uint8 (N, H, W, 3) -> float32 (N, 2048): a single clean resize per image, then the pinned network."""
import torch
from cleanfid.resize import build_resizer
if imgs_u8.dtype != np.uint8 or imgs_u8.ndim != 4 or imgs_u8.shape[-1] != 3:
raise ValueError(f"adapter output must be uint8 (N, H, W, 3); got {imgs_u8.dtype} {imgs_u8.shape}")
resizer = build_resizer("clean")
out = []
with torch.no_grad():
for i in range(0, len(imgs_u8), int(batch)):
xs = [resizer(im) for im in imgs_u8[i:i + int(batch)]] # float32 (299, 299, 3) in [0, 255]
x = torch.from_numpy(np.stack(xs).transpose(0, 3, 1, 2).copy()).to(device)
if tuple(x.shape[-2:]) != (IMAGE_SIZE, IMAGE_SIZE):
raise AssertionError("resizer did not produce 299x299")
out.append(extractor(x).cpu().numpy())
return np.concatenate(out).astype(np.float32)
def fid_from_features(f1: np.ndarray, f2: np.ndarray) -> float:
from cleanfid.fid import frechet_distance
f1 = np.asarray(f1, np.float64); f2 = np.asarray(f2, np.float64)
return float(frechet_distance(f1.mean(0), np.cov(f1, rowvar=False), f2.mean(0), np.cov(f2, rowvar=False)))
def kid_from_features(f1: np.ndarray, f2: np.ndarray, num_subsets: int = 100, max_subset_size: int = 1000) -> float:
from cleanfid.fid import kernel_distance
return float(kernel_distance(np.asarray(f1, np.float64), np.asarray(f2, np.float64),
num_subsets=int(num_subsets), max_subset_size=int(max_subset_size)))
# ── csFID feature extraction ────────────────────────────────────────────────────────────────────────────────
def _pre_mixed6_stages(layers):
"""The layer sequence of layers.forward up to (not including) mixed_6 -- mirrors its TorchScript source."""
return [layers.conv, layers.conv_1, layers.conv_2, layers.pool0, layers.conv_3, layers.conv_4, layers.pool1,
layers.mixed, layers.mixed_1, layers.mixed_2, layers.mixed_3, layers.mixed_4, layers.mixed_5]
def _post_mixed6_stages(layers):
return [layers.mixed_6, layers.mixed_7, layers.mixed_8, layers.mixed_9, layers.mixed_10, layers.pool2]
def spatial_tap(extractor, x_normalized):
"""x already resized to 299x299 and normalized as InceptionV3W does ((x - 128) / 128) -> the mixed_6/conv
activation (N, 192, 17, 17)."""
L = extractor.layers
h = x_normalized
for st in _pre_mixed6_stages(L):
h = st.forward(h)
return L.mixed_6.conv.forward(h)
def spatial_features(imgs_u8: np.ndarray, extractor, device, batch: int = 200) -> np.ndarray:
"""uint8 (N, H, W, 3) -> float32 (N, 2023): the canonical sFID feature. Same clean resize and normalization
as inception_features; only the tapped tensor differs."""
import torch
from cleanfid.resize import build_resizer
if imgs_u8.dtype != np.uint8 or imgs_u8.ndim != 4 or imgs_u8.shape[-1] != 3:
raise ValueError(f"adapter output must be uint8 (N, H, W, 3); got {imgs_u8.dtype} {imgs_u8.shape}")
resizer = build_resizer("clean")
out = []
with torch.no_grad():
for i in range(0, len(imgs_u8), int(batch)):
xs = [resizer(im) for im in imgs_u8[i:i + int(batch)]]
x = torch.from_numpy(np.stack(xs).transpose(0, 3, 1, 2).copy()).to(device)
if tuple(x.shape[-2:]) != (IMAGE_SIZE, IMAGE_SIZE):
raise AssertionError("resizer did not produce 299x299")
sp = spatial_tap(extractor, (x - 128) / 128)[:, :SPATIAL_CHANNELS]
out.append(sp.permute(0, 2, 3, 1).reshape(sp.shape[0], -1).cpu().numpy())
feats = np.concatenate(out).astype(np.float32)
if feats.shape[1] != SPATIAL_DIM:
raise AssertionError(f"spatial feature dim {feats.shape[1]} != {SPATIAL_DIM}")
return feats
def verify_spatial_tap(extractor=None) -> dict:
"""Check on CPU that (a) the stepwise traversal reproduces layers.forward exactly (so the tap sits on the
cFID path) and (b) the tapped tensor is the canonical 17x17 mixed_6/conv representation. Raises on failure;
returns the record that MANIFEST_CSFID.json carries."""
import torch
ext = extractor if extractor is not None else build_extractor(torch.device("cpu"))
L = ext.layers
if not (hasattr(L, "mixed_5") and hasattr(L, "mixed_6") and hasattr(L.mixed_6, "conv")):
raise RuntimeError("this Inception build does not expose layers.mixed_6.conv -- canonical sFID unavailable")
x = (torch.rand(3, 3, IMAGE_SIZE, IMAGE_SIZE) * 255.0 - 128) / 128
with torch.no_grad():
h = x
for st in _pre_mixed6_stages(L):
h = st.forward(h)
mixed5 = h
sp = L.mixed_6.conv.forward(mixed5)
for st in _post_mixed6_stages(L):
h = st.forward(h)
diff = float((h.view(3, POOL3_DIM) - L.forward(x).view(3, POOL3_DIM)).abs().max())
flat_dim = int(sp[:, :SPATIAL_CHANNELS].permute(0, 2, 3, 1).reshape(3, -1).shape[1])
rec = {"mixed_5_output_shape_CHW": list(sp.shape[1:]) and list(mixed5.shape[1:]),
"mixed_6_conv_output_shape_CHW": list(sp.shape[1:]), "channels_total": int(sp.shape[1]),
"channels_kept": SPATIAL_CHANNELS, "feature_dim": flat_dim, "flatten": "NHWC",
"stepwise_vs_layers_forward_max_abs_diff_cpu": diff, "bitexact_cpu": diff == 0.0}
probs = []
if tuple(mixed5.shape[1:]) != _MIXED5_SHAPE:
probs.append(f"mixed_5 output {tuple(mixed5.shape[1:])} != canonical {_MIXED5_SHAPE}")
if tuple(sp.shape[2:]) != (SPATIAL_GRID, SPATIAL_GRID):
probs.append(f"mixed_6/conv spatial {tuple(sp.shape[2:])} != {SPATIAL_GRID}x{SPATIAL_GRID}")
if flat_dim != SPATIAL_DIM:
probs.append(f"feature dim {flat_dim} != {SPATIAL_DIM}")
if diff != 0.0:
probs.append(f"stepwise traversal differs from layers.forward on CPU by {diff}")
if probs:
raise RuntimeError("spatial tap verification failed: " + "; ".join(probs))
return rec