Buckets:
| #!/usr/bin/env python3 | |
| # /// script | |
| # requires-python = ">=3.10" | |
| # dependencies = [ | |
| # "datasets>=4.0", | |
| # "numpy>=2.0", | |
| # "pillow>=11", | |
| # "torch>=2.6", | |
| # "torchvision>=0.21", | |
| # ] | |
| # /// | |
| """Scaled UTKFace feature-spectrum experiment for Figure 2(c).""" | |
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import json | |
| import math | |
| import os | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from datasets import load_dataset | |
| from torch.utils.data import DataLoader, Dataset | |
| from torchvision.models import ( | |
| ResNet18_Weights, | |
| ResNet50_Weights, | |
| ViT_B_16_Weights, | |
| resnet18, | |
| resnet50, | |
| vit_b_16, | |
| ) | |
| class Images(Dataset): | |
| def __init__(self, images, transform): | |
| self.images = images | |
| self.transform = transform | |
| def __len__(self): | |
| return len(self.images) | |
| def __getitem__(self, idx): | |
| return self.transform(self.images[idx].convert("RGB")) | |
| def html_plot(path: Path, rows: list[dict], fits: dict[str, float]) -> None: | |
| colors = {"ResNet18": "#2563eb", "ResNet50": "#dc2626", "ViT-B/16": "#059669"} | |
| width, height, left, bottom = 900, 520, 90, 70 | |
| top, right = 55, 25 | |
| grouped = {} | |
| for row in rows: | |
| grouped.setdefault(row["model"], []).append((math.log10(int(row["rank"])), math.log10(float(row["eigenvalue"])))) | |
| xs = [x for vals in grouped.values() for x, _ in vals] | |
| ys = [y for vals in grouped.values() for _, y in vals] | |
| xmin, xmax, ymin, ymax = min(xs), max(xs), min(ys), max(ys) | |
| sx = lambda x: left + (x - xmin) / (xmax - xmin) * (width - left - right) | |
| sy = lambda y: height - bottom - (y - ymin) / (ymax - ymin) * (height - top - bottom) | |
| p = ["<!doctype html><meta charset='utf-8'><div style='max-width:920px;margin:auto;font-family:system-ui'>", | |
| f"<svg viewBox='0 0 {width} {height}' role='img' aria-label='UTKFace feature spectra'>", | |
| f"<text x='{width/2}' y='30' text-anchor='middle' font-size='20' font-weight='700'>Scaled UTKFace feature spectra</text>", | |
| f"<line x1='{left}' y1='{height-bottom}' x2='{width-right}' y2='{height-bottom}' stroke='#111827'/>", | |
| f"<line x1='{left}' y1='{top}' x2='{left}' y2='{height-bottom}' stroke='#111827'/>"] | |
| for idx, (name, vals) in enumerate(grouped.items()): | |
| pts = " ".join(f"{sx(x):.2f},{sy(y):.2f}" for x, y in vals) | |
| color = colors[name] | |
| p.append(f"<polyline points='{pts}' fill='none' stroke='{color}' stroke-width='2'/>") | |
| p.append(f"<text x='{width-250}' y='{65+idx*23}' fill='{color}' font-size='13'>{name}: α={fits[name]:.3f}</text>") | |
| p += [f"<text x='{width/2}' y='{height-20}' text-anchor='middle'>log10 eigenvalue rank</text>", | |
| f"<text x='20' y='{height/2}' transform='rotate(-90 20 {height/2})' text-anchor='middle'>log10 eigenvalue</text>", | |
| "</svg><p>Dataset: <a href='https://huggingface.co/datasets/deedax/UTK-Face-Revised'>deedax/UTK-Face-Revised</a>. Pretrained weights: torchvision.</p></div>"] | |
| path.write_text("\n".join(p), encoding="utf-8") | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--output", default="/results/architecture") | |
| ap.add_argument("--samples", type=int, default=1000) | |
| ap.add_argument("--batch-size", type=int, default=32) | |
| args = ap.parse_args() | |
| out = Path(args.output) | |
| out.mkdir(parents=True, exist_ok=True) | |
| start = time.time() | |
| stream = load_dataset("deedax/UTK-Face-Revised", split="train", streaming=True) | |
| images = [] | |
| for example in stream: | |
| images.append(example["image"]) | |
| if len(images) >= args.samples: | |
| break | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| specs = [ | |
| ("ResNet18", resnet18, ResNet18_Weights.DEFAULT, "fc"), | |
| ("ResNet50", resnet50, ResNet50_Weights.DEFAULT, "fc"), | |
| ("ViT-B/16", vit_b_16, ViT_B_16_Weights.DEFAULT, "heads"), | |
| ] | |
| rows, fits, dims = [], {}, {} | |
| for name, constructor, weights, head_attr in specs: | |
| model = constructor(weights=weights) | |
| setattr(model, head_attr, torch.nn.Identity()) | |
| model.eval().to(device) | |
| loader = DataLoader(Images(images, weights.transforms()), batch_size=args.batch_size, num_workers=2) | |
| feats = [] | |
| with torch.inference_mode(): | |
| for batch in loader: | |
| batch = batch.to(device) | |
| with torch.autocast(device_type=device.type, enabled=device.type == "cuda"): | |
| feat = model(batch) | |
| feats.append(feat.float().cpu()) | |
| f = torch.cat(feats) | |
| dims[name] = int(f.shape[1]) | |
| # Paper uses uncentered empirical covariance E[ff^T]. SVD avoids a dxd matrix. | |
| s = torch.linalg.svdvals(f) | |
| eig = (s**2 / f.shape[0]).numpy() | |
| stop = min(500, len(eig) - 10) | |
| ranks = np.arange(11, stop + 1) | |
| vals = eig[10:stop] | |
| slope = np.polyfit(np.log(ranks), np.log(vals), 1)[0] | |
| fits[name] = float(-slope) | |
| for rank, val in zip(np.arange(1, min(500, len(eig)) + 1), eig[:500]): | |
| rows.append({"model": name, "rank": int(rank), "eigenvalue": f"{float(val):.12g}"}) | |
| del model, f, feats | |
| if device.type == "cuda": | |
| torch.cuda.empty_cache() | |
| # Bucket mounts can refresh while a long feature-extraction pass is running. | |
| # Recreate the leaf directory immediately before committing the results. | |
| out.mkdir(parents=True, exist_ok=True) | |
| with (out / "architecture_spectra.csv").open("w", newline="", encoding="utf-8") as f: | |
| writer = csv.DictWriter(f, fieldnames=["model", "rank", "eigenvalue"]) | |
| writer.writeheader(); writer.writerows(rows) | |
| html_plot(out / "architecture_spectra.html", rows, fits) | |
| metrics = { | |
| "dataset": "https://huggingface.co/datasets/deedax/UTK-Face-Revised", | |
| "samples": len(images), | |
| "models": list(fits), | |
| "feature_dimensions": dims, | |
| "fitted_alpha": fits, | |
| "ordering_smaller_alpha_slower_decay": [k for k, _ in sorted(fits.items(), key=lambda kv: kv[1])], | |
| "device": str(device), | |
| "gpu": torch.cuda.get_device_name(0) if device.type == "cuda" else None, | |
| "runtime_seconds": time.time() - start, | |
| "scope": "scaled proxy for paper Figure 2(c): 1000 rather than 2000 UTKFace images; three rather than four architectures", | |
| } | |
| (out / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") | |
| print(json.dumps(metrics, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 6.55 kB
- Xet hash:
- bc5618290d18eccb9ada9e4937bd5c63a0ea85eafee0e3ec9a5b1599cfbc3c9e
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.