Create optimized models

BirdNET_GLOBAL_6K_V2.4_Model_FP32.tflite

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:55f3e4055b1a13bfa9a2452731d0d34f6a02d6b775a334362665892794165e4c
+size 51726412

birdnet.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:920e2bd05aa6265e68e21711bc3bab5900cf7c4117f7d958723632ef75956295
+size 66935346

scripts/optimize.py

@@ -43,12 +43,16 @@ class ReplaceDftWithMatMulRule(onnxscript.rewriter.RewriteRuleClassBase):
 
 class ReplaceSplit(onnxscript.rewriter.RewriteRuleClassBase):
     def pattern(self, op, x):
-        return op.Split(x, _allow_other_inputs=True, _outputs=["split_out_1", "split_out_2"])
-    
+        return op.Split(
+            x, _allow_other_inputs=True, _outputs=["split_out_1", "split_out_2"]
+        )
+
     def rewrite(self, op, x: ir.Value, **kwargs):
         zero = op.initializer(ir.tensor(np.array([0], dtype=np.int64)), "zero")
         batch_size = op.Gather(x, zero)
-        sample_size = op.initializer(ir.tensor(np.array([144000], dtype=np.int32)), "sample_size")
+        sample_size = op.initializer(
+            ir.tensor(np.array([144000], dtype=np.int32)), "sample_size"
+        )
         return batch_size, sample_size
 
 
@@ -59,13 +63,17 @@ class RemoveCast(onnxscript.rewriter.RewriteRuleClassBase):
     def rewrite(self, op, x: ir.Value, **kwargs):
         return op.Identity(x)
 
+
 model = ir.load("model.onnx")
 
 # Set dynamic axes
 model.graph.inputs[0].shape = ir.Shape(["batch", 144000])
 model.graph.outputs[0].shape = ir.Shape(["batch", 6522])
 
-onnxscript.rewriter.rewrite(model, [ReplaceDftWithMatMulRule().rule(), ReplaceSplit().rule(), RemoveCast().rule()])
+onnxscript.rewriter.rewrite(
+    model,
+    [ReplaceDftWithMatMulRule().rule(), ReplaceSplit().rule(), RemoveCast().rule()],
+)
 
 # Change all int32 initializers to int64
 initializers = list(model.graph.initializers.values())
@@ -82,22 +90,39 @@ onnxscript.optimizer.optimize(
     model, input_size_limit=1024 * 1024 * 1024, output_size_limit=1024 * 1024 * 1024
 )
 
+
 # Remove Slice-Reshape
 def remove_slice_reshape(model: ir.Model):
     mul_node = model.graph.node("model/MEL_SPEC1/Mul")
     first_reshape = model.graph.node("model/MEL_SPEC1/stft/frame/Reshape_1")
-    first_shape = ir.val("first_shape", const_value=ir.tensor([-1, 72000, 2], dtype=ir.DataType.INT64))
+    first_shape = ir.val(
+        "first_shape", const_value=ir.tensor([-1, 72000, 2], dtype=ir.DataType.INT64)
+    )
     model.graph.initializers.add(first_shape)
     second_reshape = model.graph.node("model/MEL_SPEC2/stft/frame/Reshape_1")
-    second_shape = ir.val("second_shape", const_value=ir.tensor([-1, 18000, 8], dtype=ir.DataType.INT64))
+    second_shape = ir.val(
+        "second_shape", const_value=ir.tensor([-1, 18000, 8], dtype=ir.DataType.INT64)
+    )
     model.graph.initializers.add(second_shape)
 
+    third_reshape = model.graph.node("model/MEL_SPEC1/stft/frame/Reshape_4")
+    third_shape = ir.val(
+        "third_shape", const_value=ir.tensor([-1, 511, 2048], dtype=ir.DataType.INT64)
+    )
+    model.graph.initializers.add(third_shape)
+    fourth_reshape = model.graph.node("model/MEL_SPEC2/stft/frame/Reshape_4")
+    fourth_shape = ir.val(
+        "fourth_shape", const_value=ir.tensor([-1, 511, 1024], dtype=ir.DataType.INT64)
+    )
+    model.graph.initializers.add(fourth_shape)
+
     # Replace with Mul-Reshape-Gather
     first_reshape.replace_input_with(0, mul_node.outputs[0])
     first_reshape.replace_input_with(1, first_shape)
-    second_reshape.replace_input_with(0, first_reshape.outputs[0])
+    second_reshape.replace_input_with(0, mul_node.outputs[0])
     second_reshape.replace_input_with(1, second_shape)
-
+    third_reshape.replace_input_with(1, third_shape)
+    fourth_reshape.replace_input_with(1, fourth_shape)
 
 
 remove_slice_reshape(model)
@@ -108,6 +133,8 @@ onnxscript.optimizer.optimize(
 
 
 onnx_ir.passes.common.ClearMetadataAndDocStringPass()(model)
+model.graph.inputs[0].name = "input"
+model.graph.outputs[0].name = "output"
 model.ir_version = 10
 model.producer_name = "onnx-ir"
 model.graph.name = "BirdNET-v2.4"