ed33fbbaf7
research/determinism-cost.md and its record (record.json): mbt2018_mean at all eight published qualities, integer export against the model's own float codec, on the pinned reference set's 24 self-test pictures. BD-rate by PSNR and MS-SSIM per picture; each codec's decodes repeated across nine integer and seven float platforms and compared byte for byte, the float ones placed in or after the entropy decode; decode time and energy with counter, scope, idle brackets, liveness and ratio over seven rounds. determinism-tiers §7.1 quotes the record, each figure tied to its place in it by a marker the verifier checks; §7.2 says what WAI claims and does not, and makes normative that a record figure is never stated beyond what the record names, that no receipt takes its energy from it, and that a cost figure in the spec names its measurement. SPEC §5's int_hyper row points to §7. CHANGELOG entry. Co-Authored-By: Claude Opus 5.5 <[email protected]>
878 lines
41 KiB
Python
878 lines
41 KiB
Python
#!/usr/bin/env python3
|
|
"""What exactness costs `wai.neural.int_hyper`: the measurement behind
|
|
research/determinism-cost.md and its record, research/determinism-cost/record.json.
|
|
|
|
The float codec is the pretrained mean-scale hyperprior `mbt2018_mean` (MSE, qualities
|
|
1-8) as its own library codes and decodes it. The integer codec is the same model exported
|
|
the way tools/wai_sota_int_synth_export.py exports quality 6 (per-channel int12 weights,
|
|
activations at 2^12, integer IGDN, the integer hyper-decoder deciding every scale bucket),
|
|
coded as a `WIH1` payload (extensions/integer-payloads.md §4) and decoded by the reference
|
|
decoder (`wai_det_cost`, which runs `int_payload::decode_int`). Every step writes files
|
|
under one work directory, so each can be rerun alone and on another machine.
|
|
|
|
wai_determinism_cost.py select <work> <reference-root> pick and check the pictures
|
|
wai_determinism_cost.py export <work> [q ...] pinned checkpoints -> integer sets
|
|
wai_determinism_cost.py encode <work> [q ...] both codecs; the float decode on this CPU
|
|
wai_determinism_cost.py jobs <work> <out-tag> the integer decode's job list
|
|
wai_determinism_cost.py fdecode <work> <tag> [--device cpu|mps|cuda] [--threads N]
|
|
decode every float stream again
|
|
wai_determinism_cost.py fserve <work> <q> [--device D] [--threads N] [--ids a,b]
|
|
the timed float child of wai_det_energy.py
|
|
wai_determinism_cost.py metrics <work> distortion of every decode
|
|
wai_determinism_cost.py bd <work> BD-rate, integer against float
|
|
wai_determinism_cost.py record <work> <record.json> the record: raw readings and their summary
|
|
|
|
Needs the pinned stack of tools/requirements-sota-int.txt for `export`, `encode` and
|
|
`fdecode`; `select`, `jobs`, `metrics` and `bd` need numpy and Pillow only.
|
|
"""
|
|
import hashlib
|
|
import importlib.util
|
|
import io
|
|
import json
|
|
import math
|
|
import os
|
|
import struct
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
HERE = os.path.dirname(os.path.abspath(__file__))
|
|
QUALITIES = list(range(1, 9))
|
|
|
|
# The published mbt2018_mean (MSE) checkpoints, by quality: the file the library's zoo
|
|
# names, and the SHA-256 of its bytes, measured when this record was made. A checkpoint
|
|
# with other bytes is refused before it is loaded. Quality 6 is the one
|
|
# wai_sota_int_synth_export.py pins, and its digest is the same.
|
|
CKPT_SHA256 = {
|
|
1: "e522738dc4599d1538fcdb16e21aa56df9146d60e8e083542466b387c1960156",
|
|
2: "e54a039d587afddd086b836cdd8eed4e0343d7cdd91028489ad0f3e7a72d307d",
|
|
3: "723404a8a039371abe8a14f05160ee1614cf5e0a3fb2a4e5f7ba569b68c22601",
|
|
4: "6dba02a30c78ff78baafd2072dd78646c27af6c4dc8416097c2c2858c1f186ce",
|
|
5: "d504e8ebe7e6396efd0b5fd093a32136714fd7e4be03db389e6f07aee24fdb39",
|
|
6: "a19628aba19732125e02caf5c38e233edbe2a641145a5b1779c51f4527f51dd6",
|
|
7: "d5d441d16ba3e5a80ead947e35d2dddc12fb0d9785a8006e9b35a0cef7668d91",
|
|
8: "8089ae3e23d9df5b7c68107c82c31da21019147aedddb03f03350de58d5385b7",
|
|
}
|
|
|
|
# The pictures: the reference manifest's PNG entries that it marks `self_test`, each one
|
|
# picture of 768x512 or 512x768 (wai/tools/wai_visual_provenance_check.py's manifest).
|
|
MANIFEST = "wai-image-references.json"
|
|
MANIFEST_SHA256 = "db33c6633f89a66d6742a37cd348208792f6d49810da739ddf26d10dc6554b66"
|
|
|
|
|
|
def sha256(b):
|
|
return hashlib.sha256(b).hexdigest()
|
|
|
|
|
|
def jdump(obj, path):
|
|
with open(path, "w") as f:
|
|
json.dump(obj, f, indent=1, sort_keys=True)
|
|
f.write("\n")
|
|
|
|
|
|
def jload(path):
|
|
with open(path) as f:
|
|
return json.load(f)
|
|
|
|
|
|
def export_module():
|
|
spec = importlib.util.spec_from_file_location("sota_export", os.path.join(HERE, "wai_sota_int_synth_export.py"))
|
|
m = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(m)
|
|
return m
|
|
|
|
|
|
# ---------------------------------------------------------------- select
|
|
|
|
def cmd_select(work, root):
|
|
raw = open(os.path.join(root, MANIFEST), "rb").read()
|
|
if sha256(raw) != MANIFEST_SHA256:
|
|
sys.exit(f"the manifest's SHA-256 is {sha256(raw)}, not the pinned {MANIFEST_SHA256}")
|
|
man = json.loads(raw)
|
|
from PIL import Image
|
|
pics = []
|
|
for e in man["entries"]:
|
|
if not (e.get("self_test") and e["path"].endswith(".png") and len(e["images"]) == 1):
|
|
continue
|
|
im = e["images"][0]
|
|
if (im["width"], im["height"]) not in ((768, 512), (512, 768)):
|
|
continue
|
|
b = open(os.path.join(root, e["path"]), "rb").read()
|
|
if sha256(b) != e["sha256"]:
|
|
sys.exit(f"an entry's bytes do not match the manifest ({e['sha256'][:16]})")
|
|
px = np.array(Image.open(io.BytesIO(b)).convert("RGB"))
|
|
h, w = px.shape[:2]
|
|
assert h % 64 == 0 and w % 64 == 0
|
|
pid = e["sha256"][:16]
|
|
os.makedirs(os.path.join(work, "src"), exist_ok=True)
|
|
px.tofile(os.path.join(work, "src", f"{pid}.rgb"))
|
|
pics.append({"id": pid, "file_sha256": e["sha256"], "width": int(w), "height": int(h),
|
|
"rgb_sha256": sha256(px.tobytes())})
|
|
jdump({"manifest_sha256": MANIFEST_SHA256,
|
|
"rule": "PNG entries marked self_test with one picture of 768x512 or 512x768",
|
|
"pictures": pics}, os.path.join(work, "pictures.json"))
|
|
print(f"{len(pics)} pictures")
|
|
|
|
|
|
def pictures(work):
|
|
return jload(os.path.join(work, "pictures.json"))["pictures"]
|
|
|
|
|
|
def source_rgb(work, p):
|
|
return np.fromfile(os.path.join(work, "src", f"{p['id']}.rgb"), dtype=np.uint8).reshape(p["height"], p["width"], 3)
|
|
|
|
|
|
# ---------------------------------------------------------------- models
|
|
|
|
def load_model(q, allow_new=False):
|
|
"""mbt2018_mean (MSE) at quality q from the published checkpoint, its bytes checked
|
|
against CKPT_SHA256 before they are loaded, as tensors only."""
|
|
import torch
|
|
from compressai.zoo import image as zoo_image
|
|
from compressai.zoo.pretrained import load_pretrained
|
|
url = zoo_image.model_urls["mbt2018-mean"]["mse"][q]
|
|
path = os.path.join(torch.hub.get_dir(), "checkpoints", os.path.basename(url))
|
|
if not os.path.exists(path):
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
torch.hub.download_url_to_file(url, path, progress=False)
|
|
blob = open(path, "rb").read()
|
|
got = sha256(blob)
|
|
# The zoo's file names carry the first eight hex digits of the digest.
|
|
if not os.path.basename(url).split("-")[-1].startswith(got[:8]):
|
|
sys.exit(f"q{q}: {path} has SHA-256 {got}, which does not match the prefix in its name")
|
|
want = CKPT_SHA256[q]
|
|
if want is None and not allow_new:
|
|
sys.exit(f"q{q}: no pinned digest; measured {got}")
|
|
if want is not None and got != want:
|
|
sys.exit(f"q{q}: {path} has SHA-256 {got}, pinned {want}")
|
|
state = torch.load(io.BytesIO(blob), map_location="cpu", weights_only=True)
|
|
m = zoo_image.model_architectures["mbt2018-mean"].from_state_dict(load_pretrained(state)).eval()
|
|
return m, {"quality": q, "url_name": os.path.basename(url), "sha256": got}
|
|
|
|
|
|
# ---------------------------------------------------------------- integer export
|
|
|
|
def int_net(E, seq, kind):
|
|
"""One integer network from a float Sequential, quantised as E (the quality-6 export
|
|
tool) quantises g_s and h_s."""
|
|
net = {"afrac": E.AFRAC, "gf": E.GF, "shift": E.SHIFT, "layers": []}
|
|
import torch
|
|
for layer in seq.children():
|
|
t = type(layer).__name__
|
|
if t == "ConvTranspose2d":
|
|
wi, sc = E.q_weight_pc(layer.weight.detach().numpy())
|
|
net["layers"].append({"type": "convT", "in_ch": int(wi.shape[0]), "out_ch": int(wi.shape[1]),
|
|
"k": int(wi.shape[2]), "stride": int(layer.stride[0]), "pad": int(layer.padding[0]),
|
|
"opad": int(layer.output_padding[0]), "weight": wi.reshape(-1).tolist(),
|
|
"mult": np.round(sc * (2 ** E.SHIFT)).astype(np.int64).tolist(),
|
|
"bias": np.round(layer.bias.detach().numpy() * (2 ** E.AFRAC)).astype(np.int64).tolist()})
|
|
elif t == "Conv2d":
|
|
wi, sc = E.q_weight_conv(layer.weight.detach().numpy())
|
|
net["layers"].append({"type": "conv", "in_ch": int(wi.shape[1]), "out_ch": int(wi.shape[0]),
|
|
"k": int(wi.shape[2]), "stride": int(layer.stride[0]), "pad": int(layer.padding[0]),
|
|
"weight": wi.reshape(-1).tolist(),
|
|
"mult": np.round(sc * (2 ** E.SHIFT)).astype(np.int64).tolist(),
|
|
"bias": np.round(layer.bias.detach().numpy() * (2 ** E.AFRAC)).astype(np.int64).tolist()})
|
|
elif t == "GDN":
|
|
with torch.no_grad():
|
|
beta = layer.beta_reparam(layer.beta).numpy()
|
|
C = beta.shape[0]
|
|
gamma = layer.gamma_reparam(layer.gamma).reshape(C, C).numpy()
|
|
net["layers"].append({"type": "igdn", "c": int(C),
|
|
"gamma": np.round(gamma * (2 ** E.GF)).astype(np.int64).reshape(-1).tolist(),
|
|
"beta": np.round(beta * (2 ** (2 * E.AFRAC))).astype(np.int64).tolist()})
|
|
elif t == "ReLU":
|
|
net["layers"].append({"type": "relu"})
|
|
elif t == "LeakyReLU":
|
|
net["layers"].append({"type": "leaky", "slope_q": 655, "slope_frac": 16})
|
|
else:
|
|
sys.exit(f"{kind}: no integer form for a {t} layer")
|
|
return net
|
|
|
|
|
|
def requant(raw, mult, bias, shift):
|
|
"""((raw * mult) >> shift) + bias per output channel, exact: in int64 where the
|
|
product provably fits, else in Python integers."""
|
|
if np.abs(raw).max(initial=0) * np.abs(mult).max(initial=0) < 2 ** 62:
|
|
return ((raw * mult[:, None, None]) >> shift) + bias[:, None, None]
|
|
return ((raw.astype(object) * mult[:, None, None]) >> shift).astype(np.int64) + bias[:, None, None]
|
|
|
|
|
|
def conv_t_fast(x, L, shift):
|
|
"""The integer transposed convolution of `int_hyper_synth` (weight [in,out,k,k]),
|
|
tap by tap: out[:, iy*s+ky, ix*s+kx] += w[:, :, ky, kx]^T x[:, iy, ix]. Integer sums
|
|
are exact in any order, so this equals the export tool's per-pixel scatter."""
|
|
cin, co, k, s, p, op = L["in_ch"], L["out_ch"], L["k"], L["stride"], L["pad"], L["opad"]
|
|
w = np.array(L["weight"], dtype=np.int64).reshape(cin, co, k, k)
|
|
_, H, W = x.shape
|
|
Ho, Wo = (H - 1) * s + k + op - 2 * p, (W - 1) * s + k + op - 2 * p
|
|
raw = np.zeros((co, Ho + 2 * p, Wo + 2 * p), dtype=np.int64)
|
|
xf = x.reshape(cin, -1)
|
|
assert np.abs(x).max(initial=0) * np.abs(w).max(initial=0) * cin * k * k < 2 ** 63
|
|
for ky in range(k):
|
|
for kx in range(k):
|
|
c = (w[:, :, ky, kx].T @ xf).reshape(co, H, W)
|
|
raw[:, ky:ky + s * H:s, kx:kx + s * W:s] += c
|
|
raw = raw[:, p:p + Ho, p:p + Wo]
|
|
return requant(raw, np.array(L["mult"], dtype=np.int64), np.array(L["bias"], dtype=np.int64), shift)
|
|
|
|
|
|
def conv_fast(x, L, shift):
|
|
cin, co, k, s, p = L["in_ch"], L["out_ch"], L["k"], L["stride"], L["pad"]
|
|
w = np.array(L["weight"], dtype=np.int64).reshape(co, cin, k, k)
|
|
_, H, W = x.shape
|
|
Ho, Wo = (H + 2 * p - k) // s + 1, (W + 2 * p - k) // s + 1
|
|
xp = np.zeros((cin, H + 2 * p, W + 2 * p), dtype=np.int64)
|
|
xp[:, p:p + H, p:p + W] = x
|
|
assert np.abs(x).max(initial=0) * np.abs(w).max(initial=0) * cin * k * k < 2 ** 63
|
|
raw = np.zeros((co, Ho, Wo), dtype=np.int64)
|
|
for ky in range(k):
|
|
for kx in range(k):
|
|
sub = xp[:, ky:ky + s * Ho:s, kx:kx + s * Wo:s].reshape(cin, -1)
|
|
raw += (w[:, :, ky, kx] @ sub).reshape(co, Ho, Wo)
|
|
return requant(raw, np.array(L["mult"], dtype=np.int64), np.array(L["bias"], dtype=np.int64), shift)
|
|
|
|
|
|
def run_hs(net, x):
|
|
"""The integer hyper-decoder (convolutions and ReLUs only), as `int_hyper_synth`
|
|
runs it."""
|
|
for L in net["layers"]:
|
|
t = L["type"]
|
|
if t == "convT":
|
|
x = conv_t_fast(x, L, net["shift"])
|
|
elif t == "conv":
|
|
x = conv_fast(x, L, net["shift"])
|
|
elif t == "relu":
|
|
x = np.maximum(x, 0)
|
|
elif t == "leaky":
|
|
x = np.where(x >= 0, x, (x * L["slope_q"]) >> L["slope_frac"])
|
|
else:
|
|
sys.exit(f"h_s has a {t} layer")
|
|
return x
|
|
|
|
|
|
def cmd_export(work, qs, allow_new=False):
|
|
"""Per quality, the integer set (integer-payloads §9: entropy.json, g_s.whm, h_s.whm)
|
|
and the model's identity."""
|
|
import torch
|
|
from compressai.entropy_models import GaussianConditional
|
|
torch.use_deterministic_algorithms(True)
|
|
torch.set_num_threads(1)
|
|
E = export_module()
|
|
out = {}
|
|
for q in qs:
|
|
m, ident = load_model(q, allow_new)
|
|
g_s = int_net(E, m.g_s, "g_s")
|
|
h_s = int_net(E, m.h_s, "h_s")
|
|
eb = m.entropy_bottleneck
|
|
eb.update()
|
|
median = eb._get_medians().reshape(-1).detach().numpy().astype(np.float64)
|
|
scale_table = list(torch.exp(torch.linspace(np.log(0.11), np.log(256.0), 64)).numpy().astype(float))
|
|
gc = GaussianConditional(scale_table)
|
|
gc.update()
|
|
ent = {"afrac": E.AFRAC, "frac": E.AFRAC,
|
|
"scale_edges": np.floor(np.array(scale_table) * (2 ** E.AFRAC)).astype(np.int64).tolist(),
|
|
"cdfs": gc._quantized_cdf.tolist(), "cdf_lengths": gc._cdf_length.reshape(-1).int().tolist(),
|
|
"offsets": gc._offset.reshape(-1).int().tolist(),
|
|
"z_cdfs": eb._quantized_cdf.tolist(), "z_cdf_lengths": eb._cdf_length.reshape(-1).int().tolist(),
|
|
"z_offsets": eb._offset.reshape(-1).int().tolist(),
|
|
"z_median_q": np.round(median * (2 ** E.AFRAC)).astype(np.int64).tolist()}
|
|
d = os.path.join(work, "sets", f"q{q}")
|
|
os.makedirs(d, exist_ok=True)
|
|
open(os.path.join(d, "g_s.whm"), "wb").write(E.model_to_bin(g_s))
|
|
open(os.path.join(d, "h_s.whm"), "wb").write(E.model_to_bin(h_s))
|
|
with open(os.path.join(d, "entropy.json"), "w") as f:
|
|
json.dump(ent, f)
|
|
jdump(h_s, os.path.join(d, "h_s.json")) # the encoder's copy of h_s
|
|
ident["set"] = {n: sha256(open(os.path.join(d, n), "rb").read()) for n in ("entropy.json", "g_s.whm", "h_s.whm")}
|
|
ident["N"], ident["M"] = int(m.N), int(m.M)
|
|
out[q] = ident
|
|
print(f"q{q}: N={m.N} M={m.M} checkpoint {ident['sha256'][:16]}… set g_s {ident['set']['g_s.whm'][:16]}…")
|
|
path = os.path.join(work, "models.json")
|
|
prev = jload(path) if os.path.exists(path) else {}
|
|
prev.update({str(k): v for k, v in out.items()})
|
|
jdump(prev, path)
|
|
|
|
|
|
# ---------------------------------------------------------------- encode
|
|
|
|
def to_rgb8(x_hat):
|
|
"""A float reconstruction in [0, 1] as 8-bit samples: x·255 rounded half to even, as
|
|
the float library's own tensor arithmetic gives it on the device it ran on."""
|
|
import torch
|
|
return (x_hat.clamp(0, 1) * 255).round().to(torch.uint8)[0].permute(1, 2, 0).contiguous().cpu().numpy()
|
|
|
|
|
|
def float_container(strings, zshape):
|
|
y, z = strings[0][0], strings[1][0]
|
|
return b"DETF" + struct.pack("<IIII", zshape[0], zshape[1], len(y), len(z)) + y + z
|
|
|
|
|
|
def read_float_container(b):
|
|
assert b[:4] == b"DETF"
|
|
zh, zw, ly, lz = struct.unpack("<IIII", b[4:20])
|
|
y = b[20:20 + ly]
|
|
z = b[20 + ly:20 + ly + lz]
|
|
assert 20 + ly + lz == len(b)
|
|
return [[y], [z]], (zh, zw)
|
|
|
|
|
|
def float_decode(m, strings, zshape):
|
|
"""The float library's own decode of a mean-scale stream (its `decompress`), with the
|
|
decoded indexes and symbols kept, so a mismatch can be placed: in the entropy decode
|
|
(other indexes or symbols) or after it."""
|
|
import torch
|
|
with torch.no_grad():
|
|
z_hat = m.entropy_bottleneck.decompress(strings[1], zshape)
|
|
scales_hat, means_hat = m.h_s(z_hat).chunk(2, 1)
|
|
indexes = m.gaussian_conditional.build_indexes(scales_hat)
|
|
symbols = m.gaussian_conditional.decompress(strings[0], indexes)
|
|
y_hat = symbols + means_hat
|
|
x_hat = m.g_s(y_hat).clamp_(0, 1)
|
|
idx = indexes.cpu().to(torch.int32).numpy().tobytes()
|
|
sym = symbols.cpu().round().to(torch.int32).numpy().tobytes()
|
|
return to_rgb8(x_hat), sha256(idx), sha256(sym)
|
|
|
|
|
|
def tables_sha256(m):
|
|
"""The entropy tables a float decode reads (both models' quantised CDFs, lengths and
|
|
offsets): a decode on another machine is comparable only under the same tables."""
|
|
import torch
|
|
h = hashlib.sha256()
|
|
for em in (m.entropy_bottleneck, m.gaussian_conditional):
|
|
for t in (em._quantized_cdf, em._cdf_length, em._offset):
|
|
h.update(t.cpu().to(torch.int32).numpy().tobytes())
|
|
return h.hexdigest()
|
|
|
|
|
|
def wih1(m_ch, sh, sw, n_z, zh, zw, z_rans, rans):
|
|
o = b"WIH1" + bytes([1, 0, 0, 0]) + struct.pack("<6I", m_ch, sh, sw, n_z, zh, zw)
|
|
return o + struct.pack("<I", len(z_rans)) + z_rans + struct.pack("<I", len(rans)) + rans
|
|
|
|
|
|
def cmd_encode(work, qs):
|
|
import torch
|
|
from compressai.ans import RansEncoder
|
|
torch.use_deterministic_algorithms(True)
|
|
torch.set_num_threads(int(os.environ.get("DET_THREADS", "1")))
|
|
E = export_module()
|
|
pics = pictures(work)
|
|
rows = {}
|
|
for q in qs:
|
|
m, _ = load_model(q)
|
|
m.update()
|
|
tables = tables_sha256(m)
|
|
d = os.path.join(work, "sets", f"q{q}")
|
|
hs = jload(os.path.join(d, "h_s.json"))
|
|
ent = jload(os.path.join(d, "entropy.json"))
|
|
edges = np.array(ent["scale_edges"], dtype=np.int64)
|
|
mq = np.array(ent["z_median_q"], dtype=np.int64)
|
|
for sub in ("float", "int", "dec"):
|
|
os.makedirs(os.path.join(work, sub, f"q{q}"), exist_ok=True)
|
|
for p in pics:
|
|
px = source_rgb(work, p)
|
|
x = torch.from_numpy(px.astype(np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0)
|
|
with torch.no_grad():
|
|
# the float codec, as its library codes and decodes it
|
|
c = m.compress(x)
|
|
fb = float_container(c["strings"], tuple(c["shape"]))
|
|
open(os.path.join(work, "float", f"q{q}", f"{p['id']}.detf"), "wb").write(fb)
|
|
rgb, ih, sh_ = float_decode(m, c["strings"], tuple(c["shape"]))
|
|
rgb.tofile(os.path.join(work, "dec", f"q{q}", f"{p['id']}.float.rgb"))
|
|
# the integer codec: the float analysis, then everything the decoder reads in integers
|
|
y = m.g_a(x)
|
|
z = m.h_a(y)
|
|
y_np = y[0].numpy()
|
|
eb_median = m.entropy_bottleneck._get_medians().reshape(-1).detach().numpy().astype(np.float64)
|
|
z_sym = np.round(z[0].numpy() - eb_median[:, None, None]).astype(np.int64)
|
|
Cz, zh, zw = z_sym.shape
|
|
z_bits = RansEncoder().encode_with_indexes(
|
|
z_sym.reshape(-1).tolist(), np.repeat(np.arange(Cz), zh * zw).tolist(),
|
|
ent["z_cdfs"], ent["z_cdf_lengths"], ent["z_offsets"])
|
|
zi = z_sym * (2 ** E.AFRAC) + mq[:, None, None]
|
|
sm = run_hs(hs, zi)
|
|
M = sm.shape[0] // 2
|
|
sigma, mu = sm[:M], sm[M:]
|
|
resid = np.round(y_np - mu.astype(np.float64) / (2 ** E.AFRAC)).astype(np.int64)
|
|
idx = E.det_build_index_vec(sigma, edges)
|
|
y_bits = RansEncoder().encode_with_indexes(resid.reshape(-1).tolist(), idx.reshape(-1).tolist(),
|
|
ent["cdfs"], ent["cdf_lengths"], ent["offsets"])
|
|
pay = wih1(M, sigma.shape[1], sigma.shape[2], Cz, zh, zw, z_bits, y_bits)
|
|
open(os.path.join(work, "int", f"q{q}", f"{p['id']}.wih1"), "wb").write(pay)
|
|
rows[f"{p['id']}/q{q}"] = {
|
|
"float_stream_bytes": len(c["strings"][0][0]) + len(c["strings"][1][0]),
|
|
"float_stream_sha256": sha256(fb),
|
|
"float_tables_sha256": tables,
|
|
"float_cpu": {"rgb_sha256": sha256(rgb.tobytes()), "indexes_sha256": ih, "symbols_sha256": sh_},
|
|
"int_stream_bytes": len(z_bits) + len(y_bits),
|
|
"int_payload_bytes": len(pay),
|
|
"int_payload_sha256": sha256(pay),
|
|
}
|
|
row = rows[f"{p['id']}/q{q}"]
|
|
print(f"q{q} {p['id']}: float {row['float_stream_bytes']} B, integer {row['int_stream_bytes']} B", flush=True)
|
|
path = os.path.join(work, "encoded.json")
|
|
prev = jload(path) if os.path.exists(path) else {}
|
|
prev.update(rows)
|
|
jdump(prev, path)
|
|
|
|
|
|
def cmd_jobs(work, tag):
|
|
"""One line per integer payload: its set, the payload, where the decode goes."""
|
|
lines = []
|
|
for q in QUALITIES:
|
|
d = os.path.join(work, "int", f"q{q}")
|
|
if not os.path.isdir(d):
|
|
continue
|
|
os.makedirs(os.path.join(work, "dec", f"q{q}"), exist_ok=True)
|
|
for f in sorted(os.listdir(d)):
|
|
pid = f.split(".")[0]
|
|
out = os.path.join(work, "dec", f"q{q}", f"{pid}.int.{tag}.rgb") if tag != "-" else "-"
|
|
lines.append(f"{os.path.join(work, 'sets', f'q{q}')}\t{os.path.join(d, f)}\t{out}")
|
|
sys.stdout.write("\n".join(lines) + "\n")
|
|
|
|
|
|
def cmd_fdecode(work, tag, device, threads):
|
|
"""Decode every float stream in a fresh process on `device`: the RGB8, the indexes and
|
|
the symbols, by SHA-256."""
|
|
import platform
|
|
import torch
|
|
torch.use_deterministic_algorithms(True)
|
|
if threads:
|
|
torch.set_num_threads(threads)
|
|
out = {"tag": tag, "device": device, "threads": torch.get_num_threads(), "torch": torch.__version__,
|
|
"machine": platform.machine(), "system": platform.system(), "results": {}}
|
|
srcs = {p["id"]: p for p in pictures(work)}
|
|
for q in QUALITIES:
|
|
d = os.path.join(work, "float", f"q{q}")
|
|
if not os.path.isdir(d):
|
|
continue
|
|
m, _ = load_model(q)
|
|
m.update()
|
|
out.setdefault("tables_sha256", {})[f"q{q}"] = tables_sha256(m)
|
|
m = m.to(device)
|
|
for f in sorted(os.listdir(d)):
|
|
strings, zshape = read_float_container(open(os.path.join(d, f), "rb").read())
|
|
rgb, ih, sh_ = float_decode(m, strings, zshape)
|
|
pid = f.split(".")[0]
|
|
src = source_rgb(work, srcs[pid])
|
|
diff = rgb.astype(np.int64) - src.astype(np.int64)
|
|
out["results"][f"{pid}/q{q}"] = {"rgb_sha256": sha256(rgb.tobytes()), "indexes_sha256": ih,
|
|
"symbols_sha256": sh_, "sse": int((diff * diff).sum())}
|
|
print(f"q{q}: {len(os.listdir(d))} decoded on {device}", flush=True)
|
|
os.makedirs(os.path.join(work, "fdecode"), exist_ok=True)
|
|
jdump(out, os.path.join(work, "fdecode", f"{tag}.json"))
|
|
|
|
|
|
def cmd_fserve(work, q, device, threads, ids):
|
|
"""The timed child of tools/wai_det_energy.py for the float decode: the library's own
|
|
`decompress` of every float stream at quality q (or of the pictures `ids`) on
|
|
`device`, then the 8-bit samples on the host. Each decode's RGB8 is hashed and
|
|
compared with the first decode of its stream; a differing decode is counted, not
|
|
hidden (`done <decodes> <ns> mismatches=<n>`)."""
|
|
import time
|
|
import torch
|
|
torch.use_deterministic_algorithms(True, warn_only=device != "cpu")
|
|
if threads:
|
|
torch.set_num_threads(threads)
|
|
m, _ = load_model(q)
|
|
m.update()
|
|
m = m.to(device)
|
|
d = os.path.join(work, "float", f"q{q}")
|
|
files = [f for f in sorted(os.listdir(d)) if not ids or f.split(".")[0] in ids]
|
|
streams = [read_float_container(open(os.path.join(d, f), "rb").read()) for f in files]
|
|
|
|
def one(st):
|
|
with torch.no_grad():
|
|
x_hat = m.decompress(st[0], st[1])["x_hat"]
|
|
return sha256(to_rgb8(x_hat).tobytes())
|
|
|
|
want = [one(st) for st in streams]
|
|
print(f"ready {len(streams)}", flush=True)
|
|
for line in sys.stdin:
|
|
w = line.split()
|
|
if not w or w[0] == "exit":
|
|
return
|
|
k = int(w[1])
|
|
t0 = time.monotonic_ns()
|
|
n = bad = 0
|
|
for _ in range(k):
|
|
for st, h in zip(streams, want):
|
|
bad += one(st) != h
|
|
n += 1
|
|
print(f"done {n} {time.monotonic_ns() - t0} mismatches={bad}", flush=True)
|
|
|
|
|
|
# ---------------------------------------------------------------- metrics
|
|
|
|
GAUSS = None
|
|
MSSSIM_W = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)
|
|
|
|
|
|
def _gauss():
|
|
global GAUSS
|
|
if GAUSS is None:
|
|
g = np.exp(-((np.arange(11) - 5.0) ** 2) / (2 * 1.5 ** 2))
|
|
GAUSS = g / g.sum()
|
|
return GAUSS
|
|
|
|
|
|
def _filt(x):
|
|
"""Separable 11-tap Gaussian (sigma 1.5), valid region only."""
|
|
g = _gauss()
|
|
h, w = x.shape
|
|
t = sum(g[i] * x[i:h - 10 + i, :] for i in range(11))
|
|
return sum(g[i] * t[:, i:w - 10 + i] for i in range(11))
|
|
|
|
|
|
def ms_ssim_channel(a, b):
|
|
"""MS-SSIM of one channel, 8-bit samples (dynamic range 255): five scales, an 11-tap
|
|
Gaussian window of sigma 1.5, K1 0.01, K2 0.03, 2x2 mean pooling between scales
|
|
(an odd row or column dropped), contrast-structure at the first four scales and
|
|
luminance-contrast-structure at the fifth, weights 0.0448, 0.2856, 0.3001, 0.2363,
|
|
0.1333, each term clamped at zero before its power."""
|
|
C1, C2 = (0.01 * 255) ** 2, (0.03 * 255) ** 2
|
|
a, b = a.astype(np.float64), b.astype(np.float64)
|
|
val = 1.0
|
|
for s in range(5):
|
|
mu_a, mu_b = _filt(a), _filt(b)
|
|
saa = _filt(a * a) - mu_a ** 2
|
|
sbb = _filt(b * b) - mu_b ** 2
|
|
sab = _filt(a * b) - mu_a * mu_b
|
|
cs = (2 * sab + C2) / (saa + sbb + C2)
|
|
if s < 4:
|
|
val *= max(cs.mean(), 0.0) ** MSSSIM_W[s]
|
|
h, w = a.shape[0] // 2 * 2, a.shape[1] // 2 * 2
|
|
a = a[:h, :w].reshape(h // 2, 2, w // 2, 2).mean(axis=(1, 3))
|
|
b = b[:h, :w].reshape(h // 2, 2, w // 2, 2).mean(axis=(1, 3))
|
|
else:
|
|
l = (2 * mu_a * mu_b + C1) / (mu_a ** 2 + mu_b ** 2 + C1)
|
|
val *= max((l * cs).mean(), 0.0) ** MSSSIM_W[s]
|
|
return val
|
|
|
|
|
|
def ms_ssim(a, b):
|
|
return float(np.mean([ms_ssim_channel(a[:, :, c], b[:, :, c]) for c in range(3)]))
|
|
|
|
|
|
def cmd_metrics(work, tags):
|
|
pics = pictures(work)
|
|
enc = jload(os.path.join(work, "encoded.json"))
|
|
rows = {}
|
|
for p in pics:
|
|
src = source_rgb(work, p)
|
|
for q in QUALITIES:
|
|
k = f"{p['id']}/q{q}"
|
|
if k not in enc:
|
|
continue
|
|
r = {}
|
|
for path, fname in (("float", f"{p['id']}.float.rgb"), ("int", f"{p['id']}.int.{tags[0]}.rgb")):
|
|
rgb = np.fromfile(os.path.join(work, "dec", f"q{q}", fname), dtype=np.uint8)
|
|
rgb = rgb.reshape(src.shape)
|
|
d = rgb.astype(np.int64) - src.astype(np.int64)
|
|
r[path] = {"sse": int((d * d).sum()), "samples": int(d.size), "ms_ssim": ms_ssim(rgb, src),
|
|
"rgb_sha256": sha256(rgb.tobytes())}
|
|
rows[k] = r
|
|
print(p["id"], flush=True)
|
|
jdump(rows, os.path.join(work, "metrics.json"))
|
|
|
|
|
|
# ---------------------------------------------------------------- BD-rate
|
|
|
|
def psnr(sse, n):
|
|
return 10 * math.log10(255.0 ** 2 * n / sse)
|
|
|
|
|
|
def ms_ssim_db(v):
|
|
return -10 * math.log10(1 - v)
|
|
|
|
|
|
def pchip_slopes(x, y):
|
|
"""Interior slopes of the monotone piecewise-cubic Hermite interpolant (weighted
|
|
harmonic mean of the neighbouring secants, zero at a change of sign), end slopes by
|
|
the three-point shape-preserving formula."""
|
|
n = len(x)
|
|
h = [x[i + 1] - x[i] for i in range(n - 1)]
|
|
d = [(y[i + 1] - y[i]) / h[i] for i in range(n - 1)]
|
|
m = [0.0] * n
|
|
for i in range(1, n - 1):
|
|
if d[i - 1] * d[i] > 0:
|
|
w1, w2 = 2 * h[i] + h[i - 1], h[i] + 2 * h[i - 1]
|
|
m[i] = (w1 + w2) / (w1 / d[i - 1] + w2 / d[i])
|
|
def end(h0, h1, d0, d1):
|
|
s = ((2 * h0 + h1) * d0 - h0 * d1) / (h0 + h1)
|
|
if s * d0 <= 0:
|
|
return 0.0
|
|
if d0 * d1 <= 0 and abs(s) > abs(3 * d0):
|
|
return 3 * d0
|
|
return s
|
|
m[0] = end(h[0], h[1], d[0], d[1])
|
|
m[-1] = end(h[-1], h[-2], d[-1], d[-2])
|
|
return m
|
|
|
|
|
|
def pchip_integral(x, y, a, b):
|
|
"""The integral from a to b of the interpolant through (x, y): each segment's cubic
|
|
Hermite polynomial integrated exactly."""
|
|
m = pchip_slopes(x, y)
|
|
total = 0.0
|
|
for i in range(len(x) - 1):
|
|
lo, hi = max(a, x[i]), min(b, x[i + 1])
|
|
if hi <= lo:
|
|
continue
|
|
h = x[i + 1] - x[i]
|
|
def F(t): # antiderivative in t = (u - x_i)/h, times h
|
|
h00 = t - t ** 3 + t ** 4 / 2
|
|
h10 = t ** 2 / 2 - 2 * t ** 3 / 3 + t ** 4 / 4
|
|
h01 = t ** 3 - t ** 4 / 2
|
|
h11 = -t ** 3 / 3 + t ** 4 / 4
|
|
return h * (h00 * y[i] + h10 * h * m[i] + h01 * y[i + 1] + h11 * h * m[i + 1])
|
|
total += F((hi - x[i]) / h) - F((lo - x[i]) / h)
|
|
return total
|
|
|
|
|
|
def bd_rate(r_ref, q_ref, r_test, q_test):
|
|
"""Bjontegaard delta rate (ITU-T VCEG-M33), log-rate interpolated as a monotone
|
|
piecewise cubic of quality and averaged over the overlap of the two quality ranges:
|
|
the percentage by which the test needs more rate than the reference at equal quality.
|
|
None when either curve is not strictly increasing in quality."""
|
|
def prep(r, qv):
|
|
pts = sorted(zip(qv, (math.log10(v) for v in r)))
|
|
xs = [p[0] for p in pts]
|
|
if any(b <= a for a, b in zip(xs, xs[1:])):
|
|
return None
|
|
return xs, [p[1] for p in pts]
|
|
A, B = prep(r_ref, q_ref), prep(r_test, q_test)
|
|
if A is None or B is None:
|
|
return None
|
|
lo, hi = max(A[0][0], B[0][0]), min(A[0][-1], B[0][-1])
|
|
if hi <= lo:
|
|
return None
|
|
ia = pchip_integral(A[0], A[1], lo, hi)
|
|
ib = pchip_integral(B[0], B[1], lo, hi)
|
|
return (10 ** ((ib - ia) / (hi - lo)) - 1) * 100
|
|
|
|
|
|
def cmd_bd(work):
|
|
pics = pictures(work)
|
|
enc = jload(os.path.join(work, "encoded.json"))
|
|
met = jload(os.path.join(work, "metrics.json"))
|
|
per = {}
|
|
for p in pics:
|
|
n = p["width"] * p["height"]
|
|
rf, ri, pf, pi, sf, si = [], [], [], [], [], []
|
|
for q in QUALITIES:
|
|
k = f"{p['id']}/q{q}"
|
|
rf.append(enc[k]["float_stream_bytes"] * 8 / n)
|
|
ri.append(enc[k]["int_stream_bytes"] * 8 / n)
|
|
pf.append(psnr(met[k]["float"]["sse"], met[k]["float"]["samples"]))
|
|
pi.append(psnr(met[k]["int"]["sse"], met[k]["int"]["samples"]))
|
|
sf.append(ms_ssim_db(met[k]["float"]["ms_ssim"]))
|
|
si.append(ms_ssim_db(met[k]["int"]["ms_ssim"]))
|
|
per[p["id"]] = {"bd_rate_psnr": bd_rate(rf, pf, ri, pi), "bd_rate_ms_ssim": bd_rate(rf, sf, ri, si)}
|
|
print(p["id"], per[p["id"]])
|
|
jdump(per, os.path.join(work, "bd.json"))
|
|
|
|
|
|
# ---------------------------------------------------------------- the record
|
|
|
|
def stats(v):
|
|
v = sorted(v)
|
|
n = len(v)
|
|
if n == 0:
|
|
return None
|
|
med = v[n // 2] if n % 2 else (v[n // 2 - 1] + v[n // 2]) / 2
|
|
return {"n": n, "median": med, "mean": sum(v) / n, "min": v[0], "max": v[-1]}
|
|
|
|
|
|
def summarise(rec):
|
|
"""Everything the record states beyond its raw readings, computed from them alone:
|
|
tools/wai_determinism_cost_verify.py recomputes each value independently."""
|
|
pics = rec["pictures"]["items"]
|
|
pts = rec["points"]
|
|
qs = sorted(int(q) for q in rec["models"])
|
|
out = {}
|
|
bd = {"psnr": {}, "ms_ssim": {}}
|
|
for p in pics:
|
|
n = p["width"] * p["height"]
|
|
ks = [f"{p['id']}/q{q}" for q in qs]
|
|
rf = [pts[k]["float_stream_bytes"] * 8 / n for k in ks]
|
|
ri = [pts[k]["int_stream_bytes"] * 8 / n for k in ks]
|
|
for metric, f in (("psnr", lambda r: psnr(r["sse"], r["samples"])), ("ms_ssim", lambda r: ms_ssim_db(r["ms_ssim"]))):
|
|
bd[metric][p["id"]] = bd_rate(rf, [f(pts[k]["float"]) for k in ks], ri, [f(pts[k]["int"]) for k in ks])
|
|
out["bd_rate_percent"] = {m: {"per_picture": v, **stats([x for x in v.values() if x is not None]),
|
|
"undefined": sorted(k for k, x in v.items() if x is None)} for m, v in bd.items()}
|
|
perq = {}
|
|
for q in qs:
|
|
ks = [f"{p['id']}/q{q}" for p in pics]
|
|
perq[f"q{q}"] = {
|
|
"float_bpp": stats([pts[k]["float_stream_bytes"] * 8 / (pts[k]["float"]["samples"] / 3) for k in ks]),
|
|
"rate_ratio": stats([pts[k]["int_stream_bytes"] / pts[k]["float_stream_bytes"] for k in ks]),
|
|
"float_psnr_db": stats([psnr(pts[k]["float"]["sse"], pts[k]["float"]["samples"]) for k in ks]),
|
|
"delta_psnr_db": stats([psnr(pts[k]["int"]["sse"], pts[k]["int"]["samples"]) - psnr(pts[k]["float"]["sse"], pts[k]["float"]["samples"]) for k in ks]),
|
|
"delta_ms_ssim_db": stats([ms_ssim_db(pts[k]["int"]["ms_ssim"]) - ms_ssim_db(pts[k]["float"]["ms_ssim"]) for k in ks]),
|
|
}
|
|
out["per_quality"] = perq
|
|
ref_int = {k: v["int"]["rgb_sha256"] for k, v in pts.items()}
|
|
xi = {}
|
|
for tag, d in rec["integer_decodes"].items():
|
|
r = d["results"]
|
|
xi[tag] = {"decodes": len(r),
|
|
"identical": sum(1 for k, h in r.items() if h == ref_int.get(k)),
|
|
"refused": sum(1 for h in r.values() if h.startswith("refused:")),
|
|
"missing": sorted(k for k in ref_int if k not in r)}
|
|
xi[tag]["differ"] = xi[tag]["decodes"] - xi[tag]["identical"] - xi[tag]["refused"]
|
|
out["integer_cross_platform"] = xi
|
|
out["integer_cross_platform_total"] = {"platforms": len(xi), **{f: sum(v[f] for v in xi.values()) for f in ("decodes", "identical", "differ", "refused")}}
|
|
xf = {}
|
|
for tag, d in rec["float_decodes"].items():
|
|
r = d["results"]
|
|
c = {"decodes": len(r), "identical": 0, "rgb_differ": 0, "indexes_differ": 0, "symbols_differ": 0,
|
|
"tables_differ": sorted(q for q, h in d.get("tables_sha256", {}).items() if h != rec["float_tables_sha256"][q])}
|
|
dps, dsym = [], []
|
|
for k, v in r.items():
|
|
ref = pts[k]
|
|
same = v["rgb_sha256"] == ref["float"]["rgb_sha256"]
|
|
c["identical"] += same
|
|
c["rgb_differ"] += not same
|
|
c["indexes_differ"] += v["indexes_sha256"] != ref["float_cpu_indexes_sha256"]
|
|
c["symbols_differ"] += v["symbols_sha256"] != ref["float_cpu_symbols_sha256"]
|
|
dp = psnr(v["sse"], ref["float"]["samples"]) - psnr(ref["float"]["sse"], ref["float"]["samples"])
|
|
if not same:
|
|
dps.append(dp)
|
|
if v["symbols_sha256"] != ref["float_cpu_symbols_sha256"]:
|
|
dsym.append(dp)
|
|
c["delta_psnr_db_where_differ"] = stats(dps)
|
|
c["delta_psnr_db_where_symbols_differ"] = stats(dsym)
|
|
c["missing"] = sorted(k for k in pts if k not in r)
|
|
xf[tag] = c
|
|
out["float_cross_platform"] = xf
|
|
def live(b):
|
|
"""energy-measurement §3.1: the rail advanced in two consecutive intervals, in at
|
|
least five and in at least four of every five, over at least 300 ms."""
|
|
return bool(b) and b["longest_run"] >= 2 and b["advanced"] >= 5 and 5 * b["advanced"] >= 4 * b["intervals"] and b["span_ns"] >= 300_000_000
|
|
|
|
tm = {}
|
|
for run in rec["timing"]:
|
|
rs = run["rounds"]
|
|
t = {"rounds": len(rs), "decodes_per_round": sorted({r["decodes"] for r in rs}),
|
|
"ms_per_decode": stats([r["child_wall_ns"] / r["decodes"] / 1e6 for r in rs]),
|
|
"load_avg_1min": stats([r["load_avg_before"][0] for r in rs] + [r["load_avg_after"][0] for r in rs])}
|
|
if all("task" in r for r in rs):
|
|
t["task_j_per_decode"] = stats([r["task"]["energy_nj"] / r["decodes"] / 1e9 for r in rs])
|
|
t["task_cpu_ms_per_decode"] = stats([r["task"]["cpu_ns"] / r["decodes"] / 1e6 for r in rs])
|
|
t["billed_j_per_decode"] = stats([r["task"]["billed_nj"] / r["decodes"] / 1e9 for r in rs])
|
|
for cn in ("rapl", "ioreport-gpu"):
|
|
if all(cn in r for r in rs):
|
|
ok = [r for r in rs if r[cn].get("verdict") in ("clean", "relative-only")
|
|
and all(live(r[cn][b]) for b in ("busy", "idle_before", "idle_after"))]
|
|
if cn == "ioreport-gpu":
|
|
ok = [r for r in ok if (r.get("gpu_time", {}).get("share") or 0) >= 0.99]
|
|
t[cn] = {"verdicts": [r[cn].get("verdict") for r in rs],
|
|
"live_rounds": sum(all(live(r[cn][b]) for b in ("busy", "idle_before", "idle_after")) for r in rs),
|
|
"ratio": stats([r[cn]["ratio"] for r in rs if r[cn].get("ratio") is not None]),
|
|
"usable_rounds": len(ok),
|
|
"marginal_j_per_decode": stats([r[cn]["marginal_nj"] / r["decodes"] / 1e9 for r in ok]),
|
|
"publishable": len(ok) >= 7 and all(r[cn]["verdict"] == "clean" for r in ok)}
|
|
if cn == "ioreport-gpu":
|
|
t[cn]["gpu_share"] = stats([r["gpu_time"]["share"] for r in rs if r.get("gpu_time", {}).get("share") is not None])
|
|
tm[run["label"]] = t
|
|
out["timing"] = tm
|
|
return out
|
|
|
|
|
|
def cmd_record(work, dest):
|
|
pics = jload(os.path.join(work, "pictures.json"))
|
|
enc = jload(os.path.join(work, "encoded.json"))
|
|
met = jload(os.path.join(work, "metrics.json"))
|
|
models = jload(os.path.join(work, "models.json"))
|
|
points, tables = {}, {}
|
|
for k, e in sorted(enc.items()):
|
|
q = k.split("/")[1]
|
|
if q in tables and tables[q] != e["float_tables_sha256"]:
|
|
sys.exit(f"{k}: other float tables than its quality's other pictures")
|
|
tables[q] = e["float_tables_sha256"]
|
|
points[k] = {"float_stream_bytes": e["float_stream_bytes"], "int_stream_bytes": e["int_stream_bytes"],
|
|
"int_payload_bytes": e["int_payload_bytes"], "float_stream_sha256": e["float_stream_sha256"],
|
|
"int_payload_sha256": e["int_payload_sha256"],
|
|
"float_cpu_indexes_sha256": e["float_cpu"]["indexes_sha256"],
|
|
"float_cpu_symbols_sha256": e["float_cpu"]["symbols_sha256"],
|
|
"float": met[k]["float"], "int": met[k]["int"]}
|
|
if met[k]["float"]["rgb_sha256"] != e["float_cpu"]["rgb_sha256"]:
|
|
sys.exit(f"{k}: the measured float decode is not the one the encoder wrote")
|
|
rec = {"format": "wai-determinism-cost/1", "procedure": "wai/research/determinism-cost.md",
|
|
"runs": "wai/research/determinism-cost/runs/", "measured": os.environ.get("DET_MEASURED", "2026-10-22"),
|
|
"pictures": {"manifest_sha256": pics["manifest_sha256"], "rule": pics["rule"],
|
|
"items": [{kk: p[kk] for kk in ("id", "file_sha256", "width", "height", "rgb_sha256")} for p in pics["pictures"]]},
|
|
"models": {q: {"checkpoint_sha256": v["sha256"], "checkpoint_name": v["url_name"], "N": v["N"], "M": v["M"], "set_sha256": v["set"]}
|
|
for q, v in sorted(models.items(), key=lambda kv: int(kv[0]))},
|
|
"float_tables_sha256": tables, "points": points,
|
|
"integer_decodes": {}, "float_decodes": {}, "timing": []}
|
|
xd = os.path.join(work, "xplat")
|
|
for f in sorted(os.listdir(xd)) if os.path.isdir(xd) else []:
|
|
if f.endswith(".tsv"):
|
|
tag = f[:-4]
|
|
meta = jload(os.path.join(xd, tag + ".json"))
|
|
res = {}
|
|
for line in open(os.path.join(xd, f)):
|
|
c = line.rstrip("\n").split("\t")
|
|
k = f"{os.path.basename(c[0]).split('.')[0]}/{os.path.basename(os.path.dirname(c[0]))}"
|
|
res[k] = f"refused:{c[2]}" if c[1] == "REFUSED" else c[1]
|
|
rec["integer_decodes"][tag] = {"platform": meta, "results": dict(sorted(res.items()))}
|
|
fd = os.path.join(work, "fdecode")
|
|
for f in sorted(os.listdir(fd)) if os.path.isdir(fd) else []:
|
|
if not f.endswith(".json"):
|
|
continue
|
|
d = jload(os.path.join(fd, f))
|
|
meta = jload(os.path.join(fd, f[:-5] + ".meta")) if os.path.exists(os.path.join(fd, f[:-5] + ".meta")) else {}
|
|
rec["float_decodes"][d["tag"]] = {"platform": {**meta, "device": d["device"], "threads": d["threads"], "torch": d["torch"],
|
|
"machine": d["machine"], "system": d["system"]},
|
|
"tables_sha256": d.get("tables_sha256", {}), "results": d["results"]}
|
|
ed = os.path.join(work, "energy")
|
|
for f in sorted(os.listdir(ed)) if os.path.isdir(ed) else []:
|
|
if f.endswith(".json"):
|
|
run = jload(os.path.join(ed, f))
|
|
# The child's command line, relative to each host's work area (no home or volume paths).
|
|
run["child"] = [a.split("/wai-build/", 1)[1] if "/wai-build/" in a else a for a in run["child"]]
|
|
rec["timing"].append(run)
|
|
rec["summary"] = summarise(rec)
|
|
jdump(rec, dest)
|
|
print(json.dumps(rec["summary"]["bd_rate_percent"]["psnr"] | {"per_picture": None}, indent=1))
|
|
|
|
|
|
def main(argv):
|
|
if len(argv) < 3:
|
|
sys.exit(__doc__)
|
|
cmd, work = argv[1], argv[2]
|
|
rest = argv[3:]
|
|
if cmd == "select":
|
|
cmd_select(work, rest[0])
|
|
elif cmd == "export":
|
|
allow = "--allow-new" in rest
|
|
cmd_export(work, [int(a) for a in rest if a != "--allow-new"] or QUALITIES, allow)
|
|
elif cmd == "encode":
|
|
cmd_encode(work, [int(a) for a in rest] or QUALITIES)
|
|
elif cmd == "jobs":
|
|
cmd_jobs(work, rest[0] if rest else "-")
|
|
elif cmd == "fdecode":
|
|
tag = rest[0]
|
|
dev = rest[rest.index("--device") + 1] if "--device" in rest else "cpu"
|
|
th = int(rest[rest.index("--threads") + 1]) if "--threads" in rest else 0
|
|
cmd_fdecode(work, tag, dev, th)
|
|
elif cmd == "fserve":
|
|
q = int(rest[0])
|
|
dev = rest[rest.index("--device") + 1] if "--device" in rest else "cpu"
|
|
th = int(rest[rest.index("--threads") + 1]) if "--threads" in rest else 0
|
|
ids = rest[rest.index("--ids") + 1].split(",") if "--ids" in rest else []
|
|
cmd_fserve(work, q, dev, th, ids)
|
|
elif cmd == "metrics":
|
|
cmd_metrics(work, rest or ["native"])
|
|
elif cmd == "bd":
|
|
cmd_bd(work)
|
|
elif cmd == "record":
|
|
cmd_record(work, rest[0])
|
|
else:
|
|
sys.exit(__doc__)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main(sys.argv)
|