Download code/bifurcation/fid/backend.py from coarse-grained-forecasting-benchmark/forecasting-models: direct link, hf CLI and curl.
- Browser
- Download file 11.6 kB
-
https://huggingface.co/coarse-grained-forecasting-benchmark/forecasting-models/resolve/main/code/bifurcation/fid/backend.py
- Command line
-
hf download hf://coarse-grained-forecasting-benchmark/forecasting-models/code/bifurcation/fid/backend.py
-
curl -L -o backend.py https://huggingface.co/coarse-grained-forecasting-benchmark/forecasting-models/resolve/main/code/bifurcation/fid/backend.py
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 | |