File size: 5,918 Bytes
67faadf | 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 | """Remove data-dependent shapes from the duration predictor's spline flows.
Each spline flow (dp/flows.3/5/7) uses NonZero(mask) to pick the elements inside
the spline range (+-5), runs the spline on just those, then scatters results back.
That makes tensor sizes depend on the data.
Dense rewrite, same output:
- every NonZero is replaced by a constant listing all positions, so gathers and
scatters always see all N elements
- the spline input is clamped to the spline range so out-of-range elements can't
index past the bins (their results are discarded below)
- the final "inside" scatter becomes Where(inside_mask, spline_result, previous),
so outside elements keep the identity value written by the "outside" scatter
"""
import sys
from pathlib import Path
import numpy as np
import onnx
from onnx import TensorProto, helper, numpy_helper
ROOT = Path(__file__).resolve().parent
SRC = sys.argv[1] if len(sys.argv) > 1 else "data.encoder.fixed64x160.onnx"
DST = sys.argv[2] if len(sys.argv) > 2 else SRC.replace(".onnx", ".static.onnx")
def main() -> None:
m = onnx.load(ROOT / SRC)
g = m.graph
m = onnx.shape_inference.infer_shapes(m)
g = m.graph
shapes = {}
for vi in list(g.value_info) + list(g.input) + list(g.output):
shapes[vi.name] = [d.dim_value if d.HasField("dim_value") else None for d in vi.type.tensor_type.shape.dim]
producer = {o: n for n in g.node for o in n.output}
consumers = {}
for n in g.node:
for x in n.input:
consumers.setdefault(x, []).append(n)
def mask_root(name):
# Follow Expand back to the And / Not node that made the mask.
n = producer[name]
while n.op_type == "Expand":
n = producer[n.input[0]]
return n
new_nodes, extra = [], []
nz_kind = {} # NonZero output -> "in" / "out"
for n in g.node:
if n.op_type != "NonZero":
continue
root = mask_root(n.input[0])
shp = shapes.get(n.input[0])
if not (shp and all(shp)):
# Expanded masks: batch and channel are 1 here, same as the source mask.
shp = shapes.get(root.output[0])
# dp flows act on one channel of a [1, 2, N] tensor: masks are [1, 1, N].
shp = [d if d else 1 for d in (shp or [None, None, None])]
assert len(shp) == 3 and shp[-1] > 1, f"unexpected mask shape for {n.name}: {shp}"
nz_kind[n.output[0]] = "in" if root.op_type == "And" else "out"
idx = np.stack(np.meshgrid(*[np.arange(s) for s in shp], indexing="ij")).reshape(len(shp), -1).astype(np.int64)
n.op_type = "Constant"
del n.input[:]
del n.attribute[:]
n.attribute.append(helper.make_attribute("value", numpy_helper.from_array(idx, n.output[0] + "_all")))
def idx_kind(idx_name):
# GatherND/ScatterND indices = Transpose(NonZero) (possibly via Cast/Reshape).
n = producer.get(idx_name)
while n is not None and n.output[0] not in nz_kind:
n = producer.get(n.input[0]) if n.input else None
return nz_kind.get(n.output[0]) if n is not None else None
clamped = where_fixed = 0
nodes = list(g.node)
out_nodes = []
for n in nodes:
if n.op_type == "GatherND" and idx_kind(n.input[1]) == "in" and len(shapes.get(n.input[0], [])) == 3:
# Spline input x (shape [1,1,N]): clamp into the spline range.
lo = n.name + "/lo"; hi = n.name + "/hi"; cl = n.input[0] + "_clamped_" + str(clamped)
out_nodes.append(helper.make_node("Constant", [], [lo], value=numpy_helper.from_array(np.array(-5.0, np.float32))))
out_nodes.append(helper.make_node("Constant", [], [hi], value=numpy_helper.from_array(np.array(5.0, np.float32))))
out_nodes.append(helper.make_node("Clip", [n.input[0], lo, hi], [cl], name=n.name + "/Clamp"))
n.input[0] = cl
clamped += 1
out_nodes.append(n)
elif n.op_type == "ScatterND" and idx_kind(n.input[1]) == "in":
data, idx, upd = n.input
shp = shapes.get(data)
# The mask fed to the (now constant) NonZero: find it via the original Expand node.
nz = producer[idx]
while nz.output[0] not in nz_kind:
nz = producer[nz.input[0]]
# Rebuild the boolean mask: inside = And(...) expanded to data's shape.
and_node = None
for cand in g.node:
if cand.op_type == "And" and cand.name.rsplit("/", 1)[0] == n.name.rsplit("/", 1)[0]:
and_node = cand
assert and_node is not None, n.name
if not (shp and all(shp)):
shp = shapes[and_node.output[0]]
shp = [d if d else 1 for d in shp]
shp_name = n.name + "/shape"
out_nodes.append(helper.make_node("Constant", [], [shp_name], value=numpy_helper.from_array(np.array(shp, np.int64))))
mask = n.name + "/inside_mask"
out_nodes.append(helper.make_node("Expand", [and_node.output[0], shp_name], [mask], name=n.name + "/MaskExpand"))
upd_r = n.name + "/upd_dense"
out_nodes.append(helper.make_node("Reshape", [upd, shp_name], [upd_r], name=n.name + "/UpdReshape"))
out_nodes.append(helper.make_node("Where", [mask, upd_r, data], list(n.output), name=n.name + "/Where"))
where_fixed += 1
else:
out_nodes.append(n)
g.ClearField("node")
g.node.extend(out_nodes)
g.ClearField("value_info")
m = onnx.shape_inference.infer_shapes(m)
onnx.checker.check_model(m)
onnx.save(m, ROOT / DST)
left = sum(1 for n in m.graph.node if n.op_type == "NonZero")
print(f"NonZero replaced: {len(nz_kind)}; clamps: {clamped}; inside scatters -> Where: {where_fixed}; NonZero left: {left}")
print("saved", DST)
if __name__ == "__main__":
main()
|