File size: 1,780 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 | from pathlib import Path
import numpy as np
import onnx
from onnx import TensorProto, helper
ROOT = Path(__file__).resolve().parent
FRAMES = 64
def main() -> None:
model = onnx.load(ROOT / "data.decoder.onnx")
model.graph.ClearField("value_info")
model.graph.input[0].CopyFrom(
helper.make_tensor_value_info("z", TensorProto.FLOAT, [1, 192, FRAMES])
)
model.graph.output[0].CopyFrom(
helper.make_tensor_value_info(model.graph.output[0].name, TensorProto.FLOAT, [1, 1, FRAMES * 256])
)
scale = 1
nodes = []
mask_scales = set()
for node in model.graph.node:
if node.op_type == "ConvTranspose":
stride = next(attr.ints[0] for attr in node.attribute if attr.name == "strides")
scale *= stride
if node.op_type in {"Conv", "ConvTranspose"}:
original = node.output[0]
unmasked = f"{original}_unmasked"
node.output[0] = unmasked
nodes.append(node)
mask_name = f"valid_mask_{scale}"
nodes.append(helper.make_node("Mul", [unmasked, mask_name], [original], name=f"{node.name}/ValidMask"))
mask_scales.add(scale)
else:
nodes.append(node)
model.graph.ClearField("node")
model.graph.node.extend(nodes)
for current_scale in sorted(mask_scales):
model.graph.input.append(
helper.make_tensor_value_info(
f"valid_mask_{current_scale}", TensorProto.FLOAT, [1, 1, FRAMES * current_scale]
)
)
model = onnx.shape_inference.infer_shapes(model)
onnx.checker.check_model(model)
onnx.save(model, ROOT / "data.decoder.fixed64.onnx")
print("mask scales", sorted(mask_scales))
if __name__ == "__main__":
main()
|