onnxslim 0.1.80__py3-none-any.whl

onnxslim/core/optimization/__init__.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import logging
+from collections import Counter
+from typing import List, Optional, Union
+
+import onnx
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.pattern.registry import get_fusion_patterns
+from onnxslim.third_party.onnx_graphsurgeon.ir.graph import Graph
+
+logger = logging.getLogger("onnxslim")
+
+from .dead_node_elimination import dead_node_elimination
+from .subexpression_elimination import subexpression_elimination
+from .weight_tying import tie_weights
+
+
+class OptimizationSettings:
+    constant_folding = True
+    graph_fusion = True
+    dead_node_elimination = True
+    subexpression_elimination = True
+    weight_tying = True
+
+    @classmethod
+    def keys(cls):
+        return [
+            "constant_folding",
+            "graph_fusion",
+            "dead_node_elimination",
+            "subexpression_elimination",
+            "weight_tying",
+        ]
+
+    @classmethod
+    def reset(cls, skip_optimizations: list[str] | None = None):
+        for key in cls.keys():
+            if skip_optimizations and key in skip_optimizations:
+                setattr(cls, key, False)
+            else:
+                setattr(cls, key, True)
+
+    @classmethod
+    def stats(cls):
+        return {key: getattr(cls, key) for key in cls.keys()}
+
+    @classmethod
+    def enabled(cls):
+        return any([getattr(cls, key) for key in cls.keys()])
+
+
+def optimize_model(model: onnx.ModelProto | gs.Graph, skip_fusion_patterns: str | None = None) -> onnx.ModelProto:
+    """Optimize and transform the given ONNX model using various fusion patterns and graph rewriting techniques."""
+    graph = model if isinstance(model, gs.Graph) else gs.import_onnx(model)
+    if OptimizationSettings.graph_fusion:
+        logger.debug("Start graph_fusion.")
+        fusion_patterns = get_fusion_patterns(skip_fusion_patterns)
+        graph_fusion(graph, fusion_patterns)
+        logger.debug("Finish graph_fusion.")
+    if OptimizationSettings.dead_node_elimination:
+        logger.debug("Start dead_node_elimination.")
+        dead_node_elimination(graph)
+        graph.cleanup(remove_unused_graph_inputs=True).toposort()
+        logger.debug("Finish dead_node_elimination.")
+    if OptimizationSettings.subexpression_elimination:
+        logger.debug("Start subexpression_elimination.")
+        subexpression_elimination(graph)
+        graph.cleanup(remove_unused_graph_inputs=True).toposort()
+        logger.debug("Finish subexpression_elimination.")
+    if OptimizationSettings.weight_tying:
+        logger.debug("Start weight_tying.")
+        tie_weights(graph)
+        logger.debug("Finish weight_tying.")
+    model = gs.export_onnx(graph)
+
+    return model
+
+
+@gs.Graph.register()
+def replace_custom_layer(
+    self,
+    op: str,
+    inputs,
+    outputs: list[str],
+    name: str,
+    attrs: dict | None = None,
+    domain: str = "ai.onnx.contrib",
+):
+    """Replace a custom layer in the computational graph with specified parameters and domain."""
+    return self.layer(
+        op=op,
+        inputs=inputs,
+        outputs=outputs,
+        name=name,
+        attrs=attrs,
+        domain=domain,
+    )
+
+
+def graph_fusion(graph: Graph, fusion_patterns: dict, is_subgraph=False):
+    for subgraph in graph.subgraphs():
+        graph_fusion(subgraph, fusion_patterns, is_subgraph=True)
+
+    fusion_pairs = find_matches(graph, fusion_patterns)
+    for match in fusion_pairs.values():
+        graph.replace_custom_layer(**match)
+
+    graph.cleanup(remove_unused_graph_inputs=True if not is_subgraph else False).toposort()
+
+
+def find_matches(graph: Graph, fusion_patterns: dict):
+    """Find matching patterns in the graph based on provided fusion patterns."""
+    match_map = {}
+
+    counter = Counter()
+    for node in reversed(graph.nodes):
+        if node.name not in match_map:
+            for layer_type, pattern_matcher in fusion_patterns.items():
+                match = pattern_matcher.match(node)
+                if match:
+                    match_case = pattern_matcher.rewrite(opset=graph.opset)
+                    logger.debug(f"matched pattern {layer_type}")
+                    for _, match in match_case.items():
+                        if "op" not in match:
+                            match.update({"op": layer_type})
+                        if "name" not in match:
+                            match.update({"name": f"{layer_type.lower()}_{counter[layer_type]}"})
+                        counter.update([layer_type])
+                    match_map.update(match_case)
+
+    return match_map
+
+
+def get_previous_node_by_type(node, op_type, trajectory=None):
+    """Recursively find and return the first preceding node of a specified type in the computation graph."""
+    if trajectory is None:
+        trajectory = []
+    node_feeds = node.feeds
+    for node_feed in node_feeds:
+        trajectory.append(node_feed)
+        if node_feed.op == op_type:
+            return trajectory
+        else:
+            return get_previous_node_by_type(node_feed, op_type, trajectory)

onnxslim/core/optimization/dead_node_elimination.py

@@ -0,0 +1,151 @@
+import logging
+
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.third_party.onnx_graphsurgeon.exporters.onnx_exporter import dtype_to_onnx
+from onnxslim.third_party.onnx_graphsurgeon.ir.tensor import Constant, Variable
+
+logger = logging.getLogger("onnxslim")
+
+
+def dead_node_elimination(graph, is_subgraph=False):
+    """Perform in-place constant folding optimizations on the given computational graph by eliminating redundant
+    nodes.
+    """
+    for subgraph in graph.subgraphs():
+        dead_node_elimination(subgraph, is_subgraph=True)
+
+    for node in graph.nodes:
+        if node.op in {"Identity", "Dropout"}:
+            if not is_subgraph:
+                node.erase()
+                logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Pad":
+            if len(node.inputs) > 1 and isinstance(node.inputs[1], Constant):
+                pad_value = node.inputs[1].values.tolist()
+                pad_value = pad_value if isinstance(pad_value, list) else [pad_value]
+                if all(value == 0 for value in pad_value):
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Cast":
+            inp_dtype = next(dtype_to_onnx(input.dtype) for input in node.inputs)
+            if inp_dtype == node.attrs["to"]:
+                node.erase()
+                logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Reshape":
+            if (node.inputs[0].shape and len(node.inputs[0].shape) == 1) and (
+                node.outputs[0].shape and len(node.outputs[0].shape) == 1
+            ):
+                node.erase()
+                logger.debug(f"removing {node.op} op: {node.name}")
+            elif node.inputs[0].shape and node.outputs[0].shape and node.inputs[0].shape == node.outputs[0].shape:
+                node.erase()
+                logger.debug(f"removing {node.op} op: {node.name}")
+            else:
+                node_output_shape = node.outputs[0].shape
+                if node_output_shape and check_shape(node_output_shape) and not isinstance(node.inputs[1], gs.Constant):
+                    shapes = [shape if isinstance(shape, int) else -1 for shape in node_output_shape]
+                    reshape_const = gs.Constant(
+                        f"{node.inputs[1].name}_",
+                        values=np.array(shapes, dtype=np.int64),
+                    )
+                    node.inputs.pop(1)
+                    node.inputs.insert(1, reshape_const)
+                    logger.debug(f"replacing {node.op} op: {node.name}")
+        elif node.op == "Mul":
+            if (isinstance(node.inputs[1], Constant) and isinstance(node.inputs[0], Variable)) or (
+                isinstance(node.inputs[0], Constant) and isinstance(node.inputs[1], Variable)
+            ):
+                idx, constant_variable = get_constant_variable(node, return_idx=True)
+                if np.all(constant_variable.values == 1):
+                    var_idx = 0 if idx == 1 else 1
+                    node.erase(var_idx, 0)
+                    logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Add":
+            if (isinstance(node.inputs[1], Constant) and isinstance(node.inputs[0], Variable)) or (
+                isinstance(node.inputs[0], Constant) and isinstance(node.inputs[1], Variable)
+            ):
+                idx, constant_variable = get_constant_variable(node, return_idx=True)
+                value = constant_variable.values
+                var_idx = 0 if idx == 1 else 1
+                if value.ndim == 0 and value == 0:
+                    node.erase(var_idx, 0)
+                    logger.debug(f"removing {node.op} op: {node.name}")
+                elif np.all(value == 0) and (node.inputs[var_idx].shape == node.outputs[0].shape):
+                    node.erase(var_idx, 0)
+                    logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Expand":
+            # tests/test_onnx_nets.py::TestTimmClass::test_timm[lambda_resnet26rpt_256]
+            if len(node.inputs) > 1 and isinstance(node.inputs[1], Constant):
+                constant_variable = node.inputs[1]
+                value = constant_variable.values
+                if node.inputs[0].shape is not None and node.inputs[0].shape == node.outputs[0].shape:
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+                elif value.ndim == 0 and value == 1:
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Concat":
+            if len(node.inputs) == 1:
+                node.erase()
+                logger.debug(f"removing {node.op} op: {node.name}")
+            else:
+                for input in node.inputs:
+                    if isinstance(input, Constant) and input.values.size == 0:
+                        node.inputs.remove(input)
+        elif node.op == "Sub":
+            if isinstance(node.inputs[1], Constant) and isinstance(node.inputs[0], Variable):
+                constant_variable = node.inputs[1]
+                value = constant_variable.values
+                if value.ndim == 0 and value == 0:
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+                elif np.all(value == 0) and (node.inputs[0].shape == node.outputs[0].shape):
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Div":
+            if isinstance(node.inputs[1], Constant) and isinstance(node.inputs[0], Variable):
+                constant_variable = node.inputs[1]
+                value = constant_variable.values
+                if value.ndim == 0 and value == 1:
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+                elif np.all(value == 1) and (node.inputs[0].shape == node.outputs[0].shape):
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Split":
+            if (
+                len(node.outputs) == 1
+                and node.outputs[0].shape
+                and node.inputs[0].shape
+                and node.outputs[0].shape == node.inputs[0].shape
+            ):
+                node.erase()
+                logger.debug(f"removing {node.op} op: {node.name}")
+        elif node.op == "Resize":
+            mode = node.attrs.get("mode")
+            if mode is None:
+                node.attrs["mode"] = "nearest"
+                logger.debug(f"setting mode to nearest for {node.op} op: {node.name} since it is not set")
+
+
+def check_shape(shapes):
+    """Verify that 'shapes' contains exactly one string and all other elements are positive integers."""
+    string_count = 0
+    non_negative_int_count = 0
+
+    for item in shapes:
+        if isinstance(item, str):
+            string_count += 1
+        elif isinstance(item, int) and item > 0:
+            non_negative_int_count += 1
+
+    return (string_count == 1 and non_negative_int_count == len(shapes) - 1) or non_negative_int_count == len(shapes)
+
+
+def get_constant_variable(node, return_idx=False):
+    """Return the first constant variable found in a node's inputs, optionally including the index."""
+    for idx, input in enumerate(list(node.inputs)):
+        if isinstance(input, Constant):
+            return (idx, input) if return_idx else input

onnxslim/core/optimization/subexpression_elimination.py

@@ -0,0 +1,76 @@
+import logging
+
+from onnxslim.third_party.onnx_graphsurgeon.ir.tensor import Variable
+
+logger = logging.getLogger("onnxslim")
+
+
+def find_and_remove_replaceable_nodes(nodes):
+    """Find and remove duplicate or replaceable nodes in a given list of computational graph nodes."""
+
+    def get_node_key(node):
+        input_names = []
+        for input_node in node.inputs:
+            if isinstance(input_node, Variable):
+                input_names.append(input_node.name)
+        return "_".join(input_names) if input_names else None
+
+    node_dict = {}
+    for node in nodes:
+        key = get_node_key(node)
+        if key:
+            if key in node_dict:
+                node_dict[key].append(node)
+            else:
+                node_dict[key] = [node]
+
+    for key, bucketed_nodes in node_dict.items():
+        if len(bucketed_nodes) > 1:
+            keep_nodes = [True] * len(bucketed_nodes)
+            for i, node in enumerate(bucketed_nodes):
+                if keep_nodes[i]:
+                    for j in range(i + 1, len(bucketed_nodes)):
+                        if keep_nodes[j] and can_be_replaced(node, bucketed_nodes[j]):
+                            keep_nodes[j] = False
+                            existing_node = node
+                            to_be_removed_node = bucketed_nodes[j]
+                            to_be_removed_node.replace_all_uses_with(existing_node)
+                            logger.debug(
+                                f"Node {to_be_removed_node.name} Op {to_be_removed_node.op} can be replaced by {existing_node.name}"
+                            )
+
+
+def sequences_equal(seq1, seq2):
+    """Check if two sequences are equal by comparing their lengths and elements."""
+    length_match = len(seq1) == len(seq2)
+    if not length_match:
+        return False
+
+    return all(elem1 == elem2 for elem1, elem2 in zip(seq1, seq2))
+
+
+def can_be_replaced(node, other_node):
+    """Check if two nodes can be replaced based on their operations, attributes, and inputs."""
+    attrs_match = node.op == other_node.op and node.attrs == other_node.attrs
+    node_input = [input for input in node.inputs if not input.is_empty()]
+    other_node_input = [input for input in other_node.inputs if not input.is_empty()]
+    inputs_match = sequences_equal(node_input, other_node_input)
+
+    return attrs_match and inputs_match
+
+
+def subexpression_elimination(graph):
+    """Perform subexpression elimination on a computational graph to optimize node operations."""
+    nodes_by_op = {}
+
+    for subgraph in graph.subgraphs():
+        subexpression_elimination(subgraph)
+
+    for node in graph.nodes:
+        op = node.op
+        if op not in nodes_by_op:
+            nodes_by_op[op] = []
+        nodes_by_op[op].append(node)
+
+    for nodes in nodes_by_op.values():
+        find_and_remove_replaceable_nodes(nodes)

onnxslim/core/optimization/weight_tying.py

@@ -0,0 +1,59 @@
+import logging
+import os
+
+logger = logging.getLogger("onnxslim")
+from collections import defaultdict
+
+import numpy as np
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+
+THRESHOLD = int(os.getenv("ONNXSLIM_THRESHOLD")) if os.getenv("ONNXSLIM_THRESHOLD") else 1000
+
+
+def tie_weights(graph):
+    """Tie weights in a computational graph to reduce the number of parameters."""
+    tensor_map = graph.tensors()
+    constant_tensors = [tensor for tensor in tensor_map.values() if isinstance(tensor, gs.Constant)]
+
+    sub_graphs = graph.subgraphs(recursive=True)
+    sub_graphs_constant_tensors = [
+        [tensor for name, tensor in sub_graph.tensors().items() if isinstance(tensor, gs.Constant)]
+        for sub_graph in sub_graphs
+    ]
+
+    constant_tensors.extend([tensor for tensors in sub_graphs_constant_tensors for tensor in tensors])
+
+    constant_by_shape = defaultdict(list)
+
+    for constant_tensor in constant_tensors:
+        shape = tuple(constant_tensor.shape)
+        if np.prod(shape) < THRESHOLD:
+            constant_by_shape[shape].append(constant_tensor)
+
+    for nodes in constant_by_shape.values():
+        find_and_remove_replaceable_constants(nodes)
+
+
+def find_and_remove_replaceable_constants(constant_tensors):
+    def replace_constant_references(existing_constant, to_be_removed_constant):
+        users = list(to_be_removed_constant.outputs)
+
+        for user in users:
+            for idx, inp in enumerate(user.inputs):
+                if (inp == to_be_removed_constant) and (inp.name == to_be_removed_constant.name):
+                    user.inputs.pop(idx)
+                    user.inputs.insert(idx, existing_constant)
+
+    if len(constant_tensors) > 1:
+        keep_constants = [True] * len(constant_tensors)
+        for i, constant_tensor in enumerate(constant_tensors):
+            if keep_constants[i]:
+                for j in range(i + 1, len(constant_tensors)):
+                    if keep_constants[j]:
+                        if constant_tensor == constant_tensors[j]:
+                            keep_constants[j] = False
+                            replace_constant_references(constant_tensor, constant_tensors[j])
+                            logger.debug(
+                                f"Constant {constant_tensors[j].name} can be replaced by {constant_tensor.name}"
+                            )

onnxslim/core/pattern/__init__.py

@@ -0,0 +1,249 @@
+import logging
+import re
+from abc import abstractmethod
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.third_party.onnx_graphsurgeon import Constant
+
+logger = logging.getLogger("onnxslim")
+
+
+def get_name(name):
+    """Sanitizes the input string by replacing illegal characters with underscores and prefixing with an underscore if
+    numeric.
+    """
+    _illegal_char_regex = re.compile("[^0-9a-zA-Z_]+")
+    sanitized_name = _illegal_char_regex.sub("_", name)
+    if sanitized_name.isdigit():
+        sanitized_name = f"_{sanitized_name}"
+
+    return sanitized_name
+
+
+class NodeDescriptor:
+    """
+    case 0: input [1, 2, 3, 4, 5] output [0] Optype Name 5  1 i0 i1 i2 i3 i4 o0
+    case 1: input [1, ...]        output [0] Optype Name 1+ 1 i0 o0
+    case 2: input [..., 1, ...]   output [0] Optype Name 1* 1 i0 o0.
+    """
+
+    def __init__(self, node_spec):
+        """Initialize NodeDescriptor with node_spec list requiring at least 4 elements."""
+        if not isinstance(node_spec, list):
+            raise ValueError("node_spec must be a list")
+        if len(node_spec) < 4:
+            raise ValueError(f"node_spec must have at least 4 elements {node_spec}")
+
+        def get_input_info(io_spec):
+            """Parses io_spec to return a tuple of (integer, boolean) indicating the presence of a plus sign in the
+            input.
+            """
+            if not io_spec.isdigit():
+                match = re.search(r"(\d+)([+*])", io_spec)
+                if match:
+                    number = match.group(1)
+                    operator = match.group(2)
+
+                    if not number.isdigit():
+                        raise ValueError(f"input_num and output_num must be integers {io_spec}")
+
+                    if operator == "+":
+                        return int(number), True, "append"
+                    elif operator == "*":
+                        return int(number), True, "free-match"
+                    else:
+                        raise ValueError(f"operator must be + or * {io_spec}")
+
+            return int(io_spec), False, None
+
+        self.op = node_spec[0]
+        self.name = node_spec[1]
+        self.input_num, self.coarse_input_num, self.input_mode = get_input_info(node_spec[2])
+        self.output_num, self.coarse_output_num, self.output_mode = get_input_info(node_spec[3])
+        self.input_names = node_spec[4 : 4 + self.input_num]
+        self.output_names = node_spec[4 + self.input_num :]
+        assert len(self.input_names) == self.input_num
+        assert len(self.output_names) == self.output_num, f"{self.name} {len(self.output_names)} != {self.output_num}"
+
+    def __repr__(self):
+        """Return a string representation of the object, including its name, operation type, input/output counts, and
+        input/output names.
+        """
+        return f"name: {self.name}, type: {self.op}, input_num: {self.input_num}, output_num: {self.output_num}, input_names: {self.input_names}, output_names: {self.output_names}"
+
+    def __dict__(self):
+        """Returns a dictionary representation of the object, with 'name' as the key."""
+        return {
+            "name": self,
+        }
+
+
+class Pattern:
+    def __init__(self, pattern):
+        """Initialize the Pattern class with a given pattern and parse its nodes."""
+        self.pattern = pattern
+        self.nodes = self.parse_nodes()
+
+    def parse_nodes(self):
+        """Parse pattern into a list of NodeDescriptor objects from non-empty, stripped, and split lines."""
+        nodes = self.pattern.split("\n")
+        nodes = [line.strip().split() for line in nodes if line]
+        nodes = [NodeDescriptor(node) for node in nodes if node]
+        return nodes
+
+    def match(self, node):
+        """Match a node against a precompiled pattern."""
+        return self.pattern.match(node)
+
+    def __repr__(self):
+        """Return a string representation of the pattern attribute."""
+        return self.pattern
+
+
+class PatternMatcher:
+    def __init__(self, pattern, priority):
+        """Initialize the PatternMatcher with a given pattern and priority, and prepare node references and output
+        names.
+        """
+        self.pattern = pattern
+        self.priority = priority
+        self.pattern_dict = {node.name: node for node in pattern.nodes}
+        self.output_names = [node.name for node in pattern.nodes if node.op == "output"]
+
+    def get_match_point(self):
+        """Retrieve the match point node from the pattern dictionary based on output node input names."""
+        return self.pattern_dict[self.pattern_dict[self.output_names[0]].input_names[0]]
+
+    def reset_input(self):
+        for k, v in self.pattern_dict.items():
+            if v.op == "input":
+                if hasattr(self, v.name):
+                    delattr(self, v.name)
+
+    def match(self, node):
+        """Match a given node to a pattern by comparing input names with the match point node from the pattern
+        dictionary.
+        """
+        self.reset_input()
+        match_point = self.get_match_point()
+
+        def match_(node, pattern_node):
+            """Match a given node to a pattern by comparing input names with the match point node from the pattern
+            dictionary.
+            """
+            if pattern_node.op == "input":
+                if hasattr(self, pattern_node.name):
+                    if getattr(self, pattern_node.name) == node:
+                        return True
+                    else:
+                        return False
+                else:
+                    setattr(self, pattern_node.name, node)
+                    return True
+            # node is an input variable
+            if not hasattr(node, "op"):
+                return False
+
+            if node.op == pattern_node.op:
+                setattr(self, pattern_node.name, node)
+
+                node_feeds = node.feeds
+                if pattern_node.coarse_input_num:
+                    if len(node_feeds) < len(pattern_node.input_names):
+                        return False
+                else:
+                    if len(node_feeds) != len(pattern_node.input_names):
+                        return False
+
+                if pattern_node.input_mode == "append" or pattern_node.input_mode is None:
+                    pattern_nodes = [
+                        self.pattern_dict[name] if name != "?" else None for name in pattern_node.input_names
+                    ]
+                    all_match = True
+                    for node_feed, pattern_node in zip(node_feeds, pattern_nodes):
+                        if pattern_node is not None:
+                            node_match = match_(node_feed, pattern_node)
+                            if not node_match:
+                                return False
+                            setattr(self, pattern_node.name, node_feed)
+
+                    return all_match
+                elif pattern_node.input_mode == "free-match":
+                    pattern_nodes = [
+                        self.pattern_dict[name] if name != "?" else None for name in pattern_node.input_names
+                    ]
+                    all_match = True
+                    for pattern_node in pattern_nodes:
+                        if pattern_node is not None:
+                            node_match = False
+                            for node_feed in node_feeds:
+                                node_match = match_(node_feed, pattern_node)
+                                if node_match:
+                                    break
+                            if not node_match:
+                                return False
+                            setattr(self, pattern_node.name, node_feed)
+
+                    return all_match
+            return False
+
+        if match_(node, match_point):
+            setattr(self, "output", node.outputs)
+            if self.parameter_check():
+                return True
+
+        return False
+
+    @abstractmethod
+    def rewrite(self, opset=11):
+        """Abstract method to rewrite the graph based on matched patterns, to be implemented by subclasses."""
+        raise NotImplementedError("rewrite method must be implemented")
+
+    def parameter_check(self):
+        """Check and validate parameters, returning True if valid."""
+        return True
+
+
+class PatternGenerator:
+    def __init__(self, onnx_model):
+        """Initialize the PatternGenerator class with an ONNX model and process its graph."""
+        self.graph = gs.import_onnx(onnx_model)
+        self.graph.fold_constants().cleanup().toposort()
+
+    def generate(self):
+        """Generate the inputs, outputs, and nodes from the graph of the initialized ONNX model."""
+        inputs = self.graph.inputs
+        outputs = self.graph.outputs
+        nodes = self.graph.nodes
+
+        template = []
+        for input in inputs:
+            name = get_name(input.name)
+            template.append(
+                " ".join(
+                    ["input", name, "0", str(len(input.outputs))] + [get_name(output.name) for output in input.outputs]
+                )
+            )
+
+        for node in nodes:
+            if node.op != "Constant":
+                name = get_name(node.name)
+                feeds = node.feeds
+                users = node.users
+                template.append(
+                    " ".join(
+                        [node.op, name, str(len(feeds)), str(len(users))]
+                        + ["?" if isinstance(feed, Constant) else get_name(feed.name) for feed in feeds]
+                        + ["?" if isinstance(user, Constant) else get_name(user.name) for user in users]
+                    )
+                )
+
+        for output in outputs:
+            name = get_name(output.name)
+            template.append(
+                " ".join(
+                    ["output", name, str(len(output.inputs)), "0"] + [get_name(input.name) for input in output.inputs]
+                )
+            )
+
+        return "\n".join(template)

onnxslim/core/pattern/elimination/__init__.py

@@ -0,0 +1,5 @@
+from .concat import *
+from .reshape import *
+from .reshape_as import *
+from .slice import *
+from .unsqueeze import *