espnet_onnx

无需PyTorch的ESPnet!

一个用于轻松导出、量化和优化espnet模型到onnx格式的实用程序库。如果你已经有了导出的文件,就无需在你的机器上安装PyTorch或ESPnet!

Colab上的espnet_onnx演示

现在可以在Google Colab上使用演示笔记本了!

简单的ASR演示:
简单的TTS演示:

安装

可以使用pip安装espnet_onnx

pip install espnet_onnx

如果你想导出预训练模型,还需要额外安装torch>=1.11.0、espnet、espnet_model_zoo和onnx。 onnx==1.12.0可能会导致一些错误。如果你在推理或导出时遇到错误,请考虑降级onnx版本。

开发者安装指南

克隆此仓库。

git clone git@github.com:espnet/espnet_onnx.git

创建虚拟环境。

cd tools
make venv export

激活虚拟环境并根据需要安装torch。

. tools/venv/bin/activate

# 请参考PyTorch的官方安装指南。
pip install torch

克隆s3prl仓库并用pip安装。

cd tools
git clone https://github.com/s3prl/s3prl
cd s3prl
pip install .

为开发transducer模型安装warp_transducer。

cd tools
git clone --single-branch --branch espnet_v1.1 https://github.com/b-flo/warp-transducer.git
cd warp-transducer
mkdir build
# 请将WITH_OMP设置为ON或OFF。
cd build && cmake -DWITH_OMP="ON" .. && make
cd pytorch_binding && python3 -m pip install -e .

如果你想开发优化,还需要开发onnxruntime。请克隆onnxruntime仓库。
由于espnet==202308(v0.2.0发布时的最新版本)要求protobuf<=3.20.1,而最新的onnx要求protobuf>=3.20.2,你可能会在安装时遇到错误。在这种情况下,首先安装espnet==202308,将protobuf更新到3.20.3,然后安装其他库。

用法

导出模型

espnet_onnx可以导出espnet_model_zoo上发布的预训练模型。默认情况下,导出的文件将存储在${HOME}/.cache/espnet_onnx/<tag_name>中。

from espnet2.bin.asr_inference import Speech2Text
from espnet_onnx.export import ASRModelExport

m = ASRModelExport()

# 从espnet_model_zoo下载并从预训练模型导出
m.export_from_pretrained('<tag name>', quantize=True)

# 从训练好的模型导出
speech2text = Speech2Text(args)
m.export(speech2text, '<tag name>', quantize=True)

你可以从zip文件导出预训练模型。zip文件应包含meta.yaml。

from espnet_onnx.export import ASRModelExport

m = ASRModelExport()
m.export_from_zip(
  'path/to/the/zipfile',
  tag_name='tag_name_for_zipped_model',
  quantize=True
)

你可以为导出设置一些配置。每个模型的可用配置在详细说明中显示。

from espnet_onnx.export import ASRModelExport

m = ASRModelExport()
# 将最大序列长度设置为3000
m.set_export_config(max_seq_len=3000)
m.export_from_zip(
  'path/to/the/zipfile',
  tag_name='tag_name_for_zipped_model',
)

你可以使用optimize选项轻松优化你的模型。如果你想充分优化你的模型,你需要从这里安装onnxruntime的自定义版本。更多详情请阅读这个文档。

from espnet_onnx.export import ASRModelExport

m = ASRModelExport()
m.export_from_zip(
  'path/to/the/zipfile',
  tag_name='tag_name_for_zipped_model',
  optimize=True,
  quantize=True
)

你可以从命令行导出模型。

python -m espnet_onnx.export \
  --model_type asr \
  --input ${path_to_zip} \
  --tag transformer_lm \
  --apply_optimize \
  --apply_quantize

推理

对于推理,使用tag_name或model_dir来加载onnx文件。tag_name必须在tag_config.yaml中定义。

import librosa
from espnet_onnx import Speech2Text

speech2text = Speech2Text(tag_name='<tag name>')
# speech2text = Speech2Text(model_dir='path to the onnx directory')

y, sr = librosa.load('sample.wav', sr=16000)
nbest = speech2text(y)

对于流式ASR,你可以使用StreamingSpeech2Text类。语音长度应与StreamingSpeech2Text.hop_size相同。

from espnet_onnx import StreamingSpeech2Text

stream_asr = StreamingSpeech2Text(tag_name)

# 开始流式ASR
stream_asr.start()
while streaming:
  wav = <some code to get wav>
  assert len(wav) == stream_asr.hop_size
  stream_text = stream_asr(wav)[0][0]

# 你可以用end函数获取非流式ASR结果
nbest = stream_asr.end()

你也可以用simulate函数模拟你的wav文件的流式模型。将True作为第二个参数传递将显示如下代码的流式文本。

import librosa
from espnet_onnx import StreamingSpeech2Text
stream_asr = StreamingSpeech2Text(tag_name)
y, sr = librosa.load('path/to/wav', sr=16000)
nbest = stream_asr.simulate(y, True)
# 使用6个进程处理音频。
# 位置0的结果：
# 位置1的结果：
# 位置2的结果：this
# 位置3的结果：this is
# 位置4的结果：this is a
# 位置5的结果：this is a
print(nbest[0][0])
# 'this is a pen'

如果您安装了自定义版本的onnxruntime，您可以运行优化模型进行推理。您无需更改上述任何代码。如果模型已优化，espnet_onnx会自动加载优化版本。
您可以仅使用hubert模型作为前端。

from espnet_onnx.export import ASRModelExport

# 导出模型
tag_name = 'ESPnet预训练模型与hubert'
m = ASRModelExport()
m.export_from_pretrained(tag_name, optimize=True)

# 仅加载前端模型
from espnet_onnx.asr.frontend import Frontend
frontend = Frontend.get_frontend(tag_name)

# 在应用中使用模型
import librosa
y, sr = librosa.load('wav文件')
# y: (B, T)
# y_len: (B,)
feats = frontend(y[None,:], np.array([len(y)]))

如果您在环境中安装了torch，可以在训练中使用前端。

from espnet_onnx.asr.frontend import TorchFrontend
frontend = TorchFrontend.get_frontend(tag_name) # 加载预训练前端模型

# 在训练时使用模型
import librosa
y, sr = librosa.load('wav文件')

# 需要将数据放在GPU上，
# 并以元组形式指定输出形状
y = torch.Tensor(y).unsqueeze(0).to('cuda') # (1, wav_length)
output_shape = (batch_size, feat_length, feats_dims)
feats = frontend(y, y.size(1), output_shape)

文本到语音推理

您可以像导出ASR模型一样导出TTS模型。

from espnet2.bin.tts_inference import Text2Speech
from espnet_onnx.export import TTSModelExport

m = TTSModelExport()

# 从espnet_model_zoo下载并从预训练模型导出
m.export_from_pretrained('<tag name>', quantize=True)

# 从训练好的模型导出
text2speech = Text2Speech(args)
m.export(text2speech, '<tag name>', quantize=True)

您可以简单地使用Text2Speech类生成wav文件。

from espnet_onnx import Text2Speech

tag_name = 'kan-bayashi/ljspeech_vits'
text2speech = Text2Speech(tag_name, use_quantized=True)

text = 'Hello world!'
output_dict = text2speech(text) # 使用onnx模型进行推理。
wav = output_dict['wav']

如何在espnet_onnx上使用GPU

安装依赖。

首先，我们需要onnxruntime-gpu库，而不是onnxruntime。请按照这篇文章选择并安装正确版本的onnxruntime-gpu，具体取决于您的CUDA版本。

在GPU上进行推理

现在您可以使用GPU加速推理速度。您只需选择正确的提供程序，并将其提供给Speech2Text或StreamingSpeech2Text实例。有关提供程序的更多信息，请参阅这篇文章。

import librosa
from espnet_onnx import Speech2Text

PROVIDERS = ['CUDAExecutionProvider']
tag_name = 'some_tag_name'

speech2text = Speech2Text(
  tag_name,
  providers=PROVIDERS
)
y, sr = librosa.load('path/to/wav', sr=16000)
nbest = speech2text(y) # 在GPU上运行。

请注意，某些量化模型不支持GPU计算。如果使用量化模型出现错误，请尝试使用非量化模型。