use onnxsim

This commit is contained in:
Fangjun Kuang 2024-04-09 15:10:21 +08:00
parent 589eb626f9
commit e62d97eb30
15 changed files with 30 additions and 4 deletions

View File

@ -2,7 +2,7 @@
set -ex
python3 -m pip install onnxoptimizer
python3 -m pip install onnxoptimizer onnxsim
log() {
# This function is from espnet

View File

@ -50,6 +50,7 @@ RUN pip install --no-cache-dir \
numba \
numpy \
onnxoptimizer \
onnxsim \
onnx \
onnxmltools \
onnxruntime \

View File

@ -56,6 +56,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -56,6 +56,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -70,6 +70,7 @@ RUN pip uninstall -y tqdm && \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -57,6 +57,7 @@ RUN pip install --no-cache-dir \
onnxruntime \
onnxmltools \
onnxoptimizer \
onnxsim \
multi_quantization \
typeguard \
numpy \

View File

@ -42,6 +42,7 @@ import onnxoptimizer
import torch
import torch.nn as nn
from onnxruntime.quantization import QuantType, quantize_dynamic
from onnxsim import simplify
from scaling_converter import convert_scaled_to_non_scaled
from train import add_model_arguments, get_model, get_params
from zipformer import Zipformer2
@ -239,10 +240,22 @@ def optimize_model(filename):
# https://github.com/microsoft/onnxruntime/issues/1899#issuecomment-534806537
# and
# https://github.com/onnx/onnx/issues/582#issuecomment-937788108
passes = ["extract_constant_to_initializer", "eliminate_unused_initializer"]
# and
# https://github.com/onnx/optimizer/issues/110
# and
# https://qiita.com/Yossy_Hal/items/34f3b2aef2199baf7f5f
passes = ["eliminate_unused_initializer"]
onnx_model = onnx.load(filename)
optimized_model = onnxoptimizer.optimize(onnx_model, passes)
onnx.save(optimized_model, filename)
onnx_model = onnxoptimizer.optimize(onnx_model, passes)
model_simp, check = simplify(onnx_model)
if check:
logging.info("Simplified the model!")
onnx_model = model_simp
else:
logging.info("Failed to simplify the model!")
onnx.save(onnx_model, filename)
@torch.no_grad()