onnxslim 0.1.80__py3-none-any.whl

onnxslim/__init__.py

@@ -0,0 +1,16 @@
+import os
+import warnings
+
+from onnxslim.cli import slim
+from onnxslim.core.pattern.registry import (
+    register_fusion_pattern,
+)
+from onnxslim.version import __version__
+
+if os.path.dirname(os.path.realpath(__file__)) == os.path.join(os.path.realpath(os.getcwd()), "onnxslim"):
+    message = (
+        "You are importing onnxslim within its own root folder ({}). "
+        "This is not expected to work and may give errors. Please exit the "
+        "onnxslim project source and relaunch your python interpreter."
+    )
+    warnings.warn(message.format(os.getcwd()))

onnxslim/__main__.py

@@ -0,0 +1,4 @@
+from onnxslim.cli._main import main
+
+if __name__ == "__main__":
+    main()

onnxslim/argparser.py

@@ -0,0 +1,215 @@
+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, 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:
+        return next((t for t in get_args(arg_type) if t is not type(None)), str)
+    return arg_type
+
+
+@dataclass
+class ModelArguments:
+    """
+    Args:
+        model (Union[str, onnx.ModelProto]): The ONNX model to be slimmed. It can be either a file path or an `onnx.ModelProto` object.
+
+        output_model (str, optional): File path to save the slimmed model. If None, the model will not be saved.
+    """
+
+    input_model: str = field(metadata={"help": "input onnx model"})
+    output_model: Optional[str] = field(default=None, metadata={"help": "output onnx model"})
+
+
+@dataclass
+class OptimizationArguments:
+    """
+    Args:
+        no_shape_infer (bool, optional): Flag indicating whether to perform shape inference. Default is False.
+
+        no_constant_folding (bool, optional): Flag indicating whether to perform constant folding. Default is False.
+
+        skip_fusion_patterns (str, optional): String representing fusion patterns to skip. Default is None.
+    """
+
+    no_shape_infer: bool = field(default=False, metadata={"help": "whether to disable shape_infer, default false."})
+    skip_optimizations: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "whether to skip some optimizations",
+            "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(DEFAULT_FUSION_PATTERNS.keys()),
+        },
+    )
+    size_threshold: int = field(
+        default=None,
+        metadata={
+            "help": "size threshold in bytes, size larger than this value will not be folded, default None, which means fold all constants",
+        },
+    )
+
+
+@dataclass
+class ModificationArguments:
+    """
+    Args:
+        input_shapes (str, optional): String representing the input shapes. Default is None.
+
+        outputs (str, optional): String representing the outputs. Default is None.
+
+        dtype (str, optional): Data type. Default is None.
+
+        save_as_external_data (bool, optional): Flag indicating whether to split onnx as model and weight. Default is False.
+    """
+
+    input_shapes: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "input shape of the model, INPUT_NAME:SHAPE, e.g. x:1,3,224,224 or x1:1,3,224,224 x2:1,3,224,224"
+        },
+    )
+    inputs: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "input of the model, INPUT_NAME:DTYPE, e.g. y:fp32 or y1:fp32 y2:fp32. If dtype is not specified, the dtype of the input will be the same as the original model if it has dtype, otherwise it will be fp32, available dtype: fp16, fp32, int32"
+        },
+    )
+    outputs: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "output of the model, OUTPUT_NAME:DTYPE, e.g. y:fp32 or y1:fp32 y2:fp32. If dtype is not specified, the dtype of the output will be the same as the original model if it has dtype, otherwise it will be fp32, available dtype: fp16, fp32, int32"
+        },
+    )
+    dtype: Optional[str] = field(
+        default=None, metadata={"help": "convert data format to fp16 or fp32.", "choices": ["fp16", "fp32"]}
+    )
+    save_as_external_data: bool = field(
+        default=False, metadata={"help": "split onnx as model and weight, default False."}
+    )
+
+
+@dataclass
+class CheckerArguments:
+    """
+    Args:
+        model_check (bool, optional): Flag indicating whether to perform model checking. Default is False.
+
+        model_check_inputs (str, optional): The shape or tensor used for model check. Default is None.
+
+        inspect (bool, optional): Flag indicating whether to inspect the model. Default is False.
+
+        dump_to_disk (bool, optional): Flag indicating whether to dump the model detail to disk. Default is False.
+
+        verbose (bool, optional): Flag indicating whether to print verbose logs. Default is False.
+    """
+
+    model_check: bool = field(default=False, metadata={"help": "enable model check"})
+    model_check_inputs: Optional[List[str]] = field(
+        default=None,
+        metadata={
+            "help": "Works only when model_check is enabled, Input shape of the model or numpy data path, INPUT_NAME:SHAPE or INPUT_NAME:DATAPATH, e.g. x:1,3,224,224 or x1:1,3,224,224 x2:data.npy. Useful when input shapes are dynamic."
+        },
+    )
+    inspect: bool = field(default=False, metadata={"help": "inspect model, default False."})
+    dump_to_disk: bool = field(default=False, metadata={"help": "dump model info to disk, default False."})
+    verbose: bool = field(default=False, metadata={"help": "verbose mode, default False."})
+
+
+class OnnxSlimArgumentParser(ArgumentParser):
+    def __init__(self, *argument_dataclasses: Type, **kwargs):
+        if "formatter_class" not in kwargs:
+            kwargs["formatter_class"] = ArgumentDefaultsHelpFormatter
+        super().__init__(**kwargs)
+        self.argument_dataclasses = argument_dataclasses
+        self.parser = argparse.ArgumentParser(
+            description="OnnxSlim: A Toolkit to Help Optimizer Onnx Model",
+            formatter_class=argparse.RawDescriptionHelpFormatter,
+        )
+        self._add_arguments()
+
+    def _add_arguments(self):
+        for dataclass_type in self.argument_dataclasses:
+            if dataclass_type is ModelArguments:
+                continue
+            for field_name, field_def in dataclass_type.__dataclass_fields__.items():
+                arg_type = _get_inner_type(field_def.type)
+                default_value = field_def.default if field_def.default is not field_def.default_factory else None
+                help_text = field_def.metadata.get("help", "")
+                nargs = "+" if get_origin(arg_type) == list else None
+                choices = field_def.metadata.get("choices", None)
+                if choices and default_value is not None and default_value not in choices:
+                    raise ValueError(
+                        f"Invalid default value '{default_value}' for argument '{field_name}'. Must be one of {choices}."
+                    )
+                arg_type = get_args(arg_type)[0] if get_args(arg_type) else arg_type
+                if arg_type == bool:
+                    self.parser.add_argument(
+                        f"--{field_name.replace('_', '-')}",
+                        action="store_true",
+                        default=default_value,
+                        help=help_text,
+                    )
+                else:
+                    self.parser.add_argument(
+                        f"--{field_name.replace('_', '-')}",
+                        type=arg_type,
+                        default=default_value,
+                        nargs=nargs,
+                        choices=choices,
+                        help=help_text,
+                    )
+
+        # 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=__version__)
+
+    def parse_args_into_dataclasses(self):
+        # Pre-parse arguments to check for `--inspect`
+        pre_parsed_args, _ = self.parser.parse_known_args()
+        if pre_parsed_args.inspect:
+            for action in self.parser._actions:
+                if action.dest == "input_model":
+                    action.nargs = "+"
+                    break
+
+        args = self.parser.parse_args()
+        args_dict = vars(args)
+
+        outputs = []
+        for dtype in self.argument_dataclasses:
+            keys = {f.name for f in dataclasses.fields(dtype) if f.init}
+            inputs = {k: v for k, v in args_dict.items() if k in keys}
+            obj = dtype(**inputs)
+            outputs.append(obj)
+
+        return (*outputs,)

onnxslim/cli/__init__.py

@@ -0,0 +1 @@
+from onnxslim.cli._main import main, slim

onnxslim/cli/_main.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import onnx
+
+from onnxslim.argparser import OnnxSlimKwargs
+
+
+def slim(model: str | onnx.ModelProto | list[str | onnx.ModelProto], *args, **kwargs: OnnxSlimKwargs):
+    import os
+    import time
+    from pathlib import Path
+
+    from onnxslim.core import (
+        OptimizationSettings,
+        convert_data_format,
+        freeze,
+        input_modification,
+        input_shape_modification,
+        optimize,
+        output_modification,
+        shape_infer,
+    )
+    from onnxslim.utils import (
+        TensorInfo,
+        check_onnx,
+        check_point,
+        check_result,
+        dump_model_info_to_disk,
+        init_logging,
+        onnxruntime_inference,
+        print_model_info_as_table,
+        save,
+        summarize_model,
+        update_outputs_dims,
+    )
+
+    output_model = args[0] if len(args) > 0 else kwargs.get("output_model", None)
+    model_check = kwargs.get("model_check", False)
+    input_shapes = kwargs.get("input_shapes", None)
+    inputs = kwargs.get("inputs", None)
+    outputs = kwargs.get("outputs", None)
+    no_shape_infer = kwargs.get("no_shape_infer", False)
+    skip_optimizations = kwargs.get("skip_optimizations", None)
+    dtype = kwargs.get("dtype", None)
+    skip_fusion_patterns = kwargs.get("skip_fusion_patterns", None)
+    size_threshold = kwargs.get("size_threshold", None)
+    size_threshold = int(size_threshold) if size_threshold else None
+    kwargs.get("inspect", False)
+    dump_to_disk = kwargs.get("dump_to_disk", False)
+    save_as_external_data = kwargs.get("save_as_external_data", False)
+    model_check_inputs = kwargs.get("model_check_inputs", None)
+    verbose = kwargs.get("verbose", False)
+
+    logger = init_logging(verbose)
+
+    MAX_ITER = int(os.getenv("ONNXSLIM_MAX_ITER")) if os.getenv("ONNXSLIM_MAX_ITER") else 10
+
+    start_time = time.time()
+
+    def get_info(model, inspect=False):
+        if isinstance(model, str):
+            model_name = Path(model).name
+            model = onnx.load(model)
+        else:
+            model_name = "OnnxModel"
+
+        freeze(model)
+
+        if not inspect:
+            return model_name, model
+
+        model_info = summarize_model(model, model_name)
+
+        return model_info
+
+    if isinstance(model, list):
+        model_info_list = [get_info(m, inspect=True) for m in model]
+
+        if dump_to_disk:
+            [dump_model_info_to_disk(info) for info in model_info_list]
+
+        print_model_info_as_table(model_info_list)
+
+        return
+    else:
+        model_name, model = get_info(model)
+        if output_model:
+            original_info = summarize_model(model, model_name)
+
+    if inputs:
+        model = input_modification(model, inputs)
+
+    if input_shapes:
+        model = input_shape_modification(model, input_shapes)
+
+    if outputs:
+        model = output_modification(model, outputs)
+
+    if model_check:
+        input_data_dict, raw_onnx_output, model = check_onnx(model, model_check_inputs)
+
+    output_info = {TensorInfo(o).name: TensorInfo(o).shape for o in model.graph.output}
+
+    if not no_shape_infer:
+        model = shape_infer(model)
+
+    OptimizationSettings.reset(skip_optimizations)
+    if OptimizationSettings.enabled():
+        graph_check_point = check_point(model)
+        while MAX_ITER > 0:
+            logger.debug(f"iter: {MAX_ITER}")
+            model = optimize(model, skip_fusion_patterns, size_threshold)
+            if not no_shape_infer:
+                model = shape_infer(model)
+            graph = check_point(model)
+            if graph == graph_check_point:
+                logger.debug(f"converged at iter: {MAX_ITER}")
+                break
+            else:
+                graph_check_point = graph
+
+            MAX_ITER -= 1
+
+    if dtype:
+        model = convert_data_format(model, dtype)
+
+    model = update_outputs_dims(model, output_dims=output_info)
+
+    if model_check:
+        slimmed_onnx_output, model = onnxruntime_inference(model, input_data_dict)
+        if not check_result(raw_onnx_output, slimmed_onnx_output):
+            return None
+
+    if not output_model:
+        return model
+
+    slimmed_info = summarize_model(model, output_model)
+    save(model, output_model, model_check, save_as_external_data, slimmed_info)
+
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    print_model_info_as_table(
+        [original_info, slimmed_info],
+        elapsed_time,
+    )
+
+
+def main():
+    """Entry point for the OnnxSlim toolkit, processes command-line arguments and passes them to the slim function."""
+    from onnxslim.argparser import (
+        CheckerArguments,
+        ModelArguments,
+        ModificationArguments,
+        OnnxSlimArgumentParser,
+        OptimizationArguments,
+    )
+
+    argument_parser = OnnxSlimArgumentParser(
+        ModelArguments, OptimizationArguments, ModificationArguments, CheckerArguments
+    )
+    model_args, optimization_args, modification_args, checker_args = argument_parser.parse_args_into_dataclasses()
+
+    if not checker_args.inspect and checker_args.dump_to_disk:
+        argument_parser.error("dump_to_disk can only be used with --inspect")
+
+    if not optimization_args.no_shape_infer:
+        from onnxslim.utils import check_onnx_compatibility, is_onnxruntime_available
+
+        if is_onnxruntime_available():
+            check_onnx_compatibility()
+
+    slim(
+        model_args.input_model,
+        model_args.output_model,
+        **optimization_args.__dict__,
+        **modification_args.__dict__,
+        **checker_args.__dict__,
+    )
+
+    return 0

onnxslim/core/__init__.py

@@ -0,0 +1,219 @@
+from __future__ import annotations
+
+import logging
+import os
+import tempfile
+from typing import Optional
+
+import numpy as np
+import onnx
+from onnx import checker
+
+import onnxslim.third_party.onnx_graphsurgeon as gs
+from onnxslim.core.optimization import OptimizationSettings, optimize_model
+from onnxslim.third_party.onnx_graphsurgeon.exporters.onnx_exporter import dtype_to_onnx
+from onnxslim.third_party.onnx_graphsurgeon.ir.tensor import Constant
+from onnxslim.third_party.symbolic_shape_infer import SymbolicShapeInference
+from onnxslim.utils import save
+
+logger = logging.getLogger("onnxslim")
+
+
+DEBUG = bool(os.getenv("ONNXSLIM_DEBUG"))
+AUTO_MERGE = True if os.getenv("ONNXSLIM_AUTO_MERGE") is None else bool(int(os.getenv("ONNXSLIM_AUTO_MERGE")))
+FORCE_ONNXRUNTIME_SHAPE_INFERENCE = bool(os.getenv("ONNXSLIM_FORCE_ONNXRUNTIME_SHAPE_INFERENCE"))
+
+
+def input_shape_modification(model: onnx.ModelProto, input_shapes: str) -> onnx.ModelProto:
+    """Modifies input tensor shapes in the ONNX model according to the specified input_shapes string."""
+    if not input_shapes:
+        return
+
+    graph = gs.import_onnx(model)
+    input_names = [input.name for input in graph.inputs]
+    tensors = graph.tensors()
+
+    for input_shape in input_shapes:
+        key, values = input_shape.rsplit(":", 1)
+        values_list = [int(value) for value in values.split(",")]
+        if key not in input_names:
+            raise Exception(f"Input name {key} not found in model, available keys: {' '.join(input_names)}")
+        tensors[key].shape = values_list
+
+    for tensor in tensors.values():
+        if tensor.name not in input_names:
+            if isinstance(tensor, Constant):
+                continue
+            tensor.shape = None
+
+    model = gs.export_onnx(graph)
+
+    return model
+
+
+def output_modification(model: onnx.ModelProto, outputs: str) -> onnx.ModelProto:
+    """Modifies the output layers of the ONNX model based on specified output names and data types."""
+    graph = gs.import_onnx(model)
+    graph.outputs.clear()
+    tensors = graph.tensors()
+    for output in outputs:
+        values = output.rsplit(":", 1)
+        if len(values) == 1:
+            key = values[0]
+            if key not in tensors.keys():
+                raise Exception(f"Output name {key} not found in model, available keys: {' '.join(tensors.keys())}")
+            dtype = tensors[key].dtype
+            if dtype is None:
+                dtype = np.float32
+                logger.warning(f"Output layer {key} has no dtype, set to default {dtype}")
+        else:
+            key, dtype = values
+            if dtype == "fp16":
+                dtype = np.float16
+            elif dtype == "fp32":
+                dtype = np.float32
+            elif dtype == "int32":
+                dtype = np.int32
+            elif dtype == "bool":
+                dtype = bool
+            else:
+                raise Exception(f"Output layer {key} assigned unsupported dtype {dtype}")
+
+        graph.outputs.append(tensors[key].to_variable(dtype=dtype, shape=tensors[key].shape))
+
+    graph.cleanup(remove_unused_graph_inputs=True).toposort()
+    model = gs.export_onnx(graph)
+
+    return model
+
+
+def input_modification(model: onnx.ModelProto, inputs: str) -> onnx.ModelProto:
+    """Modifies the output layers of the ONNX model based on specified output names and data types."""
+    graph = gs.import_onnx(model)
+    graph.inputs.clear()
+    tensors = graph.tensors()
+    for input in inputs:
+        values = input.rsplit(":", 1)
+        if len(values) == 1:
+            key = values[0]
+            if key not in tensors.keys():
+                raise Exception(f"Input name {key} not found in model, available keys: {' '.join(tensors.keys())}")
+            dtype = tensors[key].dtype
+            if dtype is None:
+                dtype = np.float32
+                logger.warning(f"Input layer {key} has no dtype, set to default {dtype}")
+        else:
+            key, dtype = values
+            if dtype == "fp16":
+                dtype = np.float16
+            elif dtype == "fp32":
+                dtype = np.float32
+            elif dtype == "int32":
+                dtype = np.int32
+            elif dtype == "bool":
+                dtype = bool
+            else:
+                raise Exception(f"Input layer {key} assigned unsupported dtype {dtype}")
+
+        graph.inputs.append(tensors[key].to_variable(dtype=dtype, shape=tensors[key].shape))
+
+    graph.cleanup(remove_unused_graph_inputs=True).toposort()
+    model = gs.export_onnx(graph)
+
+    return model
+
+
+def shape_infer(model: onnx.ModelProto):
+    """Infer tensor shapes in an ONNX model using symbolic and static shape inference techniques."""
+    logger.debug("Start shape inference.")
+    if FORCE_ONNXRUNTIME_SHAPE_INFERENCE:
+        logger.debug("force onnxruntime shape infer.")
+        return SymbolicShapeInference.infer_shapes(model, auto_merge=AUTO_MERGE)
+    try:
+        logger.debug("try onnxruntime shape infer.")
+        model = SymbolicShapeInference.infer_shapes(model, auto_merge=AUTO_MERGE)
+    except Exception as err:
+        logger.debug(f"onnxruntime shape infer failed, try onnx shape infer. {err}")
+        if model.ByteSize() >= checker.MAXIMUM_PROTOBUF:
+            tmp_dir = tempfile.TemporaryDirectory()
+            tmp_path = os.path.join(tmp_dir.name, "tmp.onnx")
+            tmp_infer_path = os.path.join(tmp_dir.name, "tmp_infer.onnx")
+            save(model, tmp_path)
+            onnx.shape_inference.infer_shapes_path(tmp_path, tmp_infer_path)
+            model = onnx.load(tmp_infer_path)
+        else:
+            model = onnx.shape_inference.infer_shapes(model)
+    if DEBUG:
+        onnx.save(model, "debug_shape_infer.onnx")
+    logger.debug("Finish shape inference.")
+    return model
+
+
+def optimize(model: onnx.ModelProto, skip_fusion_patterns: str | None = None, size_threshold: int | None = None):
+    """Optimize the given ONNX model with options to skip specific fusion patterns and return the optimized model."""
+    logger.debug("Start converting model to gs.")
+    graph = gs.import_onnx(model).toposort()
+    logger.debug("Finish converting model to gs.")
+    if OptimizationSettings.constant_folding:
+        logger.debug("Start constant folding.")
+        graph.fold_constants(size_threshold=size_threshold).cleanup().toposort()
+        logger.debug("Finish constant folding.")
+    logger.debug("Start optimize model.")
+    model = optimize_model(graph, skip_fusion_patterns)
+    logger.debug("Finish optimize model.")
+    if DEBUG:
+        onnx.save(model, "debug_slim.onnx")
+
+    return model
+
+
+def convert_data_format(model: onnx.ModelProto, dtype: str) -> onnx.ModelProto:
+    """Convert ONNX model data format to specified dtype, supporting 'fp16' and 'fp32'."""
+    if dtype == "fp16":
+        from onnxconverter_common import float16
+
+        model = float16.convert_float_to_float16(model)
+    elif dtype == "fp32":
+        graph = gs.import_onnx(model).toposort()
+
+        for node in graph.nodes:
+            if node.op == "Cast":
+                inp_dtype = next(input.dtype for input in node.inputs)
+                if inp_dtype in [np.float16, np.float32]:
+                    node.erase()
+                else:
+                    outp_dtype = next(output.dtype for output in node.outputs)
+                    if outp_dtype == np.float16:
+                        node.attrs["to"] = dtype_to_onnx(np.float32)
+                        node.outputs[0].dtype = np.float32
+            elif node.op == "ConstantOfShape":
+                if hasattr(node, "attrs") and "value" in node.attrs:
+                    if node.attrs["value"].dtype == np.float16:
+                        node.attrs["value"].values = node.attrs["value"].values.astype(np.float32)
+                        node.outputs[0].dtype = np.float32
+
+        for tensor in graph.tensors().values():
+            if isinstance(tensor, gs.Variable) and tensor.dtype == np.float16:
+                tensor.dtype = np.float32
+            elif isinstance(tensor, gs.Constant) and tensor.dtype == np.float16:
+                tensor.values = tensor.values.astype(np.float32)
+
+        graph.cleanup(remove_unused_graph_inputs=True).toposort()
+        model = gs.export_onnx(graph)
+
+    return model
+
+
+def freeze(model: onnx.ModelProto):
+    """Freeze the input layers of an ONNX model by removing the initializers from the input graph."""
+    inputs = model.graph.input
+    name_to_input = {}
+    for input in inputs:
+        if input.name in name_to_input:
+            logger.warning(f"Duplicate input name: {input.name}")
+        name_to_input[input.name] = input
+
+    for initializer in model.graph.initializer:
+        if initializer.name in name_to_input:
+            inputs.remove(name_to_input[initializer.name])
+            name_to_input.pop(initializer.name)