| """Rebuild faceage_dino_fp32.onnx as a graph of few, large kernels for ONNX Runtime's CUDA provider.
|
|
|
| python rebuild_fused.py [OUT.onnx] [--fp32]
|
|
|
| The network and its weights are the source's; only the way the graph is written changes.
|
| The weights are read out of source/faceage_dino_fp32.onnx (the source's PyTorch
|
| checkpoint is gated on the hub), and a new graph is written by hand:
|
|
|
| - The input is Bx3x224x224 fp32 with the batch dynamic, as before. Every Reshape takes a
|
| constant shape with -1 for the batch, so no Shape, Gather or Concat runs on the host.
|
| - The rotary tables are computed once here in float64 from the source's own angle table
|
| and stored as fp16 constants. The five prefix tokens (CLS and four registers) get rows
|
| of cos 1, sin 0, so one com.microsoft RotaryEmbedding covers all 201 tokens.
|
| - Attention is com.microsoft MultiHeadAttention in fp16, fed q, k and v from three
|
| separate projections, so nothing is split (a Split on the last axis was the slowest
|
| kernel in a packed-qkv version of this graph).
|
| - The MatMuls are fp16. The residual stream and every LayerNorm stay fp32, because of
|
| DINOv3's massive activation (about 1.56e5 on channel 757, written by block 0's MLP).
|
| The residual add and the LayerNorm after it are one com.microsoft SkipLayerNormalization,
|
| which also adds the output projection's bias, in fp32. Each block has four Casts: down to
|
| fp16 after each of its two LayerNorms, and back up to fp32 after each of its two output
|
| projections.
|
| - The residual stream is carried scaled by S = 1/8, with every LayerNorm's epsilon scaled
|
| by S^2, which leaves each LayerNorm's output exactly what it was. The layer scales,
|
| times S, are folded into the output projections' weights and biases. Block 0's MLP then
|
| writes about 1.95e4 into channel 757, inside fp16's 65,504; every other branch output
|
| on the WIDER FACE crops stays under 300.
|
| - v's bias is folded into o's bias: a row of attention weights sums to one, so a constant
|
| added to every value comes through attention unchanged.
|
| - The last block computes the CLS query only, and its MLP for the CLS token only, since
|
| only the CLS token is read after it.
|
|
|
| With --fp32 the same graph is written without the casts, all in fp32, as a check on the
|
| structure against the source.
|
| """
|
|
|
| import hashlib
|
| import os
|
| import sys
|
|
|
| import numpy as np
|
| import onnx
|
| from onnx import TensorProto, helper, numpy_helper
|
|
|
| SRC = r"E:/projects/faceage-onnx/source/faceage_dino_fp32.onnx"
|
| DST = r"E:/projects/faceage-onnx/faceage-dino-vitl16-fused-fp16.onnx"
|
|
|
| S = 0.125
|
| EPS = 1e-5
|
| HEADS = 16
|
| DIM = 1024
|
| TOKENS = 201
|
| PREFIX = 5
|
| P = "model.backbone."
|
|
|
|
|
| def sha256(path):
|
| d = hashlib.sha256()
|
| with open(path, "rb") as f:
|
| for chunk in iter(lambda: f.read(1 << 20), b""):
|
| d.update(chunk)
|
| return d.hexdigest()
|
|
|
|
|
| def read_source():
|
| m = onnx.load(SRC)
|
| w = {i.name: numpy_helper.to_array(i).astype(np.float32) for i in m.graph.initializer}
|
| consumers = {}
|
| for n in m.graph.node:
|
| for i in n.input:
|
| consumers.setdefault(i, []).append(n)
|
|
|
| mm = [n for n in m.graph.node if n.op_type == "MatMul" and n.input[1] in w]
|
| assert len(mm) == 24 * 6, len(mm)
|
| roles = ["q_proj", "k_proj", "v_proj", "o_proj", "up_proj", "down_proj"]
|
| layers = []
|
| for L in range(24):
|
| layer = {}
|
| for k, role in enumerate(roles):
|
| node = mm[L * 6 + k]
|
| layer[role] = w[node.input[1]]
|
|
|
| adds = [c for c in consumers.get(node.output[0], []) if c.op_type == "Add"]
|
| sub = "mlp" if role in ("up_proj", "down_proj") else "attention"
|
| bias = f"{P}model.layer.{L}.{sub}.{role}.bias"
|
| if role == "k_proj":
|
| assert not adds and bias not in w, (L, role)
|
| else:
|
| assert adds and adds[0].input[1] == bias, (L, role, [a.input for a in adds])
|
| layer[role + ".bias"] = w[bias]
|
| for nm in ("norm1.weight", "norm1.bias", "norm2.weight", "norm2.bias",
|
| "layer_scale1.lambda1", "layer_scale2.lambda1"):
|
| layer[nm] = w[f"{P}model.layer.{L}.{nm}"]
|
| layers.append(layer)
|
| return w, layers
|
|
|
|
|
| class G:
|
| def __init__(self):
|
| self.nodes, self.inits, self.n = [], [], 0
|
|
|
| def const(self, name, arr):
|
| self.inits.append(numpy_helper.from_array(arr, name))
|
| return name
|
|
|
| def op(self, op_type, inputs, n_out=1, domain="", name=None, **attrs):
|
| self.n += 1
|
| base = name or f"{op_type.lower()}_{self.n}"
|
| outs = [base] if n_out == 1 else [f"{base}_{k}" for k in range(n_out)]
|
| self.nodes.append(helper.make_node(op_type, inputs, outs, name=base, domain=domain or None, **attrs))
|
| return outs[0] if n_out == 1 else outs
|
|
|
|
|
| def build(fp32_only=False):
|
| w, layers = read_source()
|
| lo = np.float32 if fp32_only else np.float16
|
| g = G()
|
|
|
| def down(x, name):
|
| return x if fp32_only else g.op("Cast", [x], name=name, to=TensorProto.FLOAT16)
|
|
|
| def up(x, name):
|
| return x if fp32_only else g.op("Cast", [x], name=name, to=TensorProto.FLOAT)
|
|
|
| shape = lambda name, dims: g.const(name, np.array(dims, dtype=np.int64))
|
| sh_tok3 = shape("shape_tokens3d", [-1, TOKENS, DIM])
|
| sh_cls3 = shape("shape_cls3d", [-1, 1, DIM])
|
| sh_2d = shape("shape_rows2d", [-1, DIM])
|
| sh_patch = shape("shape_patches", [-1, DIM, 196])
|
| first = g.const("slice_first", np.array([0], np.int64))
|
| one = g.const("slice_one", np.array([1], np.int64))
|
| axis1 = g.const("slice_axis1", np.array([1], np.int64))
|
|
|
|
|
| tile = w["tile"].astype(np.float64)
|
| assert tile.shape == (196, 64) and np.array_equal(tile[:, :32], tile[:, 32:])
|
| ang = tile[:, :32]
|
| cos = np.concatenate([np.ones((PREFIX, 32)), np.cos(ang)]).astype(lo)
|
| sin = np.concatenate([np.zeros((PREFIX, 32)), np.sin(ang)]).astype(lo)
|
| cos_c, sin_c = g.const("rope_cos", cos), g.const("rope_sin", sin)
|
| pos0 = g.const("rope_position_offset", np.array([0], dtype=np.int64))
|
|
|
| def rope(t, name):
|
| return g.op("RotaryEmbedding", [t, pos0, cos_c, sin_c], domain="com.microsoft",
|
| name=name, num_heads=HEADS, rotary_embedding_dim=64)
|
|
|
|
|
|
|
| conv_w = g.const("patch_embed.weight", (w[P + "embeddings.patch_embeddings.weight"] * S).astype(np.float32))
|
| table = np.zeros((1, TOKENS, DIM), dtype=np.float32)
|
| table[0, 0] = w[P + "embeddings.cls_token"][0, 0]
|
| table[0, 1:PREFIX] = w[P + "embeddings.register_tokens"][0]
|
| table[0, PREFIX:] = w[P + "embeddings.patch_embeddings.bias"][None, :]
|
| table_c = g.const("token_table", (table * S).astype(np.float32))
|
|
|
| x = g.op("Conv", ["face", conv_w], name="patch_embed", strides=[16, 16], kernel_shape=[16, 16])
|
| x = g.op("Reshape", [x, sh_patch], name="patches_flat")
|
| x = g.op("Transpose", [x], name="patches_tokens", perm=[0, 2, 1])
|
| pads = g.const("prefix_pads", np.array([0, PREFIX, 0, 0, 0, 0], dtype=np.int64))
|
| x = g.op("Pad", [x, pads], name="prefix_slots")
|
|
|
|
|
|
|
|
|
|
|
| eps = EPS * S * S
|
| resid = None
|
| pending_bias = None
|
| for L, ly in enumerate(layers):
|
| last = L == len(layers) - 1
|
| c = lambda nm, arr, dt=np.float32: g.const(f"l{L}.{nm}", np.ascontiguousarray(arr).astype(dt))
|
| n1w, n1b = c("norm1.weight", ly["norm1.weight"]), c("norm1.bias", ly["norm1.bias"])
|
| if L == 0:
|
|
|
| ln, _, _, resid = g.op("SkipLayerNormalization", [x, table_c, n1w, n1b], n_out=4,
|
| domain="com.microsoft", name="l0.norm1", epsilon=eps)
|
| else:
|
| ln, _, _, resid = g.op("SkipLayerNormalization", [branch, resid, n1w, n1b, pending_bias], n_out=4,
|
| domain="com.microsoft", name=f"l{L}.norm1", epsilon=eps)
|
| h = down(ln, f"l{L}.norm1_to_fp16")
|
|
|
| lam1 = ly["layer_scale1.lambda1"] * S
|
| lam2 = ly["layer_scale2.lambda1"] * S
|
| wq, bq = c("q.weight", ly["q_proj"], lo), c("q.bias", ly["q_proj.bias"], lo)
|
| wk, wv = c("k.weight", ly["k_proj"], lo), c("v.weight", ly["v_proj"], lo)
|
| wo = c("o.weight", ly["o_proj"] * lam1[None, :], lo)
|
| bo = c("o.bias", (ly["v_proj.bias"].astype(np.float64) @ ly["o_proj"] + ly["o_proj.bias"]) * lam1)
|
| wup, bup = c("up.weight", ly["up_proj"], lo), c("up.bias", ly["up_proj.bias"], lo)
|
| wdn = c("down.weight", ly["down_proj"] * lam2[None, :], lo)
|
| pending_bias = c("down.bias", ly["down_proj.bias"] * lam2)
|
| n2w, n2b = c("norm2.weight", ly["norm2.weight"]), c("norm2.bias", ly["norm2.bias"])
|
|
|
| k = rope(g.op("MatMul", [h, wk], name=f"l{L}.k_proj"), f"l{L}.rope_k")
|
| v = g.op("MatMul", [h, wv], name=f"l{L}.v_proj")
|
| if not last:
|
| q = g.op("Gemm", [g.op("Reshape", [h, sh_2d], name=f"l{L}.q_in_2d"), wq, bq], name=f"l{L}.q_proj")
|
| q = rope(g.op("Reshape", [q, sh_tok3], name=f"l{L}.q_3d"), f"l{L}.rope_q")
|
| skip = resid
|
| else:
|
|
|
|
|
| hc = g.op("Slice", [h, first, one, axis1], name=f"l{L}.cls_rows")
|
| q = g.op("Gemm", [g.op("Reshape", [hc, sh_2d], name=f"l{L}.q_in_2d"), wq, bq], name=f"l{L}.q_proj")
|
| q = g.op("Reshape", [q, sh_cls3], name=f"l{L}.q_3d")
|
| skip = g.op("Slice", [resid, first, one, axis1], name=f"l{L}.resid_cls")
|
| a = g.op("MultiHeadAttention", [q, k, v], domain="com.microsoft", name=f"l{L}.attention", num_heads=HEADS)
|
| o = up(g.op("MatMul", [a, wo], name=f"l{L}.o_proj"), f"l{L}.o_proj_to_fp32")
|
| ln2, _, _, resid = g.op("SkipLayerNormalization", [o, skip, n2w, n2b, bo], n_out=4,
|
| domain="com.microsoft", name=f"l{L}.norm2", epsilon=eps)
|
| m = down(ln2, f"l{L}.norm2_to_fp16")
|
| m = g.op("BiasGelu", [g.op("MatMul", [m, wup], name=f"l{L}.up_proj"), bup], domain="com.microsoft", name=f"l{L}.gelu")
|
| branch = up(g.op("MatMul", [m, wdn], name=f"l{L}.down_proj"), f"l{L}.down_proj_to_fp32")
|
|
|
|
|
| feat, _, _, _ = g.op("SkipLayerNormalization", [branch, resid, g.const("norm.weight", w[P + "norm.weight"]),
|
| g.const("norm.bias", w[P + "norm.bias"]), pending_bias],
|
| n_out=4, domain="com.microsoft", name="backbone_norm", epsilon=eps)
|
| feat = g.op("Reshape", [feat, sh_2d], name="cls_2d")
|
| feat = g.op("LayerNormalization", [feat, g.const("neck.0.weight", w["model.neck.0.weight"]),
|
| g.const("neck.0.bias", w["model.neck.0.bias"])], name="neck_norm", epsilon=EPS, axis=-1)
|
| feat = g.op("MatMul", [feat, g.const("neck.1.weight", np.ascontiguousarray(w["model.neck.1.weight"].T))], name="neck_linear")
|
| feat = g.op("BiasGelu", [feat, g.const("neck.1.bias", w["model.neck.1.bias"])], domain="com.microsoft", name="neck_gelu")
|
| g.nodes.append(helper.make_node("Gemm", [feat, g.const("age_head.weight", w["model.age_head.weight"]),
|
| g.const("age_head.bias", w["model.age_head.bias"])],
|
| ["age_logits"], name="age_head", transB=1))
|
| g.nodes.append(helper.make_node("Gemm", [feat, g.const("gender_head.weight", w["model.gender_head.weight"]),
|
| g.const("gender_head.bias", w["model.gender_head.bias"])],
|
| ["gender_logits"], name="gender_head", transB=1))
|
|
|
| graph = helper.make_graph(
|
| g.nodes, "faceage_dino_vitl16_fused",
|
| [helper.make_tensor_value_info("face", TensorProto.FLOAT, ["batch", 3, 224, 224])],
|
| [helper.make_tensor_value_info("age_logits", TensorProto.FLOAT, ["batch", 100]),
|
| helper.make_tensor_value_info("gender_logits", TensorProto.FLOAT, ["batch", 2])],
|
| g.inits)
|
| model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18), helper.make_opsetid("com.microsoft", 1)],
|
| producer_name="faceage-onnx/scripts/rebuild_fused.py")
|
| model.ir_version = 10
|
| return model
|
|
|
|
|
| def main():
|
| args = [a for a in sys.argv[1:] if not a.startswith("--")]
|
| fp32_only = "--fp32" in sys.argv
|
| out = args[0] if args else DST
|
| model = build(fp32_only)
|
| onnx.checker.check_model(model)
|
| onnx.save(model, out, save_as_external_data=False)
|
| ops = {}
|
| for n in model.graph.node:
|
| ops[n.op_type] = ops.get(n.op_type, 0) + 1
|
| print("nodes", len(model.graph.node), "casts", ops.get("Cast", 0), sorted(ops.items(), key=lambda kv: -kv[1]))
|
| print("saved", out, os.path.getsize(out), sha256(out))
|
|
|
|
|
| if __name__ == "__main__":
|
| main()
|
|
|
