use onnxsim

.github/scripts/audioset/AT/run.sh

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

.github/scripts/docker/Dockerfile

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

docker/torch1.12.1-cuda11.3.dockerfile

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

docker/torch1.13.0-cuda11.6.dockerfile

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

docker/torch1.9.0-cuda10.2.dockerfile

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

docker/torch2.0.0-cuda11.7.dockerfile

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

docker/torch2.1.0-cuda11.8.dockerfile

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

docker/torch2.1.0-cuda12.1.dockerfile

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

docker/torch2.2.0-cuda11.8.dockerfile

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

docker/torch2.2.0-cuda12.1.dockerfile

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

docker/torch2.2.1-cuda11.8.dockerfile

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

docker/torch2.2.1-cuda12.1.dockerfile

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

docker/torch2.2.2-cuda11.8.dockerfile

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

docker/torch2.2.2-cuda12.1.dockerfile

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

egs/audioset/AT/zipformer/export-onnx.py

@@ -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_model = onnxoptimizer.optimize(onnx_model, passes)
-    onnx.save(optimized_model, filename)
+
+    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()