Skip to content

Commit 261858c

Browse files
Replace ONNX simplification package from onnxsim to onnxslim (#478)
## What does this PR do? **Type of change:** Add onnxslim support **Overview:** [Onnxslim](https://github.com/inisis/OnnxSlim) is under active development and committed to long-time-support, it's easy to use and is dependent on very few packages. ## Usage ```python $ python -m modelopt.onnx.quantization --onnx_path=$MODEL_NAME.onnx --simplify ``` ## Testing <!-- Mention how have you tested your change if applicable. --> ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No <!--- If No, explain why. --> - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No <!--- Only for new features, API changes, critical bug fixes or bw breaking changes. --> ## Additional Information <!-- E.g. related issue. --> --------- Signed-off-by: inisis <[email protected]> Co-authored-by: Keval Morabia <[email protected]>
1 parent 768ee6a commit 261858c

File tree

6 files changed

+9
-26
lines changed

6 files changed

+9
-26
lines changed

.gitlab/tests.yml

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,6 @@ unit:
1919
TRANSFORMERS: latest
2020
image: python:3.$PYTHON
2121
before_script:
22-
# Install cmake to build onnxsim from sdists for Python 3.12 until http://github.com/daquexian/onnx-simplifier/pull/353
23-
- if [ "$PYTHON" = "12" ]; then apt-get update && apt-get install -y cmake; fi
2422
- pip install tox
2523
script:
2624
- tox -e py3$PYTHON-torch$TORCH-tf_$TRANSFORMERS-unit

CHANGELOG.rst

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ Model Optimizer Changelog (Linux)
2828
**Misc**
2929

3030
- Bump minimum recommended transformers version to 4.53.
31+
- Replace ONNX simplification package from ``onnxsim`` to ``onnxslim``.
3132

3233
0.39 (2025-11-11)
3334
^^^^^^^^^^^^^^^^^

modelopt/onnx/quantization/quantize.py

Lines changed: 3 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@
4141
import onnx
4242
import onnx.onnx_cpp2py_export.checker as C
4343
import onnx_graphsurgeon as gs
44+
import onnxslim
4445

4546
from modelopt.onnx.logging_config import configure_logging, logger
4647
from modelopt.onnx.op_types import is_data_dependent_shape_op
@@ -133,16 +134,8 @@ def _preprocess_onnx(
133134
if simplify:
134135
logger.info("Attempting to simplify model")
135136
try:
136-
import onnxsim
137-
except ModuleNotFoundError as e:
138-
logger.warning(
139-
"onnxsim is not installed. Please install it with 'pip install onnxsim'."
140-
)
141-
raise e
142-
143-
try:
144-
model_simp, check = onnxsim.simplify(onnx_model)
145-
if check:
137+
model_simp = onnxslim.slim(onnx_model, skip_fusion_patterns=["FusionGemm"])
138+
if model_simp:
146139
onnx_model = model_simp
147140
onnx_path = os.path.join(output_dir, f"{model_name}_simp.onnx")
148141
save_onnx(onnx_model, onnx_path, use_external_data_format)

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@
5252
"onnxruntime-gpu~=1.22.0 ; platform_machine != 'aarch64' and platform_system != 'Darwin' and platform_system != 'Windows'", # noqa: E501
5353
"onnxruntime-directml==1.20.0; platform_system == 'Windows'",
5454
"onnxscript", # For autocast opset conversion and test_onnx_dynamo_export unit test
55-
"onnxsim ; python_version < '3.12' and platform_machine != 'aarch64'",
55+
"onnxslim>=0.1.76",
5656
"polygraphy>=0.49.22",
5757
],
5858
"hf": [

tests/gpu/onnx/test_simplify.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -57,14 +57,14 @@ def test_onnx_simplification(tmp_path):
5757
assert os.path.isfile(output_onnx_path), "Quantized ONNX was not found!"
5858

5959
# Load the simplified model and check that the model doesn't contain Identity nodes,
60-
# only 3 layers (Conv->BN->Relu).
60+
# only 2 layers (Conv->Relu).
6161
graph = gs.import_onnx(onnx.load(simplified_onnx_path))
6262
identity_nodes = [n for n in graph.nodes if n.op == "Identity"]
6363
assert not identity_nodes, "Simplified ONNX model contains Identity nodes but it shouldn't."
64-
assert len(graph.nodes) == 3, (
65-
f"Number of nodes doesn't match the expected: {len(graph.nodes)} vs 3."
64+
assert len(graph.nodes) == 2, (
65+
f"Number of nodes doesn't match the expected: {len(graph.nodes)} vs 2."
6666
)
67-
assert all(n.op in ["Conv", "BatchNormalization", "Relu"] for n in graph.nodes), (
67+
assert all(n.op in ["Conv", "Relu"] for n in graph.nodes), (
6868
"Graph contains more ops than expected."
6969
)
7070

tox.ini

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,6 @@ deps =
1818
torch28: torchvision~=0.23.0
1919
torch29: torchvision~=0.24.0
2020

21-
# Build onnxsim from sdists for Python 3.12 until http://github.com/daquexian/onnx-simplifier/pull/353
22-
py312: onnxsim
23-
2421
# Install megatron-core for special unit tests
2522
megatron-core
2623

@@ -42,9 +39,6 @@ deps =
4239
# Make sure torch 2.9 is used
4340
torchvision~=0.24.0
4441

45-
# Build onnxsim from sdists for Python 3.12 until http://github.com/daquexian/onnx-simplifier/pull/353
46-
py312: onnxsim
47-
4842
# ONNX unit tests heavily rely on torch / torchvision
4943
onnx: .[onnx,dev-test]
5044
onnx: torchvision
@@ -80,9 +74,6 @@ commands_pre =
8074
# Install Eagle-3 test dependencies
8175
pip install tiktoken blobfile sentencepiece
8276

83-
# Build onnxsim from sdists for Python 3.12 until http://github.com/daquexian/onnx-simplifier/pull/353
84-
py312: pip install onnxsim
85-
8677
# NOTE: User is expected to have correct torch-cuda version pre-installed if using --current-env
8778
# to avoid possible CUDA version mismatch
8879
pip install -e .[all,dev-test]

0 commit comments

Comments
 (0)