| """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): |
| |
| n = producer[name] |
| while n.op_type == "Expand": |
| n = producer[n.input[0]] |
| return n |
|
|
| new_nodes, extra = [], [] |
| nz_kind = {} |
| 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)): |
| |
| shp = shapes.get(root.output[0]) |
| |
| 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): |
| |
| 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: |
| |
| 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) |
| |
| nz = producer[idx] |
| while nz.output[0] not in nz_kind: |
| nz = producer[nz.input[0]] |
| |
| 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() |
|
|
