导出 Whisper 模型到 ONNX

前言

在语音识别领域，Whisper 模型因其出色的性能和灵活性备受关注。为了在更多平台和环境中部署 Whisper 模型，导出为 ONNX 格式是一个有效的途径。ONNX（Open Neural Network Exchange）是一个开放格式，支持不同的深度学习框架之间的模型互操作性。本指南将详细介绍如何将 Whisper 模型导出为 ONNX 格式，并提供测试模型的步骤。

本节描述了如何将 Whisper 模型导出为 ONNX 格式。

可用模型

请注意，我们已经将 Whisper 模型导出为 ONNX 格式，它们可以从以下 Huggingface 仓库中获取：

提示：你也可以从 此处 下载它们。

如果你想自己导出模型或了解模型的导出过程，请继续阅读下文。

导出为 ONNX

我们使用 export-onnx.py 来导出 Whisper 模型到 ONNX。

首先，让我们安装依赖并下载导出脚本：

pip install torch openai-whisper onnxruntime onnx
git clone https://github.com/k2-fsa/sherpa-onnx/
cd sherpa-onnx/scripts/whisper
python3 ./export-onnx.py --help

它将打印以下信息：

usage: export-onnx.py [-h] --model {tiny,tiny.en,base,base.en,small,small.en,medium,medium.en,large,large-v1,large-v2}

optional arguments:
  -h, --help            show this help message and exit
  --model {tiny,tiny.en,base,base.en,small,small.en,medium,medium.en,large,large-v1,large-v2}

要导出 tiny.en 模型，我们可以使用：

python3 ./export-onnx.py --model tiny.en

它将生成以下文件：

(py38) fangjuns-MacBook-Pro:whisper fangjun$ ls -lh tiny.en-*
-rw-r--r--  1 fangjun  staff   105M Aug  7 15:43 tiny.en-decoder.int8.onnx
-rw-r--r--  1 fangjun  staff   185M Aug  7 15:43 tiny.en-decoder.onnx
-rw-r--r--  1 fangjun  staff    12M Aug  7 15:43 tiny.en-encoder.int8.onnx
-rw-r--r--  1 fangjun  staff    36M Aug  7 15:43 tiny.en-encoder.onnx
-rw-r--r--  1 fangjun  staff   816K Aug  7 15:43 tiny.en-tokens.txt

tiny.en-encoder.onnx 是编码器模型，tiny.en-decoder.onnx 是解码器模型。

tiny.en-encoder.int8.onnx 是量化的编码器模型，tiny.en-decoder.int8.onnx 是量化的解码器模型。

tiny.en-tokens.txt 包含了令牌表，它将整数映射到令牌和反之。

要将导出的 ONNX 模型转换为 ONNX Runtime 格式，我们可以使用：

python3 -m onnxruntime.tools.convert_onnx_models_to_ort --optimization_style=Fixed ./

到目前为止，生成的文件如下：

(py38) fangjuns-MacBook-Pro:whisper fangjun$ ls -lh tiny.en-*
-rw-r--r--  1 fangjun  staff   105M Aug  7 15:43 tiny.en-decoder.int8.onnx
-rw-r--r--  1 fangjun  staff   185M Aug  7 15:43 tiny.en-decoder.onnx
-rw-r--r--  1 fangjun  staff    12M Aug  7 15:43 tiny.en-encoder.int8.onnx
-rw-r--r--  1 fangjun  staff    36M Aug  7 15:43 tiny.en-encoder.onnx
-rw-r--r--  1 fangjun  staff   816K Aug  7 15:43 tiny.en-tokens.txt

要检查导出的模型是否工作正常，我们可以使用 test.py。

我们使用 此测试音频。

pip install kaldi-native-fbank
wget https://huggingface.co/csukuangfj/sherpa-onnx-whisper-tiny.en/resolve/main/test_wavs/0.wav

python3 ./test.py \
  --encoder ./tiny.en-encoder.onnx \
  --decoder ./tiny.en-decoder.onnx \
  --tokens ./tiny.en-tokens.txt \
  ./0.wav

要测试 int8 量化模型，我们可以使用：

python3 ./test.py \
  --encoder ./tiny.en-encoder.int8.onnx \
  --decoder ./tiny.en-decoder.int8.onnx \
  --tokens ./tiny.en-tokens.txt \
  ./0.wav

希望这篇博客能帮助你顺利导出并测试 Whisper ONNX 模型。如果你有任何问题，请随时在评论区留言。