SenseVoice模型导出教程：ONNX格式转换与优化步骤

【免费下载链接】SenseVoice Multilingual Voice Understanding Model 项目地址: https://gitcode.com/gh_mirrors/se/SenseVoice

引言

你是否在部署SenseVoice模型时遇到过推理速度慢、模型体积大的问题？本文将详细介绍如何将SenseVoice模型导出为ONNX（Open Neural Network Exchange，开放神经网络交换）格式，并进行优化，以解决这些问题。通过本文的学习，你将能够：

• 理解ONNX格式的优势及其在SenseVoice模型部署中的作用
• 掌握SenseVoice模型导出为ONNX格式的完整流程
• 学会使用量化技术优化导出的ONNX模型
• 了解导出后模型的验证方法和常见问题解决策略

ONNX格式简介

ONNX是一种开放的文件格式，用于表示机器学习模型。它允许不同的深度学习框架（如PyTorch、TensorFlow等）之间进行模型互操作，同时也为模型部署提供了便利。对于SenseVoice模型而言，导出为ONNX格式具有以下优势：

• 跨平台兼容性：ONNX模型可以在多种平台和设备上运行，包括CPU、GPU和专用硬件加速计算单元。
• 优化部署：ONNX Runtime等推理引擎针对ONNX模型进行了优化，可以显著提高推理速度。
• 减小模型体积：通过量化等技术，可以进一步减小ONNX模型的体积，便于在资源受限的环境中部署。

准备工作

环境要求

在开始导出过程之前，请确保你的环境满足以下要求：

软件/库	版本要求	说明
Python	3.8+	推荐使用3.8或更高版本
PyTorch	与训练时一致	确保与训练SenseVoice模型时使用的PyTorch版本兼容
ONNX	1.10.0+	ONNX格式支持库
ONNX Runtime	1.10.0+	ONNX模型推理引擎
funasr	>=1.1.3	SenseVoice模型依赖的语音识别工具包
funasr-onnx	0.4.0+	用于ONNX格式的SenseVoice模型推理

安装依赖

可以使用以下命令安装所需的依赖包：

pip install -r requirements.txt
pip install onnx onnxruntime funasr-onnx>=0.4.0

准备模型文件

确保你已经下载了SenseVoice模型文件。默认情况下，模型将从ModelScope或Hugging Face Hub下载。如果你已经有本地模型文件，可以直接使用本地路径。

导出流程

SenseVoice模型导出为ONNX格式的流程如下：

1. 加载预训练模型

首先，我们需要加载预训练的SenseVoice模型。以下是加载模型的代码示例：

from model import SenseVoiceSmall

model_dir = "iic/SenseVoiceSmall"  # 模型名称或本地路径
model, kwargs = SenseVoiceSmall.from_pretrained(model=model_dir, device="cuda:0")

2. 重建模型用于导出

SenseVoice模型提供了专门的导出方法，用于重建模型以适应ONNX格式：

rebuilt_model = model.export(type="onnx", optimize=False)

这里的optimize=False表示暂时不进行优化，我们将在后续步骤中单独处理优化。

3. 导出ONNX模型

接下来，我们使用PyTorch的ONNX导出功能将模型导出为ONNX格式：

import os
import torch
from utils import export_utils

model_path = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param")))
model_file = os.path.join(model_path, "model.onnx")

if not os.path.exists(model_file):
    with torch.no_grad():
        del kwargs['model']
        export_dir = export_utils.export(model=rebuilt_model, **kwargs)
        print(f"Export model onnx to {model_file}")

export_utils.export函数内部处理了ONNX导出的详细过程，包括设置输入输出名称、动态轴等。

4. 优化ONNX模型

为了减小模型体积并提高推理速度，我们可以对导出的ONNX模型进行优化。SenseVoice提供了内置的优化功能：

optimize = True  # 设置为True启用优化

if optimize:
    model_file = os.path.join(model_path, "model_optimized.onnx")
    # 优化过程会在export_utils.export中处理
    # 只需确保在调用model.export时设置optimize=True
    rebuilt_model = model.export(type="onnx", optimize=True)
    # 重新导出优化模型
    if not os.path.exists(model_file):
        with torch.no_grad():
            del kwargs['model']
            export_dir = export_utils.export(model=rebuilt_model, **kwargs)
            print(f"Export optimized model onnx to {model_file}")

优化过程的内部实现如下：

from onnxruntime.quantization import QuantType, quantize_dynamic
import onnx

def optimize_onnx_model(model_path, optimized_model_path):
    if not os.path.exists(optimized_model_path):
        onnx_model = onnx.load(model_path)
        nodes = [n.name for n in onnx_model.graph.node]
        nodes_to_exclude = [
            m for m in nodes if "output" in m or "bias_encoder" in m or "bias_decoder" in m
        ]
        quantize_dynamic(
            model_input=model_path,
            model_output=optimized_model_path,
            op_types_to_quantize=["MatMul"],
            per_channel=True,
            reduce_range=False,
            weight_type=QuantType.QUInt8,
            nodes_to_exclude=nodes_to_exclude,
        )

这段代码展示了如何使用ONNX Runtime提供的优化工具对模型进行动态优化，将权重从32位浮点数转换为8位无符号整数，同时排除了一些不适合优化的节点。

模型验证

导出完成后，我们需要验证导出的ONNX模型是否正常工作。可以使用以下代码进行验证：

from funasr_onnx import SenseVoiceSmall
from funasr_onnx.utils.postprocess_utils import rich_transcription_postprocess

# 加载ONNX模型
model = SenseVoiceSmall(model_path, batch_size=10, optimize=optimize)

# 准备测试音频
wav_path = "path/to/your/test/audio.wav"  # 替换为你的测试音频路径

# 推理
res = model(wav_path, language="auto", use_itn=True)
print([rich_transcription_postprocess(i) for i in res])

如果一切正常，你应该能看到模型对测试音频的识别结果。

优化策略

除了优化之外，还有其他一些策略可以进一步优化导出的ONNX模型：

1. 调整OPSET版本

ONNX定义了不同的操作集（OPSET）版本，较高的版本可能支持更多优化。在导出时可以指定OPSET版本：

export_utils.export(model=rebuilt_model, opset_version=14, **kwargs)

SenseVoice默认使用OPSET 14，这是一个比较稳定且广泛支持的版本。

2. 模型裁剪

如果你的应用只需要SenseVoice的部分功能（如仅需要语音识别，不需要情感识别），可以考虑裁剪模型，只保留需要的部分。这需要对模型结构有深入了解，建议谨慎操作。

3. 使用ONNX Runtime优化

ONNX Runtime提供了多种优化选项，可以在推理时启用：

import onnxruntime as ort

options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

session = ort.InferenceSession(model_file, options)

这将启用ONNX Runtime的所有可用优化。

常见问题解决

导出失败

如果导出过程中出现错误，可能的原因和解决方法：

1. PyTorch版本不兼容：确保使用与训练时相同或兼容的PyTorch版本。
2. 模型路径错误：检查模型路径是否正确，确保模型文件存在。
3. CUDA内存不足：尝试使用更小的批次大小或在CPU上导出（虽然速度较慢）。

推理结果不一致

如果ONNX模型的推理结果与原始PyTorch模型不一致：

1. 检查输入预处理：确保ONNX模型的输入预处理与PyTorch模型一致。
2. 验证动态轴设置：确保在导出时正确设置了动态轴，特别是对于可变长度的音频输入。
3. 尝试不同的OPSET版本：有时更换OPSET版本可以解决兼容性问题。

优化模型精度下降

如果优化后模型精度下降过多：

1. 选择性优化：只对部分层进行优化，保留关键层的浮点精度。
2. 使用更精细的优化策略：如感知优化（Perceptron Optimization），可能需要更多的校准数据。
3. 考虑混合精度优化：部分层使用16位浮点优化，平衡精度和性能。

部署示例

使用Python部署

导出并优化后的ONNX模型可以使用funasr-onnx库轻松部署：

from funasr_onnx import SenseVoiceSmall

# 加载模型
model = SenseVoiceSmall(model_dir, batch_size=10, optimize=True)

# 批量处理音频文件
audio_files = ["audio1.wav", "audio2.wav", ..., "audioN.wav"]
results = model(audio_files, language="auto", use_itn=True)

# 处理结果
for audio_path, result in zip(audio_files, results):
    print(f"{audio_path}: {result}")

服务化部署

可以使用FastAPI将导出的模型部署为Web服务：

from fastapi import FastAPI, File, UploadFile
from funasr_onnx import SenseVoiceSmall
import tempfile

app = FastAPI()
model = SenseVoiceSmall("iic/SenseVoiceSmall", optimize=True)

@app.post("/asr")
async def asr_endpoint(file: UploadFile = File(...)):
    with tempfile.NamedTemporaryFile(suffix=".wav") as temp:
        temp.write(await file.read())
        temp.flush()
        res = model(temp.name, language="auto", use_itn=True)
    return {"result": res[0]}

# 启动服务
# uvicorn main:app --host 0.0.0.0 --port 8000

总结与展望

本文详细介绍了将SenseVoice模型导出为ONNX格式的完整流程，包括准备工作、导出步骤、模型优化和部署示例。通过将SenseVoice导出为ONNX格式，我们可以充分利用ONNX生态系统的优势，实现跨平台部署和高性能推理。

未来，我们将继续优化导出流程，探索更多高级优化技术，如模型剪枝、知识蒸馏等，以进一步提高SenseVoice模型的部署效率。同时，我们也将提供更多平台和语言的部署示例，方便开发者在不同场景中使用SenseVoice模型。

如果你在导出或部署过程中遇到任何问题，欢迎在项目的GitHub仓库提交issue，我们的开发团队将尽快回复和解决。

附录：完整导出代码

以下是将SenseVoice模型导出为ONNX格式的完整代码：

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-

import os
import torch
from model import SenseVoiceSmall
from utils import export_utils
from utils.model_bin import SenseVoiceSmallONNX

def export_sensevoice_to_onnx(model_dir="iic/SenseVoiceSmall", optimize=False, output_dir=None):
    # 加载模型
    model, kwargs = SenseVoiceSmall.from_pretrained(model=model_dir, device="cuda:0")
    
    # 重建模型用于导出
    rebuilt_model = model.export(type="onnx", optimize=optimize)
    
    # 确定输出路径
    if output_dir:
        kwargs["output_dir"] = output_dir
    model_path = kwargs.get("output_dir", os.path.dirname(kwargs.get("init_param")))
    
    # 导出ONNX模型
    model_file = os.path.join(model_path, "model.onnx")
    if optimize:
        model_file = os.path.join(model_path, "model_optimized.onnx")
    
    if not os.path.exists(model_file):
        with torch.no_grad():
            del kwargs['model']
            export_dir = export_utils.export(model=rebuilt_model, **kwargs)
            print(f"Export model onnx to {model_file}")
    
    # 验证导出的模型
    model_bin = SenseVoiceSmallONNX(model_path)
    
    # 加载分词器
    try:
        from funasr.tokenizer.sentencepiece_tokenizer import SentencepiecesTokenizer
        tokenizer = SentencepiecesTokenizer(bpemodel=os.path.join(model_path, "chn_jpn_yue_eng_ko_spectok.bpe.model"))
    except:
        tokenizer = None
    
    # 简单推理测试
    try:
        from funasr.utils.postprocess_utils import rich_transcription_postprocess
        wav_or_scp = "/Users/shixian/Downloads/asr_example_hotword.wav"  # 替换为实际测试音频路径
        language_list = [0]
        textnorm_list = [15]
        res = model_bin(wav_or_scp, language_list, textnorm_list, tokenizer=tokenizer)
        print("模型导出成功，测试结果：")
        print([rich_transcription_postprocess(i) for i in res])
    except Exception as e:
        print(f"模型验证过程中出现错误: {e}")
        print("模型已导出，但建议检查导出结果")
    
    return model_file

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Export SenseVoice model to ONNX format")
    parser.add_argument("--model_dir", type=str, default="iic/SenseVoiceSmall", help="Model directory or name")
    parser.add_argument("--optimize", action="store_true", help="Whether to optimize the model")
    parser.add_argument("--output_dir", type=str, default=None, help="Output directory for ONNX model")
    args = parser.parse_args()
    
    export_sensevoice_to_onnx(model_dir=args.model_dir, optimize=args.optimize, output_dir=args.output_dir)

你可以将此代码保存为export_sensevoice_onnx.py，然后通过命令行运行：

# 导出基本ONNX模型
python export_sensevoice_onnx.py --model_dir "iic/SenseVoiceSmall" --output_dir "./onnx_model"

# 导出并优化ONNX模型
python export_sensevoice_onnx.py --model_dir "iic/SenseVoiceSmall" --output_dir "./onnx_model_optimized" --optimize

希望本教程能帮助你顺利将SenseVoice模型导出为ONNX格式，并在实际应用中取得良好的性能。如有任何问题或建议，欢迎在项目仓库中提出。

【免费下载链接】SenseVoice Multilingual Voice Understanding Model 项目地址: https://gitcode.com/gh_mirrors/se/SenseVoice