"""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 # scale of the residual stream EPS = 1e-5 # the source's LayerNorm epsilon HEADS = 16 DIM = 1024 TOKENS = 201 # CLS + 4 registers + 14x14 patches 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) # The MatMuls with a weight, in graph order: q, k, v, o, up, down for each block. 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]] # Check the role against the bias the source adds to it. 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): # fp32 -> matmul precision return x if fp32_only else g.op("Cast", [x], name=name, to=TensorProto.FLOAT16) def up(x, name): # matmul precision -> fp32 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)) # Rotary tables: cos/sin of the source's angle table, one half (the two halves are equal). 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) # Patch embedding (fp32), scaled by S; CLS, registers and the conv bias go in one table # that the first SkipLayerNormalization adds. 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]) # B,196,1024 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") # B,201,1024, zeros in the prefix rows # Everything below is B x tokens x 1024, and B x 1 x 1024 after the last block's attention. # A projection's bias is added in fp32 by the SkipLayerNormalization after it, except # q's, which comes before the rotary embedding. v's bias is folded into o's: the # attention weights of a row sum to one, so it passes through attention unchanged. eps = EPS * S * S resid = None pending_bias = None # the last down projection's bias, for the next SkipLayerNormalization 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: # Residual = prefix-padded patches + token table, and its LayerNorm, in one op. 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: # Only the CLS token is read after this block: its query alone, and its row of # the residual. The CLS token is not rotated. 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") # Final backbone norm (residual add fused), then the neck and the heads, all 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()