onnxslim 0.1.80__py3-none-any.whl

onnxslim/core/pattern/elimination/concat.py

@@ -0,0 +1,61 @@
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ConcatPatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the ConcatPatternMatcher with a specified priority using a predefined graph pattern."""
+        pattern = Pattern(
+            """
+            input      input       0  1 concat_0
+            Concat     concat_0    1+ 1 input concat_1
+            Concat     concat_1    1* 1 concat_0 output
+            output     output      1  0 concat_1
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the elimination pattern, 'EliminationConcat'."""
+        return "EliminationConcat"
+
+    def rewrite(self, opset=11):
+        """Rewrites an elimination pattern for concat nodes by optimizing nested slice operations."""
+        match_case = {}
+
+        node_concat_0 = self.concat_0
+        users_node_concat_0 = node_concat_0.users
+        node_concat_1 = self.concat_1
+        node_concat_0_axis = node_concat_0.attrs.get("axis", 0)
+        node_concat_1.attrs.get("axis", 0)
+
+        if all(user.op == "Concat" and user.attrs.get("axis", 0) == node_concat_0_axis for user in users_node_concat_0):
+            index = node_concat_1.inputs.index(node_concat_0.outputs[0])
+            node_concat_1.inputs.pop(index)
+            for i, item in enumerate(node_concat_0.inputs):
+                node_concat_1.inputs.insert(index + i, item)
+            inputs = list(node_concat_1.inputs)
+            outputs = list(node_concat_1.outputs)
+            node_concat_1.inputs.clear()
+            node_concat_1.outputs.clear()
+
+            if len(users_node_concat_0) == 0:
+                node_concat_0.inputs.clear()
+                node_concat_0.outputs.clear()
+
+            attrs = {"axis": node_concat_0_axis}
+
+            match_case[node_concat_1.name] = {
+                "op": "Concat",
+                "inputs": inputs,
+                "outputs": outputs,
+                "name": node_concat_1.name,
+                "attrs": attrs,
+                "domain": None,
+            }
+
+        return match_case
+
+
+register_fusion_pattern(ConcatPatternMatcher(1))

onnxslim/core/pattern/elimination/reshape.py

@@ -0,0 +1,77 @@
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ReshapePatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the ReshapePatternMatcher with a priority and a specific pattern for detecting nested reshape
+        operations.
+        """
+        pattern = Pattern(
+            """
+            input    input     0 1 reshape_0
+            Reshape  reshape_0 2 1 input   ? reshape_1
+            Reshape  reshape_1 2 1 reshape_0 ? output
+            output   output    1 0 reshape_1
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name 'EliminationReshape'."""
+        return "EliminationReshape"
+
+    def rewrite(self, opset=11):
+        """Rewrite the computational graph by eliminating redundant reshape operations when certain conditions are
+        met.
+        """
+        match_case = {}
+        node = self.reshape_1
+        first_reshape_node = node.i(0)
+        first_reshape_node_inputs = list(first_reshape_node.inputs)
+        first_reshape_node_users = first_reshape_node.users
+        if len(first_reshape_node_users) == 1:
+            second_reshape_node = node
+
+            def check_constant_mergeable(reshape_node):
+                """Check if a reshape node's shape input, containing zero dimensions, can be merged with its input
+                node's shape.
+                """
+                if isinstance(reshape_node.inputs[1], gs.Constant):
+                    input_shape = reshape_node.inputs[0].shape
+                    reshape_shape = reshape_node.inputs[1].values.tolist()
+                    if input_shape is not None and np.any(np.array(reshape_shape) == 0):
+                        shape = [
+                            input_shape[i] if dim_size == 0 else reshape_shape[i]
+                            for i, dim_size in enumerate(reshape_shape)
+                        ]
+                        if not all(isinstance(item, int) for item in shape):
+                            return False
+                return True
+
+            if check_constant_mergeable(first_reshape_node) and check_constant_mergeable(second_reshape_node):
+                inputs = []
+                inputs.append(first_reshape_node_inputs[0])
+                inputs.append(second_reshape_node.inputs[1])
+                outputs = list(second_reshape_node.outputs)
+                first_reshape_node.outputs.clear()
+                second_reshape_node.inputs.clear()
+                second_reshape_node.outputs.clear()
+
+                match_case[first_reshape_node.name] = {
+                    "op": "Reshape",
+                    "inputs": inputs,
+                    "outputs": outputs,
+                    "name": first_reshape_node.name,
+                    "attrs": first_reshape_node.attrs,
+                    "domain": None,
+                }
+
+        return match_case
+
+
+register_fusion_pattern(ReshapePatternMatcher(1))

onnxslim/core/pattern/elimination/reshape_as.py

@@ -0,0 +1,64 @@
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ReshapeAsPatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the ReshapeAsPatternMatcher with a priority and a specific pattern for reshape as operations."""
+        pattern = Pattern(
+            """
+            input     input     0  1 shape
+            Shape     shape     1+ 1 input gather
+            Gather    gather    1+ 1 shape unsqueeze
+            Unsqueeze unsqueeze 1+ 1 gather output
+            Concat    concat    1+ 1 unsqueeze output
+            output    output    1  0 concat
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name 'EliminationReshapeAs'."""
+        return "EliminationReshapeAs"
+
+    def parameter_check(self) -> bool:
+        shape_node = self.shape
+        if shape_node.outputs[0].shape is None:
+            return False
+
+        if len(shape_node.users) != shape_node.outputs[0].shape[0]:
+            return False
+
+        if not all([user.op == "Gather" for user in shape_node.users]):
+            return False
+
+        for idx, user in enumerate(shape_node.users):
+            if not isinstance(user.inputs[1], gs.Constant):
+                return False
+
+            if user.inputs[1].values.shape != ():
+                return False
+
+            if user.inputs[1].values != idx:
+                return False
+
+        concat_node = self.concat
+        if len(concat_node.inputs) != shape_node.users:
+            return False
+
+        return True
+
+    def rewrite(self, opset=11):
+        """Rewrites the pattern by replacing the Concat node with the Shape node."""
+        match_case = {}
+        shape_node = self.shape
+        concat_node = self.concat
+
+        concat_node.replace_all_uses_with(shape_node)
+
+        return match_case
+
+
+register_fusion_pattern(ReshapeAsPatternMatcher(1))

onnxslim/core/pattern/elimination/slice.py

@@ -0,0 +1,108 @@
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class SlicePatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the SlicePatternMatcher with a specified priority using a predefined graph pattern."""
+        pattern = Pattern(
+            """
+            input  input   0 1 slice_0
+            Slice  slice_0 5 1 input   ? ? ? ? slice_1
+            Slice  slice_1 5 1 slice_0 ? ? ? ? output
+            output output  1 0 slice_1
+            """
+        )  # to check here slice_0
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the elimination pattern, 'EliminationSlice'."""
+        return "EliminationSlice"
+
+    def rewrite(self, opset=11):
+        """Rewrites an elimination pattern for slice nodes by optimizing nested slice operations."""
+        match_case = {}
+        first_slice_node = self.slice_0
+        first_slice_node_inputs = list(first_slice_node.inputs)
+        if all(isinstance(input, gs.Constant) for input in first_slice_node_inputs[1:]):
+            first_slice_node_users = first_slice_node.users
+            if all(
+                user.op == "Slice" and all(isinstance(input, gs.Constant) for input in list(user.inputs)[1:])
+                for user in first_slice_node_users
+            ):
+                first_slice_node_starts = first_slice_node_inputs[1].values.tolist()
+                first_slice_node_ends = first_slice_node_inputs[2].values.tolist()
+                first_slice_node_axes = first_slice_node_inputs[3].values.tolist()
+                first_slice_node_steps = first_slice_node_inputs[4].values.tolist()
+
+                for user_node in first_slice_node_users:
+                    second_slice_node = user_node
+                    second_slice_node_inputs = list(second_slice_node.inputs)
+                    second_slice_node_starts = second_slice_node_inputs[1].values.tolist()
+                    second_slice_node_ends = second_slice_node_inputs[2].values.tolist()
+                    second_slice_node_axes = second_slice_node_inputs[3].values.tolist()
+                    second_slice_node_steps = second_slice_node_inputs[4].values.tolist()
+
+                    new_starts = first_slice_node_starts + second_slice_node_starts
+                    new_ends = first_slice_node_ends + second_slice_node_ends
+                    new_axes = first_slice_node_axes + second_slice_node_axes
+                    new_steps = first_slice_node_steps + second_slice_node_steps
+
+                    if len(new_axes) != len(set(new_axes)):
+                        continue
+
+                    inputs = []
+                    inputs.extend(
+                        (
+                            next(iter(first_slice_node.inputs)),
+                            gs.Constant(
+                                second_slice_node_inputs[1].name + "_starts",
+                                values=np.array(new_starts, dtype=np.int64),
+                            ),
+                            gs.Constant(
+                                second_slice_node_inputs[2].name + "_ends",
+                                values=np.array(new_ends, dtype=np.int64),
+                            ),
+                            gs.Constant(
+                                second_slice_node_inputs[3].name + "_axes",
+                                values=np.array(new_axes, dtype=np.int64),
+                            ),
+                            gs.Constant(
+                                second_slice_node_inputs[4].name + "_steps",
+                                values=np.array(new_steps, dtype=np.int64),
+                            ),
+                        )
+                    )
+                    outputs = list(second_slice_node.outputs)
+
+                    first_slice_node.outputs.clear()
+                    second_slice_node.inputs.clear()
+                    second_slice_node.outputs.clear()
+
+                    if len(first_slice_node_users) == 1:
+                        match_case[first_slice_node.name] = {
+                            "op": "Slice",
+                            "inputs": inputs,
+                            "outputs": outputs,
+                            "name": first_slice_node.name,
+                            "attrs": first_slice_node.attrs,
+                            "domain": None,
+                        }
+                    else:
+                        match_case[second_slice_node.name] = {
+                            "op": "Slice",
+                            "inputs": inputs,
+                            "outputs": outputs,
+                            "name": second_slice_node.name,
+                            "attrs": second_slice_node.attrs,
+                            "domain": None,
+                        }
+
+        return match_case
+
+
+register_fusion_pattern(SlicePatternMatcher(1))

onnxslim/core/pattern/elimination/unsqueeze.py

@@ -0,0 +1,92 @@
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class UnsqueezePatternMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the UnsqueezePatternMatcher with a specified priority using a predefined graph pattern."""
+        pattern = Pattern(
+            """
+            input      input       0  1 unsqueeze_0
+            Unsqueeze  unsqueeze_0 1+ 1 input unsqueeze_1
+            Unsqueeze  unsqueeze_1 1+ 1 unsqueeze_0 output
+            output     output      1  0 unsqueeze_1
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the elimination pattern, 'EliminationUnsqueeze'."""
+        return "EliminationUnsqueeze"
+
+    def rewrite(self, opset=11):
+        """Rewrites an elimination pattern for unsqueeze nodes by optimizing nested slice operations."""
+        match_case = {}
+        node_unsqueeze_0 = self.unsqueeze_0
+        users_node_unsqueeze_0 = node_unsqueeze_0.users
+        node_unsqueeze_1 = self.unsqueeze_1
+        if len(users_node_unsqueeze_0) == 1 and node_unsqueeze_0.inputs[0].shape and node_unsqueeze_1.inputs[0].shape:
+            if opset < 13 or (
+                isinstance(node_unsqueeze_0.inputs[1], gs.Constant)
+                and isinstance(node_unsqueeze_1.inputs[1], gs.Constant)
+            ):
+
+                def get_unsqueeze_axes(unsqueeze_node, opset):
+                    dim = len(unsqueeze_node.inputs[0].shape)
+                    if opset < 13:
+                        axes = unsqueeze_node.attrs["axes"]
+                    else:
+                        axes = unsqueeze_node.inputs[1].values
+                    return [axis + dim + len(axes) if axis < 0 else axis for axis in axes]
+
+                axes_node_unsqueeze_0 = get_unsqueeze_axes(node_unsqueeze_0, opset)
+                axes_node_unsqueeze_1 = get_unsqueeze_axes(node_unsqueeze_1, opset)
+
+                axes_node_unsqueeze_0 = [
+                    axis + sum(1 for axis_ in axes_node_unsqueeze_1 if axis_ <= axis) for axis in axes_node_unsqueeze_0
+                ]
+
+                inputs = [next(iter(node_unsqueeze_0.inputs))]
+                outputs = list(node_unsqueeze_1.outputs)
+
+                index = node_unsqueeze_1.inputs.index(node_unsqueeze_0.outputs[0])
+                node_unsqueeze_1.inputs.pop(index)
+                for i, item in enumerate(node_unsqueeze_0.inputs):
+                    node_unsqueeze_1.inputs.insert(index + i, item)
+                inputs = [next(iter(node_unsqueeze_1.inputs))]
+                outputs = list(node_unsqueeze_1.outputs)
+                node_unsqueeze_1.inputs.clear()
+                node_unsqueeze_1.outputs.clear()
+
+                if len(users_node_unsqueeze_0) == 0:
+                    outputs.inputs.clear()
+                    outputs.outputs.clear()
+
+                if opset < 13:
+                    attrs = {"axes": axes_node_unsqueeze_0 + axes_node_unsqueeze_1}
+                else:
+                    attrs = None
+                    inputs.append(
+                        gs.Constant(
+                            name=f"{node_unsqueeze_0.name}_axes",
+                            values=np.array(axes_node_unsqueeze_0 + axes_node_unsqueeze_1, dtype=np.int64),
+                        )
+                    )
+
+                match_case[node_unsqueeze_0.name] = {
+                    "op": "Unsqueeze",
+                    "inputs": inputs,
+                    "outputs": outputs,
+                    "name": node_unsqueeze_0.name,
+                    "attrs": attrs,
+                    "domain": None,
+                }
+
+        return match_case
+
+
+register_fusion_pattern(UnsqueezePatternMatcher(1))

onnxslim/core/pattern/fusion/__init__.py

@@ -0,0 +1,8 @@
+from .concat_reshape import *
+from .convadd import *
+from .convbn import *
+from .convmul import *
+from .gelu import *
+from .gemm import *
+from .padconv import *
+from .reduce import *

onnxslim/core/pattern/fusion/concat_reshape.py

@@ -0,0 +1,50 @@
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ConcatReshapeMatcher(PatternMatcher):
+    def __init__(self, priority):
+        pattern = Pattern(
+            """
+            input    input  0   1 concat_0
+            Concat   concat_0   1+ 1 input reshape_0
+            Reshape  reshape_0  2  1 ? concat_0 output
+            output   output     1  0 reshape_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        return "FusionConcatReshape"
+
+    def parameter_check(self):
+        concat_node = self.concat_0
+
+        def check_inputs(inputs):
+            vars = [i for i in inputs if isinstance(i, gs.Variable)]
+            consts = [i for i in inputs if isinstance(i, gs.Constant)]
+            return (
+                len(vars) == 1 and all(c.values.size == 1 and c.values != -1 for c in consts) and vars[0].shape == [1]
+            )
+
+        return check_inputs(concat_node.inputs)
+
+    def rewrite(self, opset=11):
+        match_case = {}
+        concat_node = self.concat_0
+        index = next(idx for idx, i in enumerate(concat_node.inputs) if isinstance(i, gs.Variable))
+        constant = gs.Constant(
+            concat_node.inputs[index].name + "_fixed",
+            values=np.array([-1], dtype=np.int64),
+        )
+        concat_node.inputs.pop(index)
+        concat_node.inputs.insert(index, constant)
+
+        return match_case
+
+
+register_fusion_pattern(ConcatReshapeMatcher(1))

onnxslim/core/pattern/fusion/convadd.py

@@ -0,0 +1,70 @@
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ConvAddMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the ConvAddMatcher for fusing Conv and Add layers in an ONNX graph."""
+        pattern = Pattern(
+            """
+            input    input  0  1 conv_0
+            Conv     conv_0 1+ 1 input bn_0
+            Add      add_0  2  1 conv_0 ? output
+            output   output 1  0 add_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the FusionConvAdd pattern."""
+        return "FusionConvAdd"
+
+    def rewrite(self, opset=11):
+        match_case = {}
+        conv_node = self.conv_0
+        conv_weight = list(conv_node.inputs)[1]
+        conv_node_users = conv_node.users
+        node = self.add_0
+        if (
+            len(conv_node_users) == 1
+            and isinstance(node.inputs[1], gs.Constant)
+            and isinstance(conv_weight, gs.Constant)
+            and node.inputs[1].values.squeeze().ndim == 1
+            and node.inputs[1].values.squeeze().shape[0] == conv_weight.shape[0]
+        ):
+            add_node = node
+            if len(conv_node.inputs) == 2:
+                conv_bias = node.inputs[1].values.squeeze()
+            else:
+                conv_bias = conv_node.inputs[2].values + node.inputs[1].values.squeeze()
+
+            inputs = []
+            inputs.append(next(iter(conv_node.inputs)))
+            inputs.append(conv_weight)
+            weight_name = list(conv_node.inputs)[1].name
+            if weight_name.endswith("weight"):
+                bias_name = f"{weight_name[:-6]}bias"
+            else:
+                bias_name = f"{weight_name}_bias"
+            inputs.append(gs.Constant(bias_name, values=conv_bias))
+            outputs = list(add_node.outputs)
+
+            conv_node.outputs.clear()
+            add_node.inputs.clear()
+            add_node.outputs.clear()
+
+            match_case[conv_node.name] = {
+                "op": conv_node.op,
+                "inputs": inputs,
+                "outputs": outputs,
+                "name": conv_node.name,
+                "attrs": conv_node.attrs,
+                "domain": None,
+            }
+
+        return match_case
+
+
+register_fusion_pattern(ConvAddMatcher(1))

onnxslim/core/pattern/fusion/convbn.py

@@ -0,0 +1,86 @@
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern import Pattern, PatternMatcher
+from onnxslim.core.pattern.registry import register_fusion_pattern
+
+
+class ConvBatchNormMatcher(PatternMatcher):
+    def __init__(self, priority):
+        """Initializes the ConvBatchNormMatcher for fusing Conv and BatchNormalization layers in an ONNX graph."""
+        pattern = Pattern(
+            """
+            input              input  0  1 conv_0
+            Conv               conv_0 1+ 1 input bn_0
+            BatchNormalization bn_0   5  1 conv_0 ? ? ? ? output
+            output             output 1  0 bn_0
+            """
+        )
+        super().__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the FusionConvBN pattern."""
+        return "FusionConvBN"
+
+    def rewrite(self, opset=11):
+        """Rewrites the weights and biases of a BatchNormalization layer fused with a convolution layer."""
+        match_case = {}
+        conv_transpose_node = self.conv_0
+        conv_transpose_node_users = conv_transpose_node.users
+        node = self.bn_0
+        if len(conv_transpose_node_users) == 1 and isinstance(conv_transpose_node.inputs[1], gs.Constant):
+            conv_transpose_weight = conv_transpose_node.inputs[1].values
+            bn_node = node
+            bn_scale = bn_node.inputs[1].values
+            bn_bias = bn_node.inputs[2].values
+            bn_running_mean = bn_node.inputs[3].values
+            bn_running_var = bn_node.inputs[4].values
+            bn_eps = bn_node.attrs.get("epsilon", 1.0e-5)
+
+            if len(conv_transpose_node.inputs) == 2:
+                conv_transpose_bias = np.zeros_like(bn_running_mean)
+            else:
+                conv_transpose_bias = conv_transpose_node.inputs[2].values
+
+            bn_var_rsqrt = bn_scale / np.sqrt(bn_running_var + bn_eps)
+            shape = [1] * len(conv_transpose_weight.shape)
+            if bn_node.i(0).op == "Conv":
+                shape[0] = -1
+            else:
+                shape[1] = -1
+            conv_w = conv_transpose_weight * bn_var_rsqrt.reshape(shape)
+            conv_b = (conv_transpose_bias - bn_running_mean) * bn_var_rsqrt + bn_bias
+
+            inputs = []
+            inputs.append(next(iter(conv_transpose_node.inputs)))
+            weight_name = list(conv_transpose_node.inputs)[1].name
+            if weight_name.endswith("weight"):
+                bias_name = f"{weight_name[:-6]}bias"
+            else:
+                bias_name = f"{weight_name}_bias"
+            inputs.extend(
+                (
+                    gs.Constant(weight_name + "_weight", values=conv_w),
+                    gs.Constant(bias_name, values=conv_b),
+                )
+            )
+            outputs = list(bn_node.outputs)
+
+            conv_transpose_node.outputs.clear()
+            bn_node.inputs.clear()
+            bn_node.outputs.clear()
+
+            match_case[conv_transpose_node.name] = {
+                "op": conv_transpose_node.op,
+                "inputs": inputs,
+                "outputs": outputs,
+                "name": conv_transpose_node.name,
+                "attrs": conv_transpose_node.attrs,
+                "domain": None,
+            }
+
+        return match_case
+
+
+register_fusion_pattern(ConvBatchNormMatcher(1))