147 lines
4.0 KiB
Python
Executable File
147 lines
4.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# Copyright 2023-2024 Xiaomi Corp. (authors: Fangjun Kuang)
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import pathlib
|
|
import re
|
|
from typing import Dict
|
|
|
|
import onnx
|
|
import torch
|
|
from infer_sv import supports
|
|
from modelscope.hub.snapshot_download import snapshot_download
|
|
from speakerlab.utils.builder import dynamic_import
|
|
|
|
|
|
def add_meta_data(filename: str, meta_data: Dict[str, str]):
|
|
"""Add meta data to an ONNX model. It is changed in-place.
|
|
|
|
Args:
|
|
filename:
|
|
Filename of the ONNX model to be changed.
|
|
meta_data:
|
|
Key-value pairs.
|
|
"""
|
|
model = onnx.load(filename)
|
|
for key, value in meta_data.items():
|
|
meta = model.metadata_props.add()
|
|
meta.key = key
|
|
meta.value = str(value)
|
|
|
|
onnx.save(model, filename)
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
required=True,
|
|
choices=[
|
|
"speech_campplus_sv_en_voxceleb_16k",
|
|
"speech_campplus_sv_zh-cn_16k-common",
|
|
"speech_eres2net_sv_en_voxceleb_16k",
|
|
"speech_eres2net_sv_zh-cn_16k-common",
|
|
"speech_eres2net_base_200k_sv_zh-cn_16k-common",
|
|
"speech_eres2net_base_sv_zh-cn_3dspeaker_16k",
|
|
"speech_eres2net_large_sv_zh-cn_3dspeaker_16k",
|
|
],
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
@torch.no_grad()
|
|
def main():
|
|
args = get_args()
|
|
local_model_dir = "pretrained"
|
|
model_id = f"damo/{args.model}"
|
|
conf = supports[model_id]
|
|
cache_dir = snapshot_download(
|
|
model_id,
|
|
revision=conf["revision"],
|
|
)
|
|
cache_dir = pathlib.Path(cache_dir)
|
|
|
|
save_dir = os.path.join(local_model_dir, model_id.split("/")[1])
|
|
save_dir = pathlib.Path(save_dir)
|
|
save_dir.mkdir(exist_ok=True, parents=True)
|
|
|
|
download_files = ["examples", conf["model_pt"]]
|
|
for src in cache_dir.glob("*"):
|
|
if re.search("|".join(download_files), src.name):
|
|
dst = save_dir / src.name
|
|
try:
|
|
dst.unlink()
|
|
except FileNotFoundError:
|
|
pass
|
|
dst.symlink_to(src)
|
|
pretrained_model = save_dir / conf["model_pt"]
|
|
pretrained_state = torch.load(pretrained_model, map_location="cpu")
|
|
|
|
model = conf["model"]
|
|
embedding_model = dynamic_import(model["obj"])(**model["args"])
|
|
embedding_model.load_state_dict(pretrained_state)
|
|
embedding_model.eval()
|
|
|
|
with open(f"{cache_dir}/configuration.json") as f:
|
|
json_config = json.loads(f.read())
|
|
print(json_config)
|
|
|
|
T = 100
|
|
C = 80
|
|
x = torch.rand(1, T, C)
|
|
filename = f"{args.model}.onnx"
|
|
torch.onnx.export(
|
|
embedding_model,
|
|
x,
|
|
filename,
|
|
opset_version=13,
|
|
input_names=["x"],
|
|
output_names=["embedding"],
|
|
dynamic_axes={
|
|
"x": {0: "N", 1: "T"},
|
|
"embeddings": {0: "N"},
|
|
},
|
|
)
|
|
|
|
# all models from 3d-speaker expect input samples in the range
|
|
# [-1, 1]
|
|
normalize_samples = 1
|
|
|
|
# all models from 3d-speaker normalize the features by the global mean
|
|
feature_normalize_type = "global-mean"
|
|
sample_rate = json_config["model"]["model_config"]["sample_rate"]
|
|
|
|
feat_dim = conf["model"]["args"]["feat_dim"]
|
|
assert feat_dim == 80, feat_dim
|
|
|
|
output_dim = conf["model"]["args"]["embedding_size"]
|
|
|
|
if "zh-cn" in args.model:
|
|
language = "Chinese"
|
|
elif "en" in args.model:
|
|
language = "English"
|
|
else:
|
|
raise ValueError(f"Unsupported language for model {args.model}")
|
|
|
|
comment = f"This model is from damo/{args.model}"
|
|
url = f"https://www.modelscope.cn/models/damo/{args.model}/summary"
|
|
|
|
meta_data = {
|
|
"framework": "3d-speaker",
|
|
"language": language,
|
|
"url": url,
|
|
"comment": comment,
|
|
"sample_rate": sample_rate,
|
|
"output_dim": output_dim,
|
|
"normalize_samples": normalize_samples,
|
|
"feature_normalize_type": feature_normalize_type,
|
|
}
|
|
print(meta_data)
|
|
add_meta_data(filename=filename, meta_data=meta_data)
|
|
|
|
|
|
main()
|