105 lines
3.1 KiB
Python
Executable File
105 lines
3.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
# Copyright 2024 Xiaomi Corp. (authors: Fangjun Kuang)
|
|
|
|
import argparse
|
|
from typing import Dict
|
|
|
|
import nemo.collections.asr as nemo_asr
|
|
import onnx
|
|
import torch
|
|
|
|
|
|
def get_args():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
required=True,
|
|
choices=[
|
|
"speakerverification_speakernet",
|
|
"titanet_large",
|
|
"titanet_small",
|
|
"ecapa_tdnn",
|
|
],
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
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)
|
|
|
|
|
|
@torch.no_grad()
|
|
def main():
|
|
args = get_args()
|
|
speaker_model_config = nemo_asr.models.EncDecSpeakerLabelModel.from_pretrained(
|
|
model_name=args.model, return_config=True
|
|
)
|
|
preprocessor_config = speaker_model_config["preprocessor"]
|
|
|
|
print(args.model)
|
|
print(speaker_model_config)
|
|
print(preprocessor_config)
|
|
|
|
assert preprocessor_config["n_fft"] == 512, preprocessor_config
|
|
|
|
assert (
|
|
preprocessor_config["_target_"]
|
|
== "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor"
|
|
), preprocessor_config
|
|
|
|
assert preprocessor_config["frame_splicing"] == 1, preprocessor_config
|
|
|
|
speaker_model = nemo_asr.models.EncDecSpeakerLabelModel.from_pretrained(
|
|
model_name=args.model
|
|
)
|
|
speaker_model.eval()
|
|
filename = f"nemo_en_{args.model}.onnx"
|
|
speaker_model.export(filename)
|
|
|
|
print(f"Adding metadata to {filename}")
|
|
|
|
comment = "This model is from NeMo."
|
|
url = {
|
|
"titanet_large": "https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/titanet_large",
|
|
"titanet_small": "https://catalog.ngc.nvidia.com/orgs/nvidia/teams/nemo/models/titanet_small",
|
|
"speakerverification_speakernet": "https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerverification_speakernet",
|
|
"ecapa_tdnn": "https://ngc.nvidia.com/catalog/models/nvidia:nemo:ecapa_tdnn",
|
|
}[args.model]
|
|
|
|
language = "English"
|
|
|
|
meta_data = {
|
|
"framework": "nemo",
|
|
"language": language,
|
|
"url": url,
|
|
"comment": comment,
|
|
"sample_rate": preprocessor_config["sample_rate"],
|
|
"output_dim": speaker_model_config["decoder"]["emb_sizes"],
|
|
"feature_normalize_type": preprocessor_config["normalize"],
|
|
"window_size_ms": int(float(preprocessor_config["window_size"]) * 1000),
|
|
"window_stride_ms": int(float(preprocessor_config["window_stride"]) * 1000),
|
|
"window_type": preprocessor_config["window"], # e.g., hann
|
|
"feat_dim": preprocessor_config["features"],
|
|
}
|
|
print(meta_data)
|
|
add_meta_data(filename=filename, meta_data=meta_data)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|