Create scripts/cleanup.py

scripts/cleanup.py

@@ -0,0 +1,157 @@
+"""Cleanup and optimize perch_v2_slim.onnx model.
+
+This script can be applied after completing these steps:
+
+1. Use `tf2onnx` to convert the tflite model to onnx
+2. Apply onnxslim and onnxscript.optimize.optimizer on the model
+3. Manually edit the model to remove the first DFT node (no-op) and fuse
+    the nodes that effectively takes the magnitude of the DFT output with ReduceL2.
+"""
+
+import onnx_ir as ir
+import onnx_ir.passes.common
+import onnxscript
+import numpy as np
+
+m = ir.load("perch_v2_slim.onnx")
+
+for node in m.graph:
+    if node.op_type == "MatMul":
+        print("Simplify MatMul + Reshape:", node.name)
+        if node.inputs[0].producer().op_type == "Reshape":
+            # Skip the reshape
+            input = node.inputs[0].producer().inputs[0]
+            node.replace_input_with(0, input)
+
+        for usage in node.outputs[0].uses():
+            if usage.node.op_type == "Reshape":
+                reshape_usages = list(usage.node.outputs[0].uses())
+                # Keep the last Reshape
+                if reshape_usages[0].node.op_type == "ReduceMax":
+                    shape = ir.val(
+                        "reshape_shape", const_value=ir.tensor([-1, 16, 4, 14795, 4])
+                    )
+                    m.graph.initializers.add(shape)
+                    usage.node.replace_input_with(1, shape)
+                    continue
+                reshape_node = usage.node
+                output = reshape_node.outputs[0]
+                output.replace_all_uses_with(node.outputs[0])
+
+    # Remove Expand
+    if node.op_type == "Expand":
+        print("Remove Expand:", node.name)
+        input = node.inputs[0]
+        output = node.outputs[0]
+        output.replace_all_uses_with(input)
+
+# Clean up any unused nodes
+onnx_ir.passes.common.RemoveUnusedNodesPass()(m)
+# Do some const folding
+onnxscript.optimizer.optimize(
+    m, input_size_limit=1024 * 1024 * 1024, output_size_limit=1024 * 1024 * 1024
+)
+one_1d = ir.val("1d_one", const_value=ir.tensor([1], dtype=ir.DataType.INT64))
+m.graph.initializers.add(one_1d)
+
+# Simplify Unsqueeze + Reshape
+for node in m.graph:
+    if node.op_type == "Reshape":
+        print("Simplify Unsqueeze + Reshape:", node.name)
+        if (
+            node.inputs[0].producer()
+            and node.inputs[0].producer().op_type == "Unsqueeze"
+        ):
+            unsqueeze_node = node.inputs[0].producer()
+            unsqueeze_node.replace_input_with(1, one_1d)
+            node.outputs[0].replace_all_uses_with(unsqueeze_node.outputs[0])
+            unsqueeze_node.outputs[0].shape = ir.Shape(["batch", 160000, 1])
+
+first_reshape_shape = ir.val(
+    "first_reshape_shape", const_value=ir.tensor([-1, 1, 160000, 1])
+)
+m.graph.initializers.add(first_reshape_shape)
+
+# Simplify first Reshape + Unsqueeze
+for node in m.graph:
+    if node.op_type == "Unsqueeze":
+        print("Simplify Reshape + Unsqueeze:", node.name)
+        if node.inputs[0].producer() and node.inputs[0].producer().op_type == "Reshape":
+            reshape_node = node.inputs[0].producer()
+            reshape_node.replace_input_with(1, first_reshape_shape)
+            node.outputs[0].replace_all_uses_with(reshape_node.outputs[0])
+            reshape_node.outputs[0].shape = ir.Shape(["batch", 1, 160000, 1])
+            break
+
+# Fuse Conv + Sub into Conv
+for node in m.graph:
+    if node.op_type == "Conv":
+        print("Check Conv for fusion:", node.name)
+        conv_node = node
+        assert len(conv_node.outputs[0].uses()) == 1
+        for usage in conv_node.outputs[0].uses():
+            if usage.node.op_type == "Sub":
+                sub_node = usage.node
+                print("  Fuse Sub into Conv:", sub_node.name)
+                sub_value = sub_node.inputs[1]
+                new_bias = (np.negative(sub_value.const_value.numpy())).reshape((-1,))
+                new_bias_val = ir.val(
+                    f"{sub_value.name}_neg",
+                    const_value=ir.tensor(new_bias),
+                )
+                m.graph.initializers.add(new_bias_val)
+                if len(conv_node.inputs) == 2:
+                    # Bad access of private field
+                    conv_node._inputs = conv_node._inputs + (None,)
+                conv_node.replace_input_with(2, new_bias_val)
+                sub_node.outputs[0].replace_all_uses_with(conv_node.outputs[0])
+
+# Clean up any unused nodes
+onnx_ir.passes.common.RemoveUnusedNodesPass()(m)
+
+# Clear all intermediate shapes and re-infer shapes
+for node in m.graph:
+    for output in node.outputs:
+        if output.is_graph_output():
+            continue
+        output.shape = None
+
+m.graph.inputs[0].shape = ir.Shape(["batch", *m.graph.inputs[0].shape[1:]])
+for output in m.graph.outputs:
+    output.shape = ir.Shape(["batch", *output.shape[1:]])
+
+onnxscript.optimizer.optimize(
+    m, input_size_limit=1024 * 1024 * 1024, output_size_limit=1024 * 1024 * 1024
+)
+
+onnx_ir.passes.common.ClearMetadataAndDocStringPass()(m)
+
+# Replace None dim with "batch"
+for node in m.graph:
+    for output in node.outputs:
+        if output.shape is None:
+            continue
+        shape = ir.Shape(output.shape)
+        for i in range(len(shape)):
+            dim = shape[i]
+            if isinstance(dim, ir.SymbolicDim) and dim.value is None:
+                shape[i] = ir.SymbolicDim("batch")
+        output.shape = shape
+
+# Rename IO and match the tflite model
+m.graph.inputs[0].name = "inputs"
+m.graph.outputs[0].name = "spatial_embedding"
+m.graph.outputs[1].name = "embedding"
+m.graph.outputs[2].name = "spectrogram"
+m.graph.outputs[3].name = "label"
+
+out_0 = m.graph.outputs[0]
+out_1 = m.graph.outputs[1]
+m.graph.outputs[1] = out_0
+m.graph.outputs[0] = out_1
+
+m.producer_name = "onnx-ir"
+m.producer_version = None
+m.ir_version = 10
+
+ir.save(m, "perch_v2.onnx")