File size: 8,847 Bytes
c07c0dc | 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 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """
Produce a compact fp16 variant of the exported ONNX (and best-effort onnxslim).
The 804 MB fp32 export carries ~425 MB of constant-folded fp32 tensors (Swin attention masks /
positional-embedding tables). fp16 (keep_io_types=True so I/O stay float32) ~halves the file and is
native/fast on WebGPU & DirectML. CPU runs fp16 too (via casts), just slower than fp32.
Note: onnxslim's shape-inference pass trips protobuf's 2 GB serialization limit on the fp32 model,
so we convert to fp16 first (smaller) and only then try to slim.
"""
import time
from pathlib import Path
import onnx
ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "models" / "mvanet_box_segmenter.onnx"
FP16 = ROOT / "models" / "mvanet_box_segmenter_fp16.onnx"
FP16_SLIM = ROOT / "models" / "mvanet_box_segmenter_fp16_slim.onnx"
def mb(p: Path) -> float:
return p.stat().st_size / 1e6
def main() -> None:
print(f"[src ] {SRC.name} {mb(SRC):.1f} MB")
# fp16 (input/output kept float32 so the JS harness is unchanged)
import numpy as np
from onnx import TensorProto, numpy_helper
from onnxconverter_common import float16
def fix_resize_scale_inputs(model) -> int:
"""Resize/Upsample roi & scales inputs must be float32 even when X is fp16 (ONNX spec).
The fp16 converter wrongly casts the Constants feeding them; cast those back to fp32."""
inits = {t.name: t for t in model.graph.initializer}
consts = {n.output[0]: n for n in model.graph.node if n.op_type == "Constant"}
targets = set()
for n in model.graph.node:
if n.op_type in ("Resize", "Upsample"):
for idx in (1, 2): # roi, scales
if idx < len(n.input) and n.input[idx]:
targets.add(n.input[idx])
fixed = 0
for name in targets:
if name in inits and inits[name].data_type == TensorProto.FLOAT16:
arr = numpy_helper.to_array(inits[name]).astype(np.float32)
inits[name].CopyFrom(numpy_helper.from_array(arr, name)); fixed += 1
elif name in consts:
for a in consts[name].attribute:
if a.name == "value" and a.t.data_type == TensorProto.FLOAT16:
arr = numpy_helper.to_array(a.t).astype(np.float32)
a.t.CopyFrom(numpy_helper.from_array(arr, a.t.name)); fixed += 1
return fixed
from onnx import helper
def wrap_fp32_io(model) -> None:
"""Make external input/output float32 around an all-fp16 interior, via explicit Cast nodes.
(keep_io_types=True left the fp32 input feeding dtype-agnostic reshape branches, which then
clashed with cast fp16 branches at a Concat — so we go uniform fp16 + manual I/O casts.)"""
g = model.graph
# INPUT: rename the fp16 graph input value, add a fp32 input that Casts into it.
gi = g.input[0]
old = gi.name
internal = old + "_f16"
for n in g.node:
n.input[:] = [internal if x == old else x for x in n.input]
shape = [d.dim_value for d in gi.type.tensor_type.shape.dim]
g.input.remove(gi)
g.input.insert(0, helper.make_tensor_value_info(old, TensorProto.FLOAT, shape))
g.node.insert(0, helper.make_node("Cast", [old], [internal], to=TensorProto.FLOAT16,
name="cast_input_to_fp16"))
# OUTPUT: rename the fp16 graph output, Cast it back to fp32.
go = g.output[0]
oold = go.name
ointernal = oold + "_f16"
for n in g.node:
n.output[:] = [ointernal if x == oold else x for x in n.output]
n.input[:] = [ointernal if x == oold else x for x in n.input]
oshape = [d.dim_value for d in go.type.tensor_type.shape.dim]
g.output.remove(go)
g.output.insert(0, helper.make_tensor_value_info(oold, TensorProto.FLOAT, oshape))
g.node.append(helper.make_node("Cast", [ointernal], [oold], to=TensorProto.FLOAT,
name="cast_output_to_fp32"))
# Shape inference must be ON: otherwise the converter can't see branch types (e.g. the global-view
# Resize branch) and leaves a mixed fp16/fp32 Concat in SplitMultiView. keep_io_types=True keeps
# the model's input/output float32 so the Node harness is unchanged.
# onnxconverter_common leaves several mixed fp16/fp32 boundaries (Concat in SplitMultiView, the
# WindowSDPA scale Div, ...) that onnx.checker accepts but ORT rejects. General fix: infer types,
# then Cast any fp32 input of a float-compute op to fp16 so every such op is uniformly fp16.
FLOAT_OPS = {"Add", "Sub", "Mul", "Div", "Concat", "Where", "Sum", "Min", "Max", "Pow",
"MatMul", "Gemm", "PRelu", "Sqrt", "Reciprocal", "Neg", "Softmax", "Erf", "Sigmoid"}
def reconcile_float_types(model) -> int:
"""Forward-propagate per-tensor dtypes ourselves (onnx shape-inference under-types this graph),
then Cast any fp32 input of a float-compute op to fp16 so ORT's strict type check passes."""
from onnx import TensorProto, helper
FP16, FP32, INT64, BOOL = (TensorProto.FLOAT16, TensorProto.FLOAT,
TensorProto.INT64, TensorProto.BOOL)
INT_OUT = {"Shape", "Size", "NonZero", "ArgMax", "ArgMin"}
BOOL_OUT = {"Equal", "Greater", "Less", "GreaterOrEqual", "LessOrEqual",
"And", "Or", "Not", "Xor", "IsNaN", "IsInf"}
dt = {t.name: t.data_type for t in model.graph.initializer}
for v in model.graph.input:
dt[v.name] = v.type.tensor_type.elem_type
def out_dtype(n):
op = n.op_type
if op == "Constant":
for a in n.attribute:
if a.name == "value":
return a.t.data_type
if a.name in ("value_float", "value_floats"):
return FP32
if a.name in ("value_int", "value_ints"):
return INT64
return None
if op == "ConstantOfShape":
for a in n.attribute:
if a.name == "value":
return a.t.data_type
return FP32
if op in ("Cast", "CastLike"):
return next((a.i for a in n.attribute if a.name == "to"), None)
if op in INT_OUT:
return INT64
if op in BOOL_OUT:
return BOOL
return dt.get(n.input[0]) if n.input else None
for _ in range(3): # a few passes in case nodes aren't perfectly topo-sorted
for n in model.graph.node:
d = out_dtype(n)
if d is not None:
for o in n.output:
if o:
dt[o] = d
out_nodes, c = [], 0
for n in model.graph.node:
if n.op_type in FLOAT_OPS:
for idx, i in enumerate(list(n.input)):
if i and dt.get(i) == FP32:
co = f"{i}__f16_{c}"
out_nodes.append(helper.make_node("Cast", [i], [co], to=FP16, name=f"recon_{c}"))
n.input[idx] = co
dt[co] = FP16
c += 1
out_nodes.append(n)
model.graph.ClearField("node")
model.graph.node.extend(out_nodes)
return c
# Pipeline order matters:
# convert fp16 -> fix Resize scales -> onnxslim (folds 3377 Constants, dedupes)
# -> reconcile float types -> clear stale value_info -> save.
# reconcile + value_info-clear MUST be the LAST steps: onnxslim eliminates "redundant" casts
# (undoing reconcile) and the original fp32 value_info conflicts with the fp16 tensors (ORT rejects).
tmp1 = FP16.with_name("_fp16_pre.onnx")
m = onnx.load(str(SRC))
t = time.time()
m16 = float16.convert_float_to_float16(m, keep_io_types=True, disable_shape_infer=False)
fix_resize_scale_inputs(m16)
onnx.save(m16, str(tmp1))
from onnxslim import slim
slim(str(tmp1), str(FP16)) # slim in place into the final path
ms = onnx.load(str(FP16))
n_fixed = fix_resize_scale_inputs(ms) # re-assert in case slim refolded
n_recon = reconcile_float_types(ms)
ms.graph.ClearField("value_info") # drop stale fp32 annotations; ORT re-infers
onnx.checker.check_model(ms)
onnx.save(ms, str(FP16))
tmp1.unlink(missing_ok=True)
print(f"[fp16] slim+reconcile: resize-fixes={n_fixed} reconcile-casts={n_recon}")
print(f"[fp16] {FP16.name} {mb(FP16):.1f} MB ({time.time()-t:.1f}s)")
if __name__ == "__main__":
main()
|