| |
| """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) |
| elif n.op_type == "MatMul" and k == 1 and arr.ndim == 2: |
| axes.add(1) |
| 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 |
| 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)") |
|
|
