27 lines
838 B
Python
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"任务完成")
|