Docstrings

.github/workflows/format.yml

@@ -0,0 +1,29 @@
+# Ultralytics 🚀 - AGPL-3.0 license
+# Ultralytics Actions https://github.com/ultralytics/actions
+# This workflow automatically formats code and documentation in PRs to official Ultralytics standards
+
+name: Ultralytics Actions
+
+on:
+  push:
+    branches: [main]
+  pull_request_target:
+    branches: [main]
+    types: [opened, closed, synchronize]
+
+jobs:
+  format:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Run Ultralytics Formatting
+        uses: ultralytics/actions@main
+        with:
+          token: ${{ secrets.GITHUB_TOKEN }} # automatically generated, do not modify
+          python: true # format Python code and docstrings
+          markdown: true # format Markdown
+          prettier: true # format YAML
+          spelling: true # check spelling
+          links: false # check broken links
+          # summary: true # print PR summary with GPT4 (requires 'openai_api_key' or 'openai_azure_api_key' and 'openai_azure_endpoint')
+          # openai_azure_api_key: ${{ secrets.OPENAI_AZURE_API_KEY }}
+          # openai_azure_endpoint: ${{ secrets.OPENAI_AZURE_ENDPOINT }}

README.md

@@ -12,21 +12,24 @@
 OnnxSlim can help you slim your onnx model, with less operators, but same accuracy, better inference speed.
 
 - 🚀 OnnxSlim is merged to [mnn-llm](https://github.com/wangzhaode/mnn-llm), performance increased by 5%
-- 🚀 Rank 1st in the [AICAS 2024 LLM inference optimiztion challenge](https://tianchi.aliyun.com/competition/entrance/532170/customize440) held by Arm and T-head
-
+- 🚀 Rank 1st in the [AICAS 2024 LLM inference optimization challenge](https://tianchi.aliyun.com/competition/entrance/532170/customize440) held by Arm and T-head
 
 # Installation
+
 ## Using Prebuilt
+
 ```bash
 pip install onnxslim
 ```
+
 ## Build From Source
+
 ```
 pip install .
 ```
 
-
 # How to use
+
 ```
 onnxslim your_onnx_model slimmed_onnx_model
 ```
@@ -36,12 +39,14 @@ onnxslim your_onnx_model slimmed_onnx_model
 For more usage, see onnxslim -h or refer to our [examples](./examples)
 
 # 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)
+
+> - [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  
+
+Discord: https://discord.gg/nRw2Fd3VUS\
 QQ Group: 873569894

docs/conf.py

@@ -38,4 +38,5 @@
 
 
 def setup(app):
+    """Configure Sphinx application to include custom CSS from 'style.css'."""
     app.add_css_file("style.css")

examples/common_subexpression_elimination/README.md

@@ -1,16 +1,20 @@
 # Common SubExpression Elimination
 
 ## Introduction
+
 Common Subexpression Elimination (CSE) is a powerful optimization technique commonly employed in compilers to improve the efficiency of code execution. It targets redundant computations within a program by identifying and removing duplicate expressions, thus reducing both computational overhead and memory usage. By eliminating redundant computations, CSE enhances the overall performance of slimmed onnx model.
 
 ## How CSE Works
+
 In many programs, certain expressions are computed multiple times within a given scope, even though their results remain constant across these computations. Common subexpressions refer to these redundant expressions. CSE identifies such common subexpressions and replaces subsequent occurrences with references to the original computation result. This process effectively reduces the number of computations required during program execution.
 
 For example, consider the following code snippet:
+
 ```
 int a = b + c;
 int x = b + c;
 ```
+
 In this code, b + c is a common subexpression computed twice. With CSE, the redundant computation of b + c would be eliminated, and both occurrences of x would directly reference the computation result of a.
 
 ## Running the example
@@ -31,7 +35,6 @@ After onnxslim, the output will look like this:
 
 ![../../image/after_cse.png](../../images/after_cse.png)
 
-
 and the summary is as follow:
 
-![../../image/cse.png](../../images/cse.png)
+![../../image/cse.png](../../images/cse.png)

examples/common_subexpression_elimination/cse_demo.py

@@ -5,11 +5,13 @@
 
 class Model(torch.nn.Module):
     def __init__(self):
+        """Initializes the Model class with a single LayerNorm layer of embedding dimension 10."""
         super(Model, self).__init__()
         embedding_dim = 10
         self.layer_norm = nn.LayerNorm(embedding_dim)
 
     def forward(self, x):
+        """Applies LayerNorm to the input tensor and adds it to an independently computed LayerNorm of the same tensor."""
         return self.layer_norm(x) + F.layer_norm(x, [10])
 
 

examples/input_shape_modification/README.md

@@ -1,11 +1,13 @@
 # Input Shape Modification
 
 ## Introduction
+
 OnnxSlim includes an exploration of essential input shape modification techniques for ONNX models.
 
 This concise guide unveils techniques for seamlessly adjusting input tensor dimensions, ensuring optimal compatibility and performance within the dynamic landscape of neural network architectures.
 
 ## Running the example
+
 Change the input model by running:
 
 ```bash
@@ -14,4 +16,4 @@ onnxslim UNetModel-fp16.onnx slim.onnx --input_shapes cc:1,1,768
 
 The slimmed model will look like this:
 
-![../../image/input_shape_modification.jpg](../../images/input_shape_modification.jpg)
+![../../image/input_shape_modification.jpg](../../images/input_shape_modification.jpg)

examples/model_inspect/README.md

@@ -1,9 +1,11 @@
 # Model Inspect
 
 ## Introduction
+
 Dive deep into the intricacies of your ONNX model using the powerful --inspect argument with OnnxSlim. This feature provides detailed insights into various aspects of your model, including input and output details, operator information, opset version, and more.
 
 ## Running the example
+
 Unveil the secrets of your ONNX model by executing the following command:
 
 ```bash
@@ -12,4 +14,4 @@ onnxslim --inspect UNetModel-fp16.onnx
 
 The output will look like this:
 
-![../../image/model_inspect.jpg](../../images/model_inspect.jpg)
+![../../image/model_inspect.jpg](../../images/model_inspect.jpg)

examples/output_modification/README.md

@@ -1,11 +1,13 @@
 # Output Modification
 
 ## Introduction
+
 OnnxSlim provides capabilities for modifying the output specifications of ONNX models.
 
 This section explores techniques to customize the outputs, allowing for flexibility in handling diverse model requirements.
 
 ## Running the example
+
 Change the output of one model by running:
 
 ```bash
@@ -14,4 +16,4 @@ onnxslim yolov5m.onnx slim.onnx --outputs 591 739 443
 
 The slimmed model will look like this:
 
-![../../image/output_modification.jpg](../../images/output_modification.jpg)
+![../../image/output_modification.jpg](../../images/output_modification.jpg)

onnxslim/__init__.py

@@ -5,10 +5,7 @@
 from .core.optimizer import DEFAULT_FUSION_PATTERNS
 from .version import __version__
 
-
-if os.path.dirname(os.path.realpath(__file__)) == os.path.join(
-    os.path.realpath(os.getcwd()), "onnxslim"
-):
+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/cli/_main.py

@@ -1,6 +1,7 @@
 from typing import Union
 
 import onnx
+
 from onnxslim.utils.utils import logger
 
 
@@ -83,11 +84,7 @@ def slim(
 
     init_logging(verbose)
 
-    MAX_ITER = (
-        10
-        if not os.getenv("ONNXSLIM_MAX_ITER")
-        else int(os.getenv("ONNXSLIM_MAX_ITER"))
-    )
+    MAX_ITER = 10 if not os.getenv("ONNXSLIM_MAX_ITER") else int(os.getenv("ONNXSLIM_MAX_ITER"))
 
     if isinstance(model, str):
         model_name = Path(model).name
@@ -166,6 +163,7 @@ def slim(
 
 
 def main():
+    """Entry point for the OnnxSlim toolkit, processes command-line arguments and passes them to the slim function."""
     import argparse
 
     import onnxslim
@@ -175,14 +173,10 @@ def main():
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
     parser.add_argument("input_model", help="input onnx model")
-    parser.add_argument(
-        "output_model", nargs="?", default=None, help="output onnx model"
-    )
+    parser.add_argument("output_model", nargs="?", default=None, help="output onnx model")
 
     parser.add_argument("--model_check", action="store_true", help="enable model check")
-    parser.add_argument(
-        "-v", "--version", action="version", version=onnxslim.__version__
-    )
+    parser.add_argument("-v", "--version", action="version", version=onnxslim.__version__)
 
     # Input Shape Modification
     parser.add_argument(
@@ -259,9 +253,7 @@ def main():
     )
 
     # Verbose
-    parser.add_argument(
-        "--verbose", action="store_true", help="verbose mode, default False."
-    )
+    parser.add_argument("--verbose", action="store_true", help="verbose mode, default False.")
 
     args, unknown = parser.parse_known_args()