File size: 6,199 Bytes
e2f3b24 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """
Answer the open question directly: is it better to QUANTIZE-THEN-CONVERT or
CONVERT-THEN-QUANTIZE for this model?
Variants compared, all ending in a runnable .mlpackage:
A fp16 convert only (baseline)
B convert -> int8 weights coremltools.optimize.coreml.linear_quantize_weights
C convert -> 6-bit palette coremltools.optimize.coreml.palettize_weights (kmeans)
D int8 weights -> convert coremltools.optimize.torch PostTrainingQuantizer, then trace+convert
Each is scored on real val faces by agreement with the *unquantized PyTorch*
argmax (the thing users actually see), plus mean IoU against the fp16 baseline's
labels, ANE latency, and on-disk size.
Usage: compare_quant_order.py --weights <best.pt> [--imgsz 512] [--n-val 12]
"""
import argparse, glob, json, os, time
import numpy as np
import torch
from PIL import Image
import coremltools as ct
import coreml_patch # noqa: F401
from ultralytics import YOLO
VAL = "/Users/ari/FaceSegmentation/dataset_celebamaskhq_semantic/images/val"
class LogitsOnly(torch.nn.Module):
def __init__(self, m):
super().__init__()
self.m = m
def forward(self, x):
z = self.m(x)
return z[0] if isinstance(z, (list, tuple)) else z
def convert(mod, R):
ts = torch.jit.trace(mod, torch.rand(1, 3, R, R), strict=False)
return ct.convert(
ts,
inputs=[ct.ImageType(name="image", shape=(1, 3, R, R), scale=1 / 255.0,
bias=[0, 0, 0], color_layout=ct.colorlayout.RGB)],
outputs=[ct.TensorType(name="logits")],
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT16,
compute_units=ct.ComputeUnit.CPU_AND_NE,
minimum_deployment_target=ct.target.iOS17,
)
def dirsize_mb(p):
return round(sum(os.path.getsize(f) for f in glob.glob(p + "/**/*", recursive=True)
if os.path.isfile(f)) / 1e6, 2)
def score(path, ref_labels, imgs):
m = ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_NE)
agree, ious = [], []
for img, ref in zip(imgs, ref_labels):
got = np.asarray(m.predict({"image": img})["logits"], dtype=np.float32).argmax(1)[0]
agree.append(float((got == ref).mean()))
per = []
for c in np.union1d(np.unique(ref), np.unique(got)):
inter = np.logical_and(ref == c, got == c).sum()
union = np.logical_or(ref == c, got == c).sum()
if union:
per.append(inter / union)
ious.append(float(np.mean(per)) if per else 1.0)
m.predict({"image": imgs[0]})
t0 = time.time()
for _ in range(30):
m.predict({"image": imgs[0]})
return {"argmax_agreement": round(float(np.mean(agree)), 5),
"mean_iou_vs_torch": round(float(np.mean(ious)), 5),
"latency_ms": round((time.time() - t0) / 30 * 1000, 2),
"size_mb": dirsize_mb(path)}
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--weights", required=True)
ap.add_argument("--imgsz", type=int, default=512)
ap.add_argument("--n-val", type=int, default=12)
ap.add_argument("--out", default="/Users/ari/FaceSegmentation/exports_semantic/quant_order")
args = ap.parse_args()
os.makedirs(args.out, exist_ok=True)
R = args.imgsz
files = sorted(glob.glob(os.path.join(VAL, "*.jpg")))[: args.n_val]
imgs = [Image.open(f).convert("RGB").resize((R, R), Image.BILINEAR) for f in files]
base_model = LogitsOnly(YOLO(args.weights).model).eval().float().cpu()
# reference labels from unquantized PyTorch
ref = []
for img in imgs:
x = torch.from_numpy(np.asarray(img, np.float32) / 255.0).permute(2, 0, 1)[None]
with torch.no_grad():
ref.append(base_model(x).float().numpy().argmax(1)[0])
results = {}
mlbase = convert(base_model, R)
pA = f"{args.out}/A_fp16.mlpackage"; mlbase.save(pA)
results["A_fp16_convert_only"] = score(pA, ref, imgs)
from coremltools.optimize.coreml import (
OpLinearQuantizerConfig, OpPalettizerConfig, OptimizationConfig,
linear_quantize_weights, palettize_weights)
try:
q = linear_quantize_weights(mlbase, OptimizationConfig(
global_config=OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8")))
p = f"{args.out}/B_convert_then_int8.mlpackage"; q.save(p)
results["B_convert_then_int8"] = score(p, ref, imgs)
except Exception as e:
results["B_convert_then_int8"] = {"error": repr(e)[:200]}
try:
q = palettize_weights(mlbase, OptimizationConfig(
global_config=OpPalettizerConfig(mode="kmeans", nbits=6)))
p = f"{args.out}/C_convert_then_palette6.mlpackage"; q.save(p)
results["C_convert_then_palette6"] = score(p, ref, imgs)
except Exception as e:
results["C_convert_then_palette6"] = {"error": repr(e)[:200]}
# D: quantize the TORCH model first, then convert.
try:
from coremltools.optimize.torch.quantization import (
PostTrainingQuantizer, PostTrainingQuantizerConfig)
tmod = LogitsOnly(YOLO(args.weights).model).eval().float().cpu()
cfg = PostTrainingQuantizerConfig.from_dict(
{"global_config": {"weight_dtype": "int8", "granularity": "per_channel"}})
tq = PostTrainingQuantizer(tmod, cfg).compress()
p = f"{args.out}/D_int8_then_convert.mlpackage"
convert(tq.eval(), R).save(p)
results["D_int8_then_convert"] = score(p, ref, imgs)
except Exception as e:
results["D_int8_then_convert"] = {"error": repr(e)[:300]}
print(json.dumps(results, indent=2))
with open(f"{args.out}/comparison.json", "w") as f:
json.dump({"weights": args.weights, "imgsz": R, "results": results}, f, indent=2)
ok = {k: v for k, v in results.items() if "error" not in v}
if ok:
best = max(ok.items(), key=lambda kv: (kv[1]["argmax_agreement"], -kv[1]["size_mb"]))
print("\nBEST by fidelity:", best[0], best[1])
small = min(ok.items(), key=lambda kv: kv[1]["size_mb"])
print("SMALLEST :", small[0], small[1])
if __name__ == "__main__":
main()
|