Learn.SpeakerAI/app/toONNX.py

27 lines
838 B
Python

# pip install torch onnx https://github.com/pyannote/pyannote-audio/archive/refs/heads/develop.zip
import torch
from pyannote.audio import Pipeline
import time
print(f"{time.strftime("%Y/%m/%d %H:%M:%S")}=>加载预训练模型")
model = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1",
use_auth_token="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").eval()
dummy_input = torch.zeros(2, 1, 160000)
torch.onnx.export(
model,
#输入参数
dummy_input,
#输出文件名
'speaker@3_1.onnx',
#应用常量折叠优化
do_constant_folding=True,
input_names=["input"],
output_names=["output"],
dynamic_axes={
"input": {0: "batch_size", 1: "num_channels", 2: "num_samples"},
"output": {0: "batch_size", 1: "num_frames"},
},
)
print(f"任务完成")