dtourolle
a9271c4850
Best was a mixture of two experts and a gate, 110 GMAC a megapixel; Medium a single network at 48 that was softer on real edges. fb-combo (darkroom-denoise, 20 000 steps from fb-edges2, taught by the mixture with a quarter of its crops from the edge-rich parts of the frames) is Medium's shape and holds the mixture's edges on real photographs: edge PSNR within 0.04-0.06 dB at ISO 1600/6400/25600, more sharpness kept at all three, the chart's edge 0.89 photosites wide against 0.82. It is 0.27 dB short on smooth areas at ISO 25600. It becomes Best, and the methods are Bilinear, Fast and Best. Saved edits keep their numbers: 2, which was Medium, is now Best, and 3, which was Best, is past the end and reads as the default, Best. The network ships as mosaic-hq, a new name: the result cache keys a model by name and size, and this one is byte for byte the old Medium's size. Its tablet form (A16W16) lost 0.00 dB in simulated QDQ at every ISO and at most 0.09 dB across the noise bracket.
264 lines
12 KiB
Python
264 lines
12 KiB
Python
"""The body of quantise-models.sh; see there. Run through it, not directly."""
|
||
|
||
import glob
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import onnx
|
||
import onnxruntime as ort
|
||
from onnx import version_converter
|
||
from onnxruntime.quantization import CalibrationDataReader, CalibrationMethod, QuantType, quantize_static
|
||
from onnxruntime.quantization.calibrate import create_calibrator, save_tensors_data
|
||
from onnxruntime.quantization.execution_providers.qnn import get_qnn_qdq_config
|
||
from PIL import Image
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
import htp_graph # noqa: E402
|
||
|
||
MODELS = Path(__file__).resolve().parents[1] / "models"
|
||
PHOTOS = 300 # calibration photographs; face crops and keypoints come from fewer
|
||
CHUNK = 4 # inputs whose activations are held at once (scrfd_10g: ~1 GB each)
|
||
|
||
Q = QuantType
|
||
FORMS = {"int8": (Q.QUInt8, Q.QInt8), "a16w8": (Q.QUInt16, Q.QInt8), "a16w16": (Q.QUInt16, Q.QInt16)}
|
||
|
||
# Per model: where it lives, the form `Rung::form` gives its role, the exact
|
||
# rewrites its graph needs, and nodes that stay float on the CPU because one
|
||
# scale cannot serve the tensor (the segmenter's rows: boxes in pixels beside
|
||
# scores in 0..1) or because the HTP's 16-bit arithmetic drifts there (the
|
||
# scene model's attention). Measured, inference.md §1.5.
|
||
TABLE = {
|
||
"scrfd_500m_640": dict(dir="face", form="a16w8", feed="scrfd"),
|
||
"scrfd_2.5g_640": dict(dir="face", form="a16w8", feed="scrfd"),
|
||
"scrfd_10g_640": dict(dir="face", form="a16w8", feed="scrfd"),
|
||
"2d106det_b1": dict(dir="face", form="a16w8", feed="landmarks"),
|
||
"yolo26n-seg": dict(dir="segment", form="a16w16", feed="yolo", float_from="/model.23/Concat_4"),
|
||
"yolo26s-sem-ade20k": dict(dir="scene", form="a16w16", feed="yolo",
|
||
float_nodes=["/model.10/m/m.0/attn/MatMul", "/model.10/m/m.0/attn/Softmax",
|
||
"/model.10/m/m.0/attn/MatMul_1"]),
|
||
"migan-512": dict(dir="inpaint", form="a16w16", feed="migan"),
|
||
"xfeat-1024": dict(dir="keypoints", form="int8", feed="xfeat", rewrites=["unfold", "resize"]),
|
||
"xfeat-768": dict(dir="keypoints", form="int8", feed="xfeat", rewrites=["unfold", "resize"]),
|
||
"mosaic-fast-1408": dict(dir="denoise", form="a16w16", feed=None, rewrites=["bayer"]),
|
||
# Best since 0.24: one network, exported as Fast is, so the same rewrite.
|
||
"mosaic-hq-1408": dict(dir="denoise", form="a16w16", feed=None, rewrites=["bayer"]),
|
||
}
|
||
|
||
|
||
# ---- the app's samplers (dr-face Letterbox, dr-segment Letterbox::sample, align.rs) ----
|
||
|
||
def load(p):
|
||
return np.asarray(Image.open(p).convert("RGB"), np.float32) / 255.0
|
||
|
||
|
||
def bilinear(img, sx, sy):
|
||
h, w = img.shape[:2]
|
||
x0, y0 = np.floor(sx).astype(np.int64), np.floor(sy).astype(np.int64)
|
||
fx, fy = (sx - x0)[..., None], (sy - y0)[..., None]
|
||
c = lambda a, n: np.clip(a, 0, n - 1) # noqa: E731 - neighbours clamp at the edge
|
||
top = img[c(y0, h), c(x0, w)] * (1 - fx) + img[c(y0, h), c(x0 + 1, w)] * fx
|
||
bot = img[c(y0 + 1, h), c(x0, w)] * (1 - fx) + img[c(y0 + 1, h), c(x0 + 1, w)] * fx
|
||
return top * (1 - fy) + bot * fy
|
||
|
||
|
||
def chw(x):
|
||
return np.ascontiguousarray(x.transpose(2, 0, 1))[None].astype(np.float32)
|
||
|
||
|
||
def letterbox(img, edge, yolo):
|
||
h, w = img.shape[:2]
|
||
s = min(edge / w, edge / h)
|
||
px, py = (edge - w * s) / 2, (edge - h * s) / 2
|
||
ix, iy = np.meshgrid(np.arange(edge) + 0.5, np.arange(edge) + 0.5)
|
||
if yolo: # semantic.rs: no -0.5, pad 0.5, 0..1
|
||
sx, sy = (ix - px) / s, (iy - py) / s
|
||
out = bilinear(img, sx, sy)
|
||
out[(sx < 0) | (sx >= w) | (sy < 0) | (sy >= h)] = 0.5
|
||
else: # detect.rs: -0.5, pad 114, (v·255 − 127.5)/128
|
||
sx, sy = (ix - px) / s - 0.5, (iy - py) / s - 0.5
|
||
out = (bilinear(img, sx, sy) * 255 - 127.5) / 128
|
||
out[(sx < -0.5) | (sx > w - 0.5) | (sy < -0.5) | (sy > h - 0.5)] = (114 - 127.5) / 128
|
||
return chw(out), (s, px, py)
|
||
|
||
|
||
def crop_box(img, x0, y0, bw, bh, ow, oh):
|
||
"""align.rs crop_box: output (u+.5) → source, −0.5, bilinear, outside black."""
|
||
u, v = np.meshgrid(np.arange(ow) + 0.5, np.arange(oh) + 0.5)
|
||
sx, sy = x0 + u * bw / ow - 0.5, y0 + v * bh / oh - 0.5
|
||
out = bilinear(img, sx, sy)
|
||
h, w = img.shape[:2]
|
||
out[(sx < -1) | (sx > w) | (sy < -1) | (sy > h)] = 0
|
||
return out
|
||
|
||
|
||
def scrfd_boxes(outs, s, px, py):
|
||
"""detect.rs decode: score ≥ 0.5, greedy NMS at 0.4, min side 24 px."""
|
||
fmc = len(outs) // 3
|
||
boxes, scores = [], []
|
||
for i, st in enumerate([8, 16, 32, 64][:fmc]):
|
||
sc, bx = outs[i].reshape(-1), outs[fmc + i].reshape(-1, 4)
|
||
idx = np.nonzero(sc >= 0.5)[0]
|
||
cell = idx // 2
|
||
cx, cy = (cell % (640 // st)) * st, (cell // (640 // st)) * st
|
||
boxes.append(np.stack([cx - bx[idx, 0] * st, cy - bx[idx, 1] * st, cx + bx[idx, 2] * st, cy + bx[idx, 3] * st], 1))
|
||
scores.append(sc[idx])
|
||
b, sc = np.concatenate(boxes), np.concatenate(scores)
|
||
keep = []
|
||
for i in np.argsort(-sc):
|
||
x0 = np.maximum(b[i, 0], b[keep, 0]); y0 = np.maximum(b[i, 1], b[keep, 1])
|
||
x1 = np.minimum(b[i, 2], b[keep, 2]); y1 = np.minimum(b[i, 3], b[keep, 3])
|
||
inter = np.clip(x1 - x0, 0, None) * np.clip(y1 - y0, 0, None)
|
||
area = lambda r: (r[..., 2] - r[..., 0]) * (r[..., 3] - r[..., 1]) # noqa: E731
|
||
if not keep or (inter / (area(b[i]) + area(b[keep]) - inter)).max() <= 0.4:
|
||
keep.append(i)
|
||
b = (b[keep] - [px, py, px, py]) / s
|
||
return b[np.minimum(b[:, 2] - b[:, 0], b[:, 3] - b[:, 1]) >= 32]
|
||
|
||
|
||
# ---- one calibration input per photograph (or per face), as the app makes it ----
|
||
|
||
def feeds(kind, photos, model):
|
||
name = model.get_inputs()[0].name
|
||
if kind == "scrfd":
|
||
for p in photos:
|
||
yield {name: letterbox(load(p), 640, False)[0]}
|
||
elif kind == "landmarks": # landmarks.rs: 1.5× the box, square, 0..255
|
||
det = ort.InferenceSession(str(MODELS / "face/scrfd_10g_640.onnx"), providers=["CPUExecutionProvider"])
|
||
for p in photos:
|
||
img = load(p)
|
||
x, ctx = letterbox(img, 640, False)
|
||
for b in scrfd_boxes(det.run(None, {"input.1": x}), *ctx):
|
||
cx, cy, side = (b[0] + b[2]) / 2, (b[1] + b[3]) / 2, 1.5 * max(b[2] - b[0], b[3] - b[1])
|
||
yield {name: chw(crop_box(img, cx - side / 2, cy - side / 2, side, side, 192, 192) * 255)}
|
||
elif kind == "yolo":
|
||
for p in photos:
|
||
yield {name: letterbox(load(p), 640, True)[0]}
|
||
elif kind == "migan": # migan.rs: ch0 = known − 0.5, ch1–3 = (rgb·2 − 1)·known
|
||
rng = np.random.default_rng(7)
|
||
for p in photos:
|
||
img = load(p)
|
||
h, w = img.shape[:2]
|
||
e = min(h, w)
|
||
sq = Image.fromarray((img[(h - e) // 2:(h + e) // 2, (w - e) // 2:(w + e) // 2] * 255).astype(np.uint8))
|
||
img = np.asarray(sq.resize((512, 512), Image.BILINEAR), np.float32) / 255
|
||
known = np.ones((512, 512), np.float32)
|
||
for _ in range(rng.integers(1, 3)): # a panorama's unknown border: a wedge along one edge
|
||
side = rng.integers(4)
|
||
depth = np.linspace(rng.integers(20, 110), rng.integers(20, 110), 512).astype(int)
|
||
edge = np.arange(512)[:, None] < depth[None, :] # [depth, along]: inside the wedge
|
||
wedge = edge if side % 2 == 0 else edge[::-1] # top / bottom of a column
|
||
known[wedge if side < 2 else wedge.T] = 0 # or left / right of a row
|
||
x = np.concatenate([(known - 0.5)[None], ((img * 2 - 1) * known[..., None]).transpose(2, 0, 1)])[None]
|
||
yield {name: x.astype(np.float32)}
|
||
elif kind == "xfeat": # xfeat.rs: grey 0..1, shrink to fit, top-left, zero pad
|
||
_, _, H, W = [d if isinstance(d, int) else 1 for d in model.get_inputs()[0].shape]
|
||
for p in photos:
|
||
g = load(p).mean(2) # a display-rendered photograph is already the app's (R+G+B)/3 ^ 1/2.2
|
||
h, w = g.shape
|
||
s = min(W / w, H / h, 1.0)
|
||
if s < 1:
|
||
g = np.asarray(Image.fromarray(g).resize((round(w * s), round(h * s)), Image.BOX))
|
||
pad = np.zeros((H, W), np.float32)
|
||
pad[:g.shape[0], :g.shape[1]] = g
|
||
yield {name: pad[None, None]}
|
||
|
||
|
||
class Items(CalibrationDataReader):
|
||
def __init__(self, items):
|
||
self.it = iter(items)
|
||
|
||
def get_next(self):
|
||
return next(self.it, None)
|
||
|
||
|
||
def calibrate(path, items, cache):
|
||
"""Min/max ranges in chunks — every ORT calibrator holds all activations
|
||
until it folds them, and the others measurably degrade the result."""
|
||
cal = create_calibrator(Path(path), None, augmented_model_path=f"{path}.aug.onnx",
|
||
calibrate_method=CalibrationMethod.MinMax)
|
||
batch, n = [], 0
|
||
for item in items:
|
||
batch.append(item)
|
||
n += 1
|
||
if len(batch) == CHUNK:
|
||
cal.collect_data(Items(batch))
|
||
batch = []
|
||
if batch:
|
||
cal.collect_data(Items(batch))
|
||
save_tensors_data(cal.compute_data(), cache)
|
||
os.remove(f"{path}.aug.onnx")
|
||
return n
|
||
|
||
|
||
def downstream(m, start):
|
||
names, live = set(), set()
|
||
for n in m.graph.node:
|
||
if n.name == start or any(i in live for i in n.input):
|
||
names.add(n.name)
|
||
live.update(n.output)
|
||
return sorted(names)
|
||
|
||
|
||
def main():
|
||
args = sys.argv[1:]
|
||
ranges = None
|
||
if args[:1] == ["--ranges"]:
|
||
ranges, args = args[1], args[2:]
|
||
photo_dir = None
|
||
else:
|
||
photo_dir, args = args[0], args[1:]
|
||
names = args or [n for n in TABLE if TABLE[n]["feed"]]
|
||
photos = []
|
||
if photo_dir:
|
||
photos = sorted(p for e in ("jpg", "jpeg", "JPG", "JPEG", "png")
|
||
for p in glob.glob(os.path.join(photo_dir, "**", f"*.{e}"), recursive=True))[:PHOTOS]
|
||
if len(photos) < 50:
|
||
sys.exit(f"only {len(photos)} photographs under {photo_dir}; calibration wants hundreds")
|
||
print(f"==> calibrating on {len(photos)} photographs")
|
||
|
||
for stem in names:
|
||
spec = TABLE[stem]
|
||
src = MODELS / spec["dir"] / f"{stem}.onnx"
|
||
out = src.with_name(f"{stem}.{spec['form']}.onnx")
|
||
work = src.with_name(f"{stem}.quant-work.onnx")
|
||
print(f" {stem} -> {out.name}")
|
||
m = onnx.load(src)
|
||
if next(o.version for o in m.opset_import if o.domain in ("", "ai.onnx")) < 13:
|
||
m = version_converter.convert_version(m, 17) # per-channel QDQ needs 13
|
||
m.ir_version = 8
|
||
onnx.save(m, work)
|
||
if spec.get("rewrites"):
|
||
onnx.save(htp_graph.rewrite(str(work), spec["rewrites"]), work)
|
||
model = ort.InferenceSession(str(work), providers=["CPUExecutionProvider"])
|
||
cache = str(work) + ".ranges"
|
||
if spec["feed"] is None:
|
||
if not ranges:
|
||
sys.exit(f"{stem} is calibrated on mosaics, not photographs: pass --ranges")
|
||
cache = ranges
|
||
else:
|
||
n = calibrate(str(work), feeds(spec["feed"], photos, model), cache)
|
||
print(f" {n} calibration inputs")
|
||
act, wt = FORMS[spec["form"]]
|
||
# The config needs a reader only to exist; the ranges come from `cache`.
|
||
zeros = {i.name: np.zeros([d if isinstance(d, int) else 1 for d in i.shape], np.float32)
|
||
for i in model.get_inputs()}
|
||
cfg = get_qnn_qdq_config(str(work), Items([zeros]), activation_type=act, weight_type=wt, per_channel=True)
|
||
exclude = list(cfg.nodes_to_exclude or []) + spec.get("float_nodes", [])
|
||
if spec.get("float_from"):
|
||
exclude += downstream(onnx.load(work), spec["float_from"])
|
||
quantize_static(str(work), str(out), None, quant_format=cfg.quant_format,
|
||
op_types_to_quantize=cfg.op_types_to_quantize, per_channel=True,
|
||
activation_type=act, weight_type=wt, nodes_to_exclude=exclude,
|
||
calibrate_method=CalibrationMethod.MinMax, extra_options=cfg.extra_options,
|
||
calibration_cache_path=cache)
|
||
os.remove(work)
|
||
if cache != ranges:
|
||
os.remove(cache)
|
||
print(f" {out.name}: {out.stat().st_size // 1024} KB")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|