onnxslim 0.1.69__py3-none-any.whl → 0.1.71__py3-none-any.whl

onnxslim/__init__.py

@@ -2,9 +2,7 @@ import os
 import warnings
 
 from onnxslim.cli import slim
-from onnxslim.core.optimization import OptimizationSettings
 from onnxslim.core.pattern.registry import (
-    DEFAULT_FUSION_PATTERNS,
     register_fusion_pattern,
 )
 from onnxslim.version import __version__

onnxslim/argparser.py

@@ -2,10 +2,28 @@ import argparse
 import dataclasses
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from dataclasses import dataclass, field
-from typing import List, Optional, Type, Union, get_args, get_origin
-
-import onnxslim
-
+from typing import List, Optional, Type, Union, get_args, get_origin, TypedDict, Dict, Literal
+
+from .core.optimization import OptimizationSettings
+from .core.pattern.registry import DEFAULT_FUSION_PATTERNS
+from .version import __version__
+
+
+class OnnxSlimKwargs(TypedDict, total=False):
+    model_check: bool
+    input_shapes: Dict[str, List[int]]
+    inputs: List[str]
+    outputs: List[str]
+    no_shape_infer: bool
+    skip_optimizations: List[str]
+    dtype: Literal["float16", "float32", "uint8", "int8"]
+    skip_fusion_patterns: List[str]
+    size_threshold: int
+    inspect: bool
+    dump_to_disk: bool
+    save_as_external_data: bool
+    model_check_inputs: Optional[List[str]]
+    verbose: bool
 
 def _get_inner_type(arg_type):
     if get_origin(arg_type) is Union:
@@ -42,14 +60,14 @@ class OptimizationArguments:
         default=None,
         metadata={
             "help": "whether to skip some optimizations",
-            "choices": list(onnxslim.OptimizationSettings.keys()),
+            "choices": list(OptimizationSettings.keys()),
         },
     )
     skip_fusion_patterns: Optional[List[str]] = field(
         default=None,
         metadata={
             "help": "whether to skip the fusion of some patterns",
-            "choices": list(onnxslim.DEFAULT_FUSION_PATTERNS.keys()),
+            "choices": list(DEFAULT_FUSION_PATTERNS.keys()),
         },
     )
     size_threshold: int = field(
@@ -173,7 +191,7 @@ class OnnxSlimArgumentParser(ArgumentParser):
         # Add positional arguments separately for ModelArguments
         self.parser.add_argument("input_model", help="input onnx model")
         self.parser.add_argument("output_model", nargs="?", default=None, help="output onnx model")
-        self.parser.add_argument("-v", "--version", action="version", version=onnxslim.__version__)
+        self.parser.add_argument("-v", "--version", action="version", version=__version__)
 
     def parse_args_into_dataclasses(self):
         # Pre-parse arguments to check for `--inspect`

onnxslim/cli/_main.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import onnx
 
+from onnxslim.argparser import OnnxSlimKwargs
 
-def slim(model: str | onnx.ModelProto | list[str | onnx.ModelProto], *args, **kwargs):
+
+def slim(model: str | onnx.ModelProto | list[str | onnx.ModelProto], *args, **kwargs: OnnxSlimKwargs):
     import os
     import time
     from pathlib import Path

onnxslim/core/optimization/__init__.py

@@ -57,10 +57,7 @@ def optimize_model(model: onnx.ModelProto | gs.Graph, skip_fusion_patterns: str
     if OptimizationSettings.graph_fusion:
         logger.debug("Start graph_fusion.")
         fusion_patterns = get_fusion_patterns(skip_fusion_patterns)
-        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).toposort()
+        graph_fusion(graph, fusion_patterns)
         logger.debug("Finish graph_fusion.")
     if OptimizationSettings.dead_node_elimination:
         logger.debug("Start dead_node_elimination.")
@@ -102,9 +99,21 @@ def replace_custom_layer(
     )
 
 
+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:

onnxslim/core/optimization/subexpression_elimination.py

@@ -63,6 +63,9 @@ 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:

onnxslim/third_party/symbolic_shape_infer.py

@@ -530,7 +530,7 @@ class SymbolicShapeInference:
             # (2) opset version >= 9. In older version, initializer is required in graph input by onnx spec.
             # (3) The initializer is not in graph input. The means the node input is "constant" in inference.
             initializers = []
-            if (get_opset(self.out_mp_) >= 9) and node.op_type == "Unsqueeze":
+            if (get_opset(self.out_mp_) >= 9) and (node.op_type == "Unsqueeze" or node.op_type == "ReduceMax"):
                 initializers = [
                     self.initializers_[name]
                     for name in node.input
@@ -1042,6 +1042,9 @@ class SymbolicShapeInference:
     def _infer_Constant(self, node):  # noqa: N802
         """Infer the constant value for a given node and store it in sympy_data_."""
         t = get_attribute(node, "value")
+        # Lower constant nodes to initializers
+        t.name = node.output[0]
+        self.initializers_[node.output[0]] = t
         self.sympy_data_[node.output[0]] = numpy_helper.to_array(t)
 
     def _infer_ConstantOfShape(self, node):  # noqa: N802
@@ -1944,7 +1947,25 @@ class SymbolicShapeInference:
 
     def _infer_Shape(self, node):  # noqa: N802
         """Infers and sets the symbolic shape for the output node in the computation graph."""
-        self.sympy_data_[node.output[0]] = self._get_sympy_shape(node, 0)
+        start = get_attribute(node, "start", 0)
+        end = get_attribute(node, "end", None)
+
+        full_sympy_shape = self._get_sympy_shape(node, 0)
+        num_dims = len(full_sympy_shape)
+
+        if start < 0:
+            start = num_dims + start
+        if end is None:
+            end = num_dims
+        elif end < 0:
+            end = num_dims + end
+
+        assert 0 <= start <= end <= num_dims, (
+            f"reshape start/end invalid: start={start}, end={end}, total_dims={num_dims}"
+        )
+
+        target_sympy_shape = full_sympy_shape[start:end]
+        self.sympy_data_[node.output[0]] = target_sympy_shape
 
     def _infer_Size(self, node):  # noqa: N802
         """Infers and sets the size of the output node by computing the product of its shape in the computation

onnxslim/utils.py

@@ -599,8 +599,8 @@ def get_itemsize(dtype):
     ]:
         return 1
 
-    print(dtype)
-    raise
+    print(f"Unknown ONNX dtype: {dtype}")
+    raise ValueError(f"Unsupported TensorProto dtype: {dtype}")
 
 
 def calculate_tensor_size(tensor):

onnxslim/version.py

@@ -1 +1 @@
-__version__ = "0.1.69"
+__version__ = "0.1.71"

{onnxslim-0.1.69.dist-info → onnxslim-0.1.71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxslim
-Version: 0.1.69
+Version: 0.1.71
 Summary: OnnxSlim: A Toolkit to Help Optimize Onnx Model
 Home-page: https://github.com/inisis/OnnxSlim
 Author: inisis

{onnxslim-0.1.69.dist-info → onnxslim-0.1.71.dist-info}/RECORD

@@ -1,14 +1,14 @@
-onnxslim/__init__.py,sha256=-zmz0A4bd_t8VgNXZVMizIpKJswIlC99hr8juWYm2kc,660
+onnxslim/__init__.py,sha256=ECHGdxzg4b-4SZaPhxM_KulBi-xDbVcVUbpJc8i6a60,571
 onnxslim/__main__.py,sha256=FgDcl6xX8kV_52rB-jPVsmGqidlVhkpe_YhXK75-nFU,75
-onnxslim/argparser.py,sha256=aLVOV7c_HmsvuuAeb1FMRyqXTPRtjwm-bGt4O-7AxMg,8350
-onnxslim/utils.py,sha256=uNh3Mr4DberXtKL6p649QrrLR7gYbldYuIm6QLo14Zg,28043
-onnxslim/version.py,sha256=3nwo1wJ5elZRHnCuCkc-pqzqRZO_38VMOqiKtCotf8M,23
+onnxslim/argparser.py,sha256=pFv3nEZH2BiHO9ejS4Iq5ZuZ3GrpdyRQJypAyR0xF7w,8942
+onnxslim/utils.py,sha256=1CDFOOvZVl6pj1K6eR44IagxbM44UBzZGTbcTgW_ovw,28122
+onnxslim/version.py,sha256=ebu6Nblu_UmpciwG6xnbUJm-16F-ZA6L-sagDt37smo,23
 onnxslim/cli/__init__.py,sha256=kxK27cDgWotBOWRs86rbRQf_dtmniKr1GZJeasxfESE,42
-onnxslim/cli/_main.py,sha256=CPitq8khtOLCrJxvmam61ONppACfjQUzRshZtyabZak,5487
+onnxslim/cli/_main.py,sha256=jEKv9q7y_y9g0GsrfXcnk_wyMVej6jhe9QNPChE7yTs,5550
 onnxslim/core/__init__.py,sha256=d6zfD4rcR51yJIFg1LpGQt2tK4g1WLmA87_7Gna89OU,8469
-onnxslim/core/optimization/__init__.py,sha256=cAzv2YoZMS2krWoylrqDd9HU0aAhSCAM29dZclKqOQo,4777
+onnxslim/core/optimization/__init__.py,sha256=M72h_Jyj9j0bAWI5GRt5MkAnXledYWvH-y9f9FG36Nw,5020
 onnxslim/core/optimization/dead_node_elimination.py,sha256=xAO37JvW7Yb3mT5cXY6K6cWMAvbmFL0pczpcRQfd5Ek,7557
-onnxslim/core/optimization/subexpression_elimination.py,sha256=g_AypWGzW60QaHMqxFZ9elCo8CxqrshzUwtVSPhpgDA,2754
+onnxslim/core/optimization/subexpression_elimination.py,sha256=wauPADo0msVifYGgJ_fiNmPSe7LrJ7A-D3ur5z1wif8,2838
 onnxslim/core/optimization/weight_tying.py,sha256=uAz_2AJa-s-budTcDXIr6ZcME1hXPdm54CzR-u850X4,2347
 onnxslim/core/pattern/__init__.py,sha256=mIqdX7JC7oiypOBbaCC5bG3km7x6MIvqrz4koXzRX1U,9434
 onnxslim/core/pattern/registry.py,sha256=Dv7FqS9E-zLfvWCTh682PNRyTFcZoYiuH2MnfjMNpec,831
@@ -29,7 +29,7 @@ onnxslim/core/pattern/fusion/reduce.py,sha256=dMC7CPlFglrJxugsJWjcc-jQCIa_GIbW1y
 onnxslim/misc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnxslim/misc/tabulate.py,sha256=pytOhVcWPW6B464BvG0rycbz8QA0MAeCSqQ7EO93tqs,99648
 onnxslim/third_party/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-onnxslim/third_party/symbolic_shape_infer.py,sha256=bAyxkpF0ahHvdThvE9Lpbs70pqV6jAivfCcKQHlBTeE,153036
+onnxslim/third_party/symbolic_shape_infer.py,sha256=mxpOiv0QYEWZ0Pm8NW7ABRv2X-L_luBRCCKDri0JlmA,153751
 onnxslim/third_party/_sympy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnxslim/third_party/_sympy/functions.py,sha256=FrrTHPkDhwh_EccLGugbcodeOjfGeM17T5_oVPUjOlA,8877
 onnxslim/third_party/_sympy/numbers.py,sha256=w1dJJcQkKRzLDCJMn70YTwPtrEWPFrhCJpAsZhukJOk,11383
@@ -55,10 +55,10 @@ onnxslim/third_party/onnx_graphsurgeon/logger/logger.py,sha256=L12rrwn33RHH-2WLv
 onnxslim/third_party/onnx_graphsurgeon/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnxslim/third_party/onnx_graphsurgeon/util/exception.py,sha256=KrsHbKEQ4237UbjlODsUzvkXoAY72LZi23ApBeFANWg,786
 onnxslim/third_party/onnx_graphsurgeon/util/misc.py,sha256=kyxInD2SCRLU4wHMeiDEYEHB3871fGks6kQTuF9uATY,8960
-onnxslim-0.1.69.dist-info/licenses/LICENSE,sha256=oHZXw-yrBwdNVGu4JtlZhMgmQHKIZ7BJJlJdhu1HKvI,1062
-onnxslim-0.1.69.dist-info/METADATA,sha256=NMqsOJCdOhIfLwnZnOmPSlUw7PV1Vu8YhIoG1zyR-10,7621
-onnxslim-0.1.69.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnxslim-0.1.69.dist-info/entry_points.txt,sha256=O2QgceCVeGeRhnxRSDRcGiFd0ZNfElwrTiRo1W2V7KA,47
-onnxslim-0.1.69.dist-info/top_level.txt,sha256=EWFTb99i0kc6cC9akqNKp88ipzg17_VZzYN7z1kQNlA,9
-onnxslim-0.1.69.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-onnxslim-0.1.69.dist-info/RECORD,,
+onnxslim-0.1.71.dist-info/licenses/LICENSE,sha256=oHZXw-yrBwdNVGu4JtlZhMgmQHKIZ7BJJlJdhu1HKvI,1062
+onnxslim-0.1.71.dist-info/METADATA,sha256=pYpwf2phxrDay3yTJceIS_MedaWzMTJxhGfdi9hbWcE,7621
+onnxslim-0.1.71.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnxslim-0.1.71.dist-info/entry_points.txt,sha256=O2QgceCVeGeRhnxRSDRcGiFd0ZNfElwrTiRo1W2V7KA,47
+onnxslim-0.1.71.dist-info/top_level.txt,sha256=EWFTb99i0kc6cC9akqNKp88ipzg17_VZzYN7z1kQNlA,9
+onnxslim-0.1.71.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+onnxslim-0.1.71.dist-info/RECORD,,