diff --git a/.github/scripts/audioset/AT/run.sh b/.github/scripts/audioset/AT/run.sh index f09969ec9..436c3c144 100755 --- a/.github/scripts/audioset/AT/run.sh +++ b/.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 diff --git a/.github/scripts/docker/Dockerfile b/.github/scripts/docker/Dockerfile index b918a10d8..15f49f826 100644 --- a/.github/scripts/docker/Dockerfile +++ b/.github/scripts/docker/Dockerfile @@ -50,6 +50,7 @@ RUN pip install --no-cache-dir \ numba \ numpy \ onnxoptimizer \ + onnxsim \ onnx \ onnxmltools \ onnxruntime \ diff --git a/docker/torch1.12.1-cuda11.3.dockerfile b/docker/torch1.12.1-cuda11.3.dockerfile index 1c83d4877..9815a8ec7 100644 --- a/docker/torch1.12.1-cuda11.3.dockerfile +++ b/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 \ diff --git a/docker/torch1.13.0-cuda11.6.dockerfile b/docker/torch1.13.0-cuda11.6.dockerfile index 450513e6a..d13d2a7cb 100644 --- a/docker/torch1.13.0-cuda11.6.dockerfile +++ b/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 \ diff --git a/docker/torch1.9.0-cuda10.2.dockerfile b/docker/torch1.9.0-cuda10.2.dockerfile index 36a533a31..5936fe06a 100644 --- a/docker/torch1.9.0-cuda10.2.dockerfile +++ b/docker/torch1.9.0-cuda10.2.dockerfile @@ -70,6 +70,7 @@ RUN pip uninstall -y tqdm && \ onnxruntime \ onnxmltools \ onnxoptimizer \ + onnxsim \ multi_quantization \ typeguard \ numpy \ diff --git a/docker/torch2.0.0-cuda11.7.dockerfile b/docker/torch2.0.0-cuda11.7.dockerfile index 75335b932..e2e27b55d 100644 --- a/docker/torch2.0.0-cuda11.7.dockerfile +++ b/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 \ diff --git a/docker/torch2.1.0-cuda11.8.dockerfile b/docker/torch2.1.0-cuda11.8.dockerfile index 1d6deae0d..de1e07e69 100644 --- a/docker/torch2.1.0-cuda11.8.dockerfile +++ b/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 \ diff --git a/docker/torch2.1.0-cuda12.1.dockerfile b/docker/torch2.1.0-cuda12.1.dockerfile index f6dcd2a2f..89303797a 100644 --- a/docker/torch2.1.0-cuda12.1.dockerfile +++ b/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 \ diff --git a/docker/torch2.2.0-cuda11.8.dockerfile b/docker/torch2.2.0-cuda11.8.dockerfile index 72c1add60..3364477a8 100644 --- a/docker/torch2.2.0-cuda11.8.dockerfile +++ b/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 \ diff --git a/docker/torch2.2.0-cuda12.1.dockerfile b/docker/torch2.2.0-cuda12.1.dockerfile index ca81afcf4..3cc41902d 100644 --- a/docker/torch2.2.0-cuda12.1.dockerfile +++ b/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 \ diff --git a/docker/torch2.2.1-cuda11.8.dockerfile b/docker/torch2.2.1-cuda11.8.dockerfile index 713fa1a09..76b785622 100644 --- a/docker/torch2.2.1-cuda11.8.dockerfile +++ b/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 \ diff --git a/docker/torch2.2.1-cuda12.1.dockerfile b/docker/torch2.2.1-cuda12.1.dockerfile index 5006a2180..55bdfa4d7 100644 --- a/docker/torch2.2.1-cuda12.1.dockerfile +++ b/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 \ diff --git a/docker/torch2.2.2-cuda11.8.dockerfile b/docker/torch2.2.2-cuda11.8.dockerfile index 78ac20bba..02de82c50 100644 --- a/docker/torch2.2.2-cuda11.8.dockerfile +++ b/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 \ diff --git a/docker/torch2.2.2-cuda12.1.dockerfile b/docker/torch2.2.2-cuda12.1.dockerfile index 655bcf933..44ad38b8e 100644 --- a/docker/torch2.2.2-cuda12.1.dockerfile +++ b/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 \ diff --git a/egs/audioset/AT/zipformer/export-onnx.py b/egs/audioset/AT/zipformer/export-onnx.py index e2afeeffd..9476dac62 100755 --- a/egs/audioset/AT/zipformer/export-onnx.py +++ b/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.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()