File size: 2,917 Bytes
1004324 | 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 | #!/usr/bin/env python3
"""Weight-only per-channel INT8 for ddcolor-paper-tiny.onnx (OptiMix).
Only weight initializers consumed by Conv, MatMul or Gemm are quantised (symmetric INT8, one scale per
output channel). Each becomes int8 data + float scale + DequantizeLinear(axis) feeding the original
tensor name, so compute stays float32. Einsum, Erf, LayerNormalization, norms, biases and small tensors
are left untouched in float32. DequantizeLinear with axis needs opset >= 13; the graph is opset 17.
Usage: quantize_ddcolor_int8.py SRC.onnx DEST.onnx
"""
import hashlib, sys
import numpy as np, onnx
from onnx import helper, numpy_helper
src, dest = sys.argv[1], sys.argv[2]
m = onnx.load(src)
g = m.graph
inits = {i.name: i for i in g.initializer}
consumers = {}
for n in g.node:
for k, name in enumerate(n.input):
consumers.setdefault(name, []).append((n, k))
def axis_for(name, arr):
uses = consumers.get(name, [])
if not uses:
return None
axes = set()
for n, k in uses:
if n.op_type == "Conv" and k == 1 and arr.ndim == 4:
axes.add(0) # [out, in, kh, kw]
elif n.op_type == "MatMul" and k == 1 and arr.ndim == 2:
axes.add(1) # B is [K, N]; channel = N
elif n.op_type == "Gemm" and k == 1 and arr.ndim == 2:
tb = next((a.i for a in n.attribute if a.name == "transB"), 0)
axes.add(0 if tb else 1)
else:
return None # any other consumer: keep float
return axes.pop() if len(axes) == 1 else None
new_inits, new_nodes, done, kept = [], [], 0, 0
for init in list(g.initializer):
arr = numpy_helper.to_array(init)
ax = axis_for(init.name, arr) if arr.dtype == np.float32 and arr.size >= 4096 else None
if ax is None:
new_inits.append(init); kept += arr.nbytes; continue
red = tuple(i for i in range(arr.ndim) if i != ax)
scale = np.maximum(np.abs(arr).max(axis=red) / 127.0, 1e-12).astype(np.float32)
shp = [1] * arr.ndim; shp[ax] = -1
q = np.clip(np.round(arr / scale.reshape(shp)), -127, 127).astype(np.int8)
new_inits += [numpy_helper.from_array(q, init.name + "_q"),
numpy_helper.from_array(scale, init.name + "_scale"),
numpy_helper.from_array(np.zeros_like(scale, dtype=np.int8), init.name + "_zp")]
new_nodes.append(helper.make_node("DequantizeLinear",
[init.name + "_q", init.name + "_scale", init.name + "_zp"], [init.name],
name=init.name + "_dq", axis=ax))
done += 1
del g.initializer[:]; g.initializer.extend(new_inits)
nodes = new_nodes + list(g.node)
del g.node[:]; g.node.extend(nodes)
onnx.checker.check_model(m)
onnx.save(m, dest)
data = open(dest, "rb").read()
print(f"quantised {done} weights; float32 kept {kept/1e6:.1f} MB")
print(f"{hashlib.sha256(data).hexdigest()} {dest} ({len(data)} bytes)")
|