onnxslim 0.1.38__tar.gz → 0.1.75__tar.gz

onnxslim-0.1.75/PKG-INFO

@@ -0,0 +1,146 @@
+Metadata-Version: 2.4
+Name: onnxslim
+Version: 0.1.75
+Summary: OnnxSlim: A Toolkit to Help Optimize Onnx Model
+Home-page: https://github.com/inisis/OnnxSlim
+Author: inisis
+Author-email: desmond.yao@buaa.edu.cn
+License: MIT
+Project-URL: Bug Tracker, https://github.com/inisis/OnnxSlim/issues
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.6
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: onnx
+Requires-Dist: sympy>=1.13.3
+Requires-Dist: packaging
+Requires-Dist: colorama
+Requires-Dist: ml_dtypes
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license
+Dynamic: license-file
+Dynamic: project-url
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# OnnxSlim
+
+<p align="center">
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://img.shields.io/pypi/v/onnxslim?color=blue" />
+    </a>
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://static.pepy.tech/badge/onnxslim/week" />
+    </a>
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://static.pepy.tech/badge/onnxslim/month" />
+    </a>    
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://static.pepy.tech/badge/onnxslim" />
+    </a>   
+    <a href="https://github.com/inisis/onnxslim/actions/workflows/ci.yaml">
+        <img src="https://github.com/inisis/onnxslim/actions/workflows/ci.yml/badge.svg" />
+    </a>
+    <a href="https://codecov.io/gh/inisis/onnxslim" > 
+        <img src="https://codecov.io/gh/inisis/onnxslim/branch/main/graph/badge.svg?token=C69ZH6802N"/> 
+    </a>    
+    <a href="https://muhammadrizwanmunawar.medium.com/boost-onnx-load-speed-by-10-15-with-onnxslims-python-package-d401eb8c2e69">
+        <img src="https://img.shields.io/badge/Blog-OnnxSlim?style=flat&label=OnnxSlim" />
+    </a>
+    <a href="https://deepwiki.com/inisis/OnnxSlim"><img src="https://img.shields.io/badge/DeepWiki-inisis%2FOnnxSlim-blue.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACwAAAAyCAYAAAAnWDnqAAAAAXNSR0IArs4c6QAAA05JREFUaEPtmUtyEzEQhtWTQyQLHNak2AB7ZnyXZMEjXMGeK/AIi+QuHrMnbChYY7MIh8g01fJoopFb0uhhEqqcbWTp06/uv1saEDv4O3n3dV60RfP947Mm9/SQc0ICFQgzfc4CYZoTPAswgSJCCUJUnAAoRHOAUOcATwbmVLWdGoH//PB8mnKqScAhsD0kYP3j/Yt5LPQe2KvcXmGvRHcDnpxfL2zOYJ1mFwrryWTz0advv1Ut4CJgf5uhDuDj5eUcAUoahrdY/56ebRWeraTjMt/00Sh3UDtjgHtQNHwcRGOC98BJEAEymycmYcWwOprTgcB6VZ5JK5TAJ+fXGLBm3FDAmn6oPPjR4rKCAoJCal2eAiQp2x0vxTPB3ALO2CRkwmDy5WohzBDwSEFKRwPbknEggCPB/imwrycgxX2NzoMCHhPkDwqYMr9tRcP5qNrMZHkVnOjRMWwLCcr8ohBVb1OMjxLwGCvjTikrsBOiA6fNyCrm8V1rP93iVPpwaE+gO0SsWmPiXB+jikdf6SizrT5qKasx5j8ABbHpFTx+vFXp9EnYQmLx02h1QTTrl6eDqxLnGjporxl3NL3agEvXdT0WmEost648sQOYAeJS9Q7bfUVoMGnjo4AZdUMQku50McDcMWcBPvr0SzbTAFDfvJqwLzgxwATnCgnp4wDl6Aa+Ax283gghmj+vj7feE2KBBRMW3FzOpLOADl0Isb5587h/U4gGvkt5v60Z1VLG8BhYjbzRwyQZemwAd6cCR5/XFWLYZRIMpX39AR0tjaGGiGzLVyhse5C9RKC6ai42ppWPKiBagOvaYk8lO7DajerabOZP46Lby5wKjw1HCRx7p9sVMOWGzb/vA1hwiWc6jm3MvQDTogQkiqIhJV0nBQBTU+3okKCFDy9WwferkHjtxib7t3xIUQtHxnIwtx4mpg26/HfwVNVDb4oI9RHmx5WGelRVlrtiw43zboCLaxv46AZeB3IlTkwouebTr1y2NjSpHz68WNFjHvupy3q8TFn3Hos2IAk4Ju5dCo8B3wP7VPr/FGaKiG+T+v+TQqIrOqMTL1VdWV1DdmcbO8KXBz6esmYWYKPwDL5b5FA1a0hwapHiom0r/cKaoqr+27/XcrS5UwSMbQAAAABJRU5ErkJggg==" alt="DeepWiki"></a> 
+</p>
+
+OnnxSlim can help you slim your onnx model, with less operators, but same accuracy, better inference speed.
+
+- 🚀 2025/05/17: OnnxSlim is merged into [optimum](https://github.com/huggingface/optimum) 🤗🤗🤗
+- 🚀 2025/04/30: Rank 1st in the [AICAS 2025 LLM inference optimization challenge](https://tianchi.aliyun.com/competition/entrance/532289/customize588)
+- 🚀 2025/01/28: Achieved 1M downloads
+- 🚀 2024/06/23: OnnxSlim is merged into [transformers.js](https://github.com/huggingface/transformers.js) 🤗🤗🤗
+- 🚀 2024/06/02: OnnxSlim is merged into [ultralytics](https://github.com/ultralytics/ultralytics) ❤️❤️❤️
+- 🚀 2024/04/30: Rank 1st in the [AICAS 2024 LLM inference optimization challenge](https://tianchi.aliyun.com/competition/entrance/532170/customize440) held by Arm and T-head
+- 🚀 2024/01/25: OnnxSlim is merged to [mnn-llm](https://github.com/wangzhaode/mnn-llm), performance increased by 5%
+
+# Benchmark
+
+![Image](https://github.com/user-attachments/assets/fefc79f1-5d8d-486b-935a-a088846b3900)
+
+# Installation
+
+## Using Prebuilt
+
+```bash
+pip install onnxslim
+```
+
+## Install From Source
+
+```bash
+pip install git+https://github.com/inisis/OnnxSlim@main
+```
+
+## Install From Local
+
+```bash
+git clone https://github.com/inisis/OnnxSlim && cd OnnxSlim/
+pip install .
+```
+
+# How to use
+
+## Bash
+
+```bash
+onnxslim your_onnx_model slimmed_onnx_model
+```
+
+<div align=left><img src="https://raw.githubusercontent.com/inisis/onnxslim/main/images/onnxslim.gif"></div>
+
+## Inscript
+
+```inscript
+import onnx
+import onnxslim
+
+model = onnx.load("model.onnx")
+slimmed_model = onnxslim.slim(model)
+onnx.save(slimmed_model, "slimmed_model.onnx")
+```
+
+For more usage, see onnxslim -h or refer to our [examples](./examples)
+
+# Projects using OnnxSlim
+
+- <img src="https://avatars.githubusercontent.com/u/131524?s=48&v=4" width="22" height="22"/>[Mozilla/smart_autofill](https://github.com/mozilla/smart_autofill)
+- <img src="https://avatars.githubusercontent.com/u/1961952?s=48&v=4" width="22" height="22"/>[alibaba/MNN](https://github.com/alibaba/MNN)
+- <img src="https://avatars.githubusercontent.com/u/23534030?s=48&v=4" width="22" height="22"/>[PaddlePaddle/PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR)
+- <img src="https://avatars.githubusercontent.com/u/25720743?s=48&v=4" width="22" height="22"/>[huggingface/transformers.js](https://github.com/huggingface/transformers.js)
+- <img src="https://avatars.githubusercontent.com/u/25720743?s=48&v=4" width="22" height="22"/>[huggingface/optimum](https://github.com/huggingface/optimum)
+- <img src="https://avatars.githubusercontent.com/u/86091366?s=48&v=4" width="22" height="22"/>[THU-MIG/yolov10](https://github.com/THU-MIG/yolov10)
+- <img src="https://avatars.githubusercontent.com/u/26833451?s=48&v=4" width="22" height="22"/>[ultralytics/ultralytics](https://github.com/ultralytics/ultralytics)
+- <img src="https://avatars.githubusercontent.com/u/109945100?s=48&v=4" width="22" height="22"/>[ModelScope/FunASR](https://github.com/modelscope/FunASR)
+- <img src="https://avatars.githubusercontent.com/u/1961952?s=48&v=4" width="22" height="22"/>[alibaba/MNN-LLM](https://github.com/wangzhaode/mnn-llm)
+- <img src="https://avatars.githubusercontent.com/u/126587470?s=48&v=4" width="22" height="22"/>[deepghs/imgutils](https://github.com/deepghs/imgutils)
+- <img src="https://avatars.githubusercontent.com/u/48153283?s=48&v=4" width="22" height="22"/>[sunsmarterjie/yolov12](https://github.com/sunsmarterjie/yolov12)
+- <img src="https://avatars.githubusercontent.com/u/147458884?s=48&v=4" width="22" height="22"/>[nndeploy/nndeploy](https://github.com/nndeploy/nndeploy)
+- <img src="https://avatars.githubusercontent.com/u/111754012?s=48&v=4" width="22" height="22"/>[CVCUDA/CV-CUDA](https://github.com/CVCUDA/CV-CUDA)
+
+# References
+
+> - [onnx-graphsurgeon](https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon)
+> - [Polygraphy](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/polygraphy)
+> - [onnx-simplifier](https://github.com/daquexian/onnx-simplifier)
+> - [tabulate](https://github.com/astanin/python-tabulate)
+> - [onnxruntime](https://github.com/microsoft/onnxruntime)
+
+# Contact
+
+Discord: https://discord.gg/nRw2Fd3VUS QQ Group: `873569894`

onnxslim-0.1.75/README.md

@@ -0,0 +1,112 @@
+# OnnxSlim
+
+<p align="center">
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://img.shields.io/pypi/v/onnxslim?color=blue" />
+    </a>
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://static.pepy.tech/badge/onnxslim/week" />
+    </a>
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://static.pepy.tech/badge/onnxslim/month" />
+    </a>    
+    <a href="https://pypi.org/project/onnxslim">
+        <img src="https://static.pepy.tech/badge/onnxslim" />
+    </a>   
+    <a href="https://github.com/inisis/onnxslim/actions/workflows/ci.yaml">
+        <img src="https://github.com/inisis/onnxslim/actions/workflows/ci.yml/badge.svg" />
+    </a>
+    <a href="https://codecov.io/gh/inisis/onnxslim" > 
+        <img src="https://codecov.io/gh/inisis/onnxslim/branch/main/graph/badge.svg?token=C69ZH6802N"/> 
+    </a>    
+    <a href="https://muhammadrizwanmunawar.medium.com/boost-onnx-load-speed-by-10-15-with-onnxslims-python-package-d401eb8c2e69">
+        <img src="https://img.shields.io/badge/Blog-OnnxSlim?style=flat&label=OnnxSlim" />
+    </a>
+    <a href="https://deepwiki.com/inisis/OnnxSlim"><img src="https://img.shields.io/badge/DeepWiki-inisis%2FOnnxSlim-blue.svg?logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAACwAAAAyCAYAAAAnWDnqAAAAAXNSR0IArs4c6QAAA05JREFUaEPtmUtyEzEQhtWTQyQLHNak2AB7ZnyXZMEjXMGeK/AIi+QuHrMnbChYY7MIh8g01fJoopFb0uhhEqqcbWTp06/uv1saEDv4O3n3dV60RfP947Mm9/SQc0ICFQgzfc4CYZoTPAswgSJCCUJUnAAoRHOAUOcATwbmVLWdGoH//PB8mnKqScAhsD0kYP3j/Yt5LPQe2KvcXmGvRHcDnpxfL2zOYJ1mFwrryWTz0advv1Ut4CJgf5uhDuDj5eUcAUoahrdY/56ebRWeraTjMt/00Sh3UDtjgHtQNHwcRGOC98BJEAEymycmYcWwOprTgcB6VZ5JK5TAJ+fXGLBm3FDAmn6oPPjR4rKCAoJCal2eAiQp2x0vxTPB3ALO2CRkwmDy5WohzBDwSEFKRwPbknEggCPB/imwrycgxX2NzoMCHhPkDwqYMr9tRcP5qNrMZHkVnOjRMWwLCcr8ohBVb1OMjxLwGCvjTikrsBOiA6fNyCrm8V1rP93iVPpwaE+gO0SsWmPiXB+jikdf6SizrT5qKasx5j8ABbHpFTx+vFXp9EnYQmLx02h1QTTrl6eDqxLnGjporxl3NL3agEvXdT0WmEost648sQOYAeJS9Q7bfUVoMGnjo4AZdUMQku50McDcMWcBPvr0SzbTAFDfvJqwLzgxwATnCgnp4wDl6Aa+Ax283gghmj+vj7feE2KBBRMW3FzOpLOADl0Isb5587h/U4gGvkt5v60Z1VLG8BhYjbzRwyQZemwAd6cCR5/XFWLYZRIMpX39AR0tjaGGiGzLVyhse5C9RKC6ai42ppWPKiBagOvaYk8lO7DajerabOZP46Lby5wKjw1HCRx7p9sVMOWGzb/vA1hwiWc6jm3MvQDTogQkiqIhJV0nBQBTU+3okKCFDy9WwferkHjtxib7t3xIUQtHxnIwtx4mpg26/HfwVNVDb4oI9RHmx5WGelRVlrtiw43zboCLaxv46AZeB3IlTkwouebTr1y2NjSpHz68WNFjHvupy3q8TFn3Hos2IAk4Ju5dCo8B3wP7VPr/FGaKiG+T+v+TQqIrOqMTL1VdWV1DdmcbO8KXBz6esmYWYKPwDL5b5FA1a0hwapHiom0r/cKaoqr+27/XcrS5UwSMbQAAAABJRU5ErkJggg==" alt="DeepWiki"></a> 
+</p>
+
+OnnxSlim can help you slim your onnx model, with less operators, but same accuracy, better inference speed.
+
+- 🚀 2025/05/17: OnnxSlim is merged into [optimum](https://github.com/huggingface/optimum) 🤗🤗🤗
+- 🚀 2025/04/30: Rank 1st in the [AICAS 2025 LLM inference optimization challenge](https://tianchi.aliyun.com/competition/entrance/532289/customize588)
+- 🚀 2025/01/28: Achieved 1M downloads
+- 🚀 2024/06/23: OnnxSlim is merged into [transformers.js](https://github.com/huggingface/transformers.js) 🤗🤗🤗
+- 🚀 2024/06/02: OnnxSlim is merged into [ultralytics](https://github.com/ultralytics/ultralytics) ❤️❤️❤️
+- 🚀 2024/04/30: Rank 1st in the [AICAS 2024 LLM inference optimization challenge](https://tianchi.aliyun.com/competition/entrance/532170/customize440) held by Arm and T-head
+- 🚀 2024/01/25: OnnxSlim is merged to [mnn-llm](https://github.com/wangzhaode/mnn-llm), performance increased by 5%
+
+# Benchmark
+
+![Image](https://github.com/user-attachments/assets/fefc79f1-5d8d-486b-935a-a088846b3900)
+
+# Installation
+
+## Using Prebuilt
+
+```bash
+pip install onnxslim
+```
+
+## Install From Source
+
+```bash
+pip install git+https://github.com/inisis/OnnxSlim@main
+```
+
+## Install From Local
+
+```bash
+git clone https://github.com/inisis/OnnxSlim && cd OnnxSlim/
+pip install .
+```
+
+# How to use
+
+## Bash
+
+```bash
+onnxslim your_onnx_model slimmed_onnx_model
+```
+
+<div align=left><img src="https://raw.githubusercontent.com/inisis/onnxslim/main/images/onnxslim.gif"></div>
+
+## Inscript
+
+```inscript
+import onnx
+import onnxslim
+
+model = onnx.load("model.onnx")
+slimmed_model = onnxslim.slim(model)
+onnx.save(slimmed_model, "slimmed_model.onnx")
+```
+
+For more usage, see onnxslim -h or refer to our [examples](./examples)
+
+# Projects using OnnxSlim
+
+- <img src="https://avatars.githubusercontent.com/u/131524?s=48&v=4" width="22" height="22"/>[Mozilla/smart_autofill](https://github.com/mozilla/smart_autofill)
+- <img src="https://avatars.githubusercontent.com/u/1961952?s=48&v=4" width="22" height="22"/>[alibaba/MNN](https://github.com/alibaba/MNN)
+- <img src="https://avatars.githubusercontent.com/u/23534030?s=48&v=4" width="22" height="22"/>[PaddlePaddle/PaddleOCR](https://github.com/PaddlePaddle/PaddleOCR)
+- <img src="https://avatars.githubusercontent.com/u/25720743?s=48&v=4" width="22" height="22"/>[huggingface/transformers.js](https://github.com/huggingface/transformers.js)
+- <img src="https://avatars.githubusercontent.com/u/25720743?s=48&v=4" width="22" height="22"/>[huggingface/optimum](https://github.com/huggingface/optimum)
+- <img src="https://avatars.githubusercontent.com/u/86091366?s=48&v=4" width="22" height="22"/>[THU-MIG/yolov10](https://github.com/THU-MIG/yolov10)
+- <img src="https://avatars.githubusercontent.com/u/26833451?s=48&v=4" width="22" height="22"/>[ultralytics/ultralytics](https://github.com/ultralytics/ultralytics)
+- <img src="https://avatars.githubusercontent.com/u/109945100?s=48&v=4" width="22" height="22"/>[ModelScope/FunASR](https://github.com/modelscope/FunASR)
+- <img src="https://avatars.githubusercontent.com/u/1961952?s=48&v=4" width="22" height="22"/>[alibaba/MNN-LLM](https://github.com/wangzhaode/mnn-llm)
+- <img src="https://avatars.githubusercontent.com/u/126587470?s=48&v=4" width="22" height="22"/>[deepghs/imgutils](https://github.com/deepghs/imgutils)
+- <img src="https://avatars.githubusercontent.com/u/48153283?s=48&v=4" width="22" height="22"/>[sunsmarterjie/yolov12](https://github.com/sunsmarterjie/yolov12)
+- <img src="https://avatars.githubusercontent.com/u/147458884?s=48&v=4" width="22" height="22"/>[nndeploy/nndeploy](https://github.com/nndeploy/nndeploy)
+- <img src="https://avatars.githubusercontent.com/u/111754012?s=48&v=4" width="22" height="22"/>[CVCUDA/CV-CUDA](https://github.com/CVCUDA/CV-CUDA)
+
+# References
+
+> - [onnx-graphsurgeon](https://github.com/NVIDIA/TensorRT/tree/main/tools/onnx-graphsurgeon)
+> - [Polygraphy](https://github.com/NVIDIA/TensorRT/tree/main/tools/Polygraphy/polygraphy)
+> - [onnx-simplifier](https://github.com/daquexian/onnx-simplifier)
+> - [tabulate](https://github.com/astanin/python-tabulate)
+> - [onnxruntime](https://github.com/microsoft/onnxruntime)
+
+# Contact
+
+Discord: https://discord.gg/nRw2Fd3VUS QQ Group: `873569894`

onnxslim-0.1.75/VERSION

@@ -0,0 +1 @@
+0.1.75

{onnxslim-0.1.38 → onnxslim-0.1.75}/onnxslim/__init__.py

@@ -3,7 +3,6 @@ import warnings
 
 from onnxslim.cli import slim
 from onnxslim.core.pattern.registry import (
-    DEFAULT_FUSION_PATTERNS,
     register_fusion_pattern,
 )
 from onnxslim.version import __version__

{onnxslim-0.1.38 → onnxslim-0.1.75}/onnxslim/argparser.py

@@ -1,9 +1,34 @@
 import argparse
 import dataclasses
+from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from dataclasses import dataclass, field
-from typing import List, Optional, Type
-
-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:
+        return next((t for t in get_args(arg_type) if t is not type(None)), str)
+    return arg_type
 
 
 @dataclass
@@ -31,14 +56,24 @@ class OptimizationArguments:
     """
 
     no_shape_infer: bool = field(default=False, metadata={"help": "whether to disable shape_infer, default false."})
-    no_constant_folding: bool = field(
-        default=False, metadata={"help": "whether to disable constant_folding, 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(onnxslim.DEFAULT_FUSION_PATTERNS.keys()),
+            "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",
         },
     )
 
@@ -109,8 +144,11 @@ class CheckerArguments:
     verbose: bool = field(default=False, metadata={"help": "verbose mode, default False."})
 
 
-class ArgumentParser:
-    def __init__(self, *argument_dataclasses: Type):
+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",
@@ -120,13 +158,19 @@ class ArgumentParser:
 
     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 = field_def.type
+                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 arg_type == Optional[List[str]] else None
+                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('_', '-')}",
@@ -137,7 +181,7 @@ class ArgumentParser:
                 else:
                     self.parser.add_argument(
                         f"--{field_name.replace('_', '-')}",
-                        type=arg_type if arg_type != Optional[List[str]] else str,
+                        type=arg_type,
                         default=default_value,
                         nargs=nargs,
                         choices=choices,
@@ -147,9 +191,17 @@ class 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`
+        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)
 

{onnxslim-0.1.38 → onnxslim-0.1.75}/onnxslim/cli/_main.py

@@ -1,14 +1,17 @@
-from typing import Union
+from __future__ import annotations
 
 import onnx
 
+from onnxslim.argparser import OnnxSlimKwargs
 
-def slim(model: Union[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
 
     from onnxslim.core import (
+        OptimizationSettings,
         convert_data_format,
         freeze,
         input_modification,
@@ -18,6 +21,7 @@ def slim(model: Union[str, onnx.ModelProto], *args, **kwargs):
         shape_infer,
     )
     from onnxslim.utils import (
+        TensorInfo,
         check_onnx,
         check_point,
         check_result,
@@ -27,6 +31,7 @@ def slim(model: Union[str, onnx.ModelProto], *args, **kwargs):
         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)
@@ -35,10 +40,12 @@ def slim(model: Union[str, onnx.ModelProto], *args, **kwargs):
     inputs = kwargs.get("inputs", None)
     outputs = kwargs.get("outputs", None)
     no_shape_infer = kwargs.get("no_shape_infer", False)
-    no_constant_folding = kwargs.get("no_constant_folding", False)
+    skip_optimizations = kwargs.get("skip_optimizations", None)
     dtype = kwargs.get("dtype", None)
     skip_fusion_patterns = kwargs.get("skip_fusion_patterns", None)
-    inspect = kwargs.get("inspect", False)
+    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)
@@ -48,24 +55,37 @@ def slim(model: Union[str, onnx.ModelProto], *args, **kwargs):
 
     MAX_ITER = int(os.getenv("ONNXSLIM_MAX_ITER")) if os.getenv("ONNXSLIM_MAX_ITER") else 10
 
-    if isinstance(model, str):
-        model_name = Path(model).name
-        model = onnx.load(model)
-    else:
-        model_name = "OnnxModel"
+    start_time = time.time()
 
-    freeze(model)
+    def get_info(model, inspect=False):
+        if isinstance(model, str):
+            model_name = Path(model).name
+            model = onnx.load(model)
+        else:
+            model_name = "OnnxModel"
 
-    start_time = time.time()
+        freeze(model)
+
+        if not inspect:
+            return model_name, model
+
+        model_info = summarize_model(model, model_name)
+
+        return model_info
 
-    if output_model or inspect:
-        float_info = summarize_model(model)
+    if isinstance(model, list):
+        model_info_list = [get_info(m, inspect=True) for m in model]
 
-    if inspect:
-        print_model_info_as_table(model_name, [float_info])
         if dump_to_disk:
-            dump_model_info_to_disk(model_name, float_info)
-        return None
+            [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)
@@ -79,14 +99,17 @@ def slim(model: Union[str, onnx.ModelProto], *args, **kwargs):
     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)
 
-    if not no_constant_folding:
+    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)
+            model = optimize(model, skip_fusion_patterns, size_threshold)
             if not no_shape_infer:
                 model = shape_infer(model)
             graph = check_point(model)
@@ -101,21 +124,23 @@ def slim(model: Union[str, onnx.ModelProto], *args, **kwargs):
     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)
-        check_result(raw_onnx_output, slimmed_onnx_output)
+        if not check_result(raw_onnx_output, slimmed_onnx_output):
+            return None
 
     if not output_model:
         return model
 
-    slimmed_info = summarize_model(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(
-        model_name,
-        [float_info, slimmed_info],
+        [original_info, slimmed_info],
         elapsed_time,
     )
 
@@ -123,26 +148,26 @@ def slim(model: Union[str, onnx.ModelProto], *args, **kwargs):
 def main():
     """Entry point for the OnnxSlim toolkit, processes command-line arguments and passes them to the slim function."""
     from onnxslim.argparser import (
-        ArgumentParser,
         CheckerArguments,
         ModelArguments,
         ModificationArguments,
+        OnnxSlimArgumentParser,
         OptimizationArguments,
     )
 
-    argument_parser = ArgumentParser(ModelArguments, OptimizationArguments, ModificationArguments, CheckerArguments)
+    argument_parser = OnnxSlimArgumentParser(
+        ModelArguments, OptimizationArguments, ModificationArguments, CheckerArguments
+    )
     model_args, optimization_args, modification_args, checker_args = argument_parser.parse_args_into_dataclasses()
 
-    if checker_args.inspect and model_args.output_model:
-        argument_parser.error("--inspect and output_model are mutually exclusive")
-
     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 or optimization_args.no_constant_folding:
-        from onnxslim.utils import check_onnx_compatibility
+    if not optimization_args.no_shape_infer:
+        from onnxslim.utils import check_onnx_compatibility, is_onnxruntime_available
 
-        check_onnx_compatibility()
+        if is_onnxruntime_available():
+            check_onnx_compatibility()
 
     slim(
         model_args.input_model,