Support exporting models to onnx from 3D-Speaker (#522)

scripts/3dspeaker/export-onnx.py

@@ -0,0 +1,146 @@
+#!/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()