update

examples/fsmn_vad_by_webrtcvad/run.sh

@@ -127,13 +127,11 @@ fi
 
 
 if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
-  $verbose && echo "stage 4: test model"
+  $verbose && echo "stage 4: export model"
   cd "${work_dir}" || exit 1
-  python3 step_3_evaluation.py \
-  --valid_dataset "${valid_dataset}" \
+  python3 step_5_export_model.py \
   --model_dir "${file_dir}/best" \
-  --evaluation_audio_dir "${evaluation_audio_dir}" \
-  --limit "${limit}" \
+  --output_dir "${file_dir}/best" \
 
 fi
 
@@ -145,7 +143,6 @@ if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
   mkdir -p ${final_model_dir}
 
   cp "${file_dir}/best"/* "${final_model_dir}"
-  cp -r "${file_dir}/evaluation_audio" "${final_model_dir}"
 
   cd "${final_model_dir}/.." || exit 1;
 

examples/fsmn_vad_by_webrtcvad/step_5_export_model.py

@@ -0,0 +1,77 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import os
+from pathlib import Path
+import sys
+
+pwd = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(pwd, "../../"))
+
+import torch
+
+from toolbox.torchaudio.models.vad.fsmn_vad.modeling_fsmn_vad import FSMNVadModel, FSMNVadPretrainedModel, FSMNVadModelExport
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    # parser.add_argument("--model_dir", default="file_dir/best", type=str)
+    # parser.add_argument("--output_dir", default="file_dir/best", type=str)
+
+    parser.add_argument(
+        "--model_dir",
+        default=r"D:\Users\tianx\HuggingSpaces\cc_vad\trained_models\fsmn-vad-by-webrtcvad-nx2-dns3\fsmn-vad-by-webrtcvad-nx2-dns3",
+        type=str
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=r"D:\Users\tianx\HuggingSpaces\cc_vad\trained_models\fsmn-vad-by-webrtcvad-nx2-dns3\fsmn-vad-by-webrtcvad-nx2-dns3",
+        type=str
+    )
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = get_args()
+
+    output_dir = Path(args.output_dir)
+    output_file = output_dir / "model.onnx"
+
+    model = FSMNVadPretrainedModel.from_pretrained(args.model_dir)
+    model.eval()
+    config = model.config
+
+    basic_block_layers = config.fsmn_basic_block_layers
+    hidden_size = config.fsmn_basic_block_hidden_size
+    basic_block_lorder = config.fsmn_basic_block_lorder
+    basic_block_lstride = config.fsmn_basic_block_lstride
+
+    model_export = FSMNVadModelExport(model)
+
+    b = 1
+    inputs = torch.randn(size=(b, 1, 16000), dtype=torch.float32)
+    cache_list = [
+        torch.zeros(size=(b, hidden_size, (basic_block_lorder - 1) * basic_block_lstride, 1)),
+    ] * basic_block_layers
+    cache_list = torch.stack(cache_list, dim=0)
+
+    torch.onnx.export(model_export,
+                      args=(inputs, cache_list),
+                      f=output_file.as_posix(),
+                      input_names=["inputs", "cache_list"],
+                      output_names=["logits", "probs", "lsnr", "new_cache_list"],
+                      dynamic_axes={
+                          "inputs": {0: "batch_size", 2: "num_samples"},
+                          "cache_list": {0: "basic_block_layers", 1: "batch_size"},
+                          "logits": {0: "batch_size"},
+                          "probs": {0: "batch_size"},
+                          "lsnr": {0: "batch_size"},
+                          "new_cache_list": {0: "basic_block_layers", 1: "batch_size"},
+                      })
+
+    return
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -4,21 +4,24 @@ import argparse
 from functools import lru_cache
 import json
 import logging
+from pathlib import Path
 import platform
+import shutil
 import tempfile
 import time
 from typing import Dict, Tuple
+import zipfile
 
 import gradio as gr
-import librosa
-import librosa.display
+from huggingface_hub import snapshot_download
 import matplotlib.pyplot as plt
 import numpy as np
 
 import log
 from project_settings import environment, project_path, log_directory, time_zone_info
 from toolbox.os.command import Command
-from toolbox.torchaudio.models.vad.fsmn_vad.inference_fsmn_vad import InferenceFSMNVad
+from toolbox.torchaudio.models.vad.fsmn_vad.inference_fsmn_vad_onnx import InferenceFSMNVadOnnx
+from toolbox.torchaudio.models.vad.silero_vad.inference_silero_vad import InferenceSileroVad
 from toolbox.torchaudio.utils.visualization import process_speech_probs
 
 log.setup_size_rotating(log_directory=log_directory, tz_info=time_zone_info)
@@ -28,6 +31,22 @@ logger = logging.getLogger("main")
 
 def get_args():
     parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--examples_dir",
+        # default=(project_path / "data").as_posix(),
+        default=(project_path / "data/examples").as_posix(),
+        type=str
+    )
+    parser.add_argument(
+        "--models_repo_id",
+        default="qgyd2021/cc_vad",
+        type=str
+    )
+    parser.add_argument(
+        "--trained_model_dir",
+        default=(project_path / "trained_models").as_posix(),
+        type=str
+    )
     parser.add_argument(
         "--hf_token",
         default=environment.get("hf_token"),
@@ -49,7 +68,9 @@ def shell(cmd: str):
 
 def get_infer_cls_by_model_name(model_name: str):
     if model_name.__contains__("fsmn"):
-        infer_cls = InferenceFSMNVad
+        infer_cls = InferenceFSMNVadOnnx
+    elif model_name.__contains__("silero"):
+        infer_cls = InferenceSileroVad
     else:
         raise AssertionError
     return infer_cls
@@ -111,7 +132,8 @@ def when_click_vad_button(audio_file_t = None, audio_microphone_t = None, engine
 
         probs = vad_info["probs"]
         lsnr = vad_info["lsnr"]
-        lsnr = lsnr / np.max(np.abs(lsnr))
+        # lsnr = lsnr / np.max(np.abs(lsnr))
+        lsnr = lsnr / 30
 
         frame_step = infer_engine.config.hop_size
         probs = process_speech_probs(audio, probs, frame_step)
@@ -128,6 +150,18 @@ def when_click_vad_button(audio_file_t = None, audio_microphone_t = None, engine
 def main():
     args = get_args()
 
+    examples_dir = Path(args.examples_dir)
+    trained_model_dir = Path(args.trained_model_dir)
+
+    # download models
+    if not trained_model_dir.exists():
+        trained_model_dir.mkdir(parents=True, exist_ok=True)
+        _ = snapshot_download(
+            repo_id=args.models_repo_id,
+            local_dir=trained_model_dir.as_posix(),
+            token=args.hf_token,
+        )
+
     # engines
     global vad_engines
     vad_engines = {
@@ -152,6 +186,25 @@ def main():
     # choices
     vad_engine_choices = list(vad_engines.keys())
 
+    # examples
+    if not examples_dir.exists():
+        example_zip_file = trained_model_dir / "examples.zip"
+        with zipfile.ZipFile(example_zip_file.as_posix(), "r") as f_zip:
+            out_root = examples_dir
+            if out_root.exists():
+                shutil.rmtree(out_root.as_posix())
+            out_root.mkdir(parents=True, exist_ok=True)
+            f_zip.extractall(path=out_root)
+
+    # examples
+    examples = list()
+    for filename in examples_dir.glob("**/*.wav"):
+        examples.append([
+            filename.as_posix(),
+            None,
+            vad_engine_choices[0],
+        ])
+
     # ui
     with gr.Blocks() as blocks:
         gr.Markdown(value="vad.")
@@ -175,7 +228,15 @@ def main():
                 vad_button.click(
                     when_click_vad_button,
                     inputs=[vad_audio_file, vad_audio_microphone, vad_engine],
-                    outputs=[vad_vad_image, vad_lsnr_image, vad_message]
+                    outputs=[vad_vad_image, vad_lsnr_image, vad_message],
+                )
+                gr.Examples(
+                    examples=examples,
+                    inputs=[vad_audio_file, vad_audio_microphone, vad_engine],
+                    outputs=[vad_vad_image, vad_lsnr_image, vad_message],
+                    fn=when_click_vad_button,
+                    # cache_examples=True,
+                    # cache_mode="lazy",
                 )
             with gr.TabItem("shell"):
                 shell_text = gr.Textbox(label="cmd")

requirements.txt

@@ -12,3 +12,5 @@ overrides==7.7.0
 webrtcvad==2.0.10
 matplotlib==3.10.3
 google-genai
+onnx==1.18.0
+onnxruntime==1.22.1

toolbox/torchaudio/models/vad/fsmn_vad/fsmn_encoder.py

@@ -183,6 +183,29 @@ class BasicBlock(nn.Module):
         return x4, new_cache
 
 
+class BasicBlockExport(nn.Module):
+    def __init__(self, model: BasicBlock):
+        super(BasicBlockExport, self).__init__()
+        self.linear = model.linear
+        self.fsmn_block = model.fsmn_block
+        self.affine = model.affine
+        self.relu = model.relu
+
+    def forward(self, inputs: torch.Tensor, cache: torch.Tensor):
+        # inputs shape: [b, t, f]
+        x1 = self.linear.forward(inputs)
+        # x1 shape: [b, t, f']
+
+        x2, new_cache = self.fsmn_block.forward(x1, cache=cache)
+        # x2 shape: [b, t, f']
+
+        x3 = self.affine.forward(x2)
+        # x3 shape: [b, t, f]
+
+        x4 = self.relu(x3)
+        return x4, new_cache
+
+
 class FSMN(nn.Module):
     def __init__(
         self,
@@ -251,28 +274,193 @@ class FSMN(nn.Module):
         return outputs, new_cache_list
 
 
-def main():
+class FSMNExport(nn.Module):
+    def __init__(self, model: FSMN):
+        super(FSMNExport, self).__init__()
+        self.in_linear1 = model.in_linear1
+        self.in_linear2 = model.in_linear2
+        self.relu = model.relu
+
+        self.out_linear1 = model.out_linear1
+        self.out_linear2 = model.out_linear2
+
+        self.fsmn_basic_block_list = nn.ModuleList(modules=[])
+        for i, d in enumerate(model.fsmn_basic_block_list):
+            if isinstance(d, BasicBlock):
+                self.fsmn_basic_block_list.append(BasicBlockExport(d))
+
+    def forward(self,
+                inputs: torch.Tensor,
+                cache_list: torch.Tensor,
+                ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # cache_list shape: [basic_block_layers, b, hidden_size, (basic_block_lorder - 1) * basic_block_lstride, 1]
+
+        # inputs shape: [b, t, f]
+        x = self.in_linear1.forward(inputs)
+        # x shape: [b, t, input_affine_dim]
+        x = self.in_linear2.forward(x)
+        # x shape: [b, t, f]
+
+        x = self.relu(x)
+
+        new_cache_list = list()
+        for idx, fsmn_basic_block in enumerate(self.fsmn_basic_block_list):
+            cache = cache_list[idx]
+            x, new_cache = fsmn_basic_block.forward(x, cache)
+            new_cache_list.append(new_cache)
+        new_cache_list = torch.stack(new_cache_list, dim=0)
+
+        # x shape: [b, t, f]
+        x = self.out_linear1.forward(x)
+        outputs = self.out_linear2.forward(x)
+        # outputs shape: [b, t, f]
+
+        return outputs, new_cache_list
+
+
+def main1():
+    import onnx
+    import onnxruntime as ort
+
+    input_size = 32
+    input_affine_size = 16
+    hidden_size = 16
+    basic_block_layers = 3
+    basic_block_hidden_size = 16
+    basic_block_lorder = 3
+    basic_block_rorder = 0
+    basic_block_lstride = 1
+    basic_block_rstride = 1
+    output_affine_size = 16
+    output_size = 32
+
+    basic_block = BasicBlock(
+        input_size=hidden_size,
+        hidden_size=basic_block_hidden_size,
+        lorder=basic_block_lorder,
+        rorder=basic_block_rorder,
+        lstride=basic_block_lstride,
+        rstride=basic_block_rstride,
+    )
+
+    b = 1
+    t = 198
+    f = hidden_size
+    inputs = torch.randn(size=(b, t, f), dtype=torch.float32)
+
+    result, _ = basic_block.forward(inputs)
+    print(f"result.shape: {result.shape}")
+
+    basic_block_export = BasicBlockExport(model=basic_block)
+
+    cache = torch.zeros(size=(b, hidden_size, (basic_block_lorder - 1) * basic_block_lstride, 1))
+    result, new_cache = basic_block_export.forward(inputs, cache)
+    print(f"result.shape: {result.shape}")
+    print(f"new_cache.shape: {new_cache.shape}")
+
+    torch.onnx.export(basic_block_export,
+                      args=(inputs, cache),
+                      f="basic_block.onnx",
+                      input_names=["inputs", "cache"],
+                      output_names=["outputs", "new_cache"],
+                      dynamic_axes={
+                          "inputs": {0: "batch_size"},
+                          "cache": {0: "batch_size"},
+                          "outputs": {0: "batch_size"},
+                          "new_cache": {0: "batch_size"},
+                      })
+
+    ort_session = ort.InferenceSession("basic_block.onnx")
+    input_feed = {
+        "inputs": inputs.numpy(),
+        "cache": cache.numpy(),
+    }
+    output_names = [
+        "outputs",
+        "new_cache"
+    ]
+    outputs = ort_session.run(output_names, input_feed)
+    print(outputs)
+    print(len(outputs))
+    return
+
+
+def main2():
+    import onnx
+    import onnxruntime as ort
+
+    input_size = 32
+    input_affine_size = 16
+    hidden_size = 16
+    basic_block_layers = 3
+    basic_block_hidden_size = 16
+    basic_block_lorder = 3
+    basic_block_rorder = 0
+    basic_block_lstride = 1
+    basic_block_rstride = 1
+    output_affine_size = 16
+    output_size = 32
+
     fsmn = FSMN(
-        input_size=32,
-        input_affine_size=16,
-        hidden_size=16,
-        basic_block_layers=3,
-        basic_block_hidden_size=16,
-        basic_block_lorder=3,
-        basic_block_rorder=0,
-        basic_block_lstride=1,
-        basic_block_rstride=1,
-        output_affine_size=16,
-        output_size=32,
+        input_size=input_size,
+        input_affine_size=input_affine_size,
+        hidden_size=hidden_size,
+        basic_block_layers=basic_block_layers,
+        basic_block_hidden_size=basic_block_hidden_size,
+        basic_block_lorder=basic_block_lorder,
+        basic_block_rorder=basic_block_rorder,
+        basic_block_lstride=basic_block_lstride,
+        basic_block_rstride=basic_block_rstride,
+        output_affine_size=output_affine_size,
+        output_size=output_size,
     )
 
-    inputs = torch.randn(size=(1, 198, 32), dtype=torch.float32)
+    b = 1
+    t = 198
+    f = input_size
+    inputs = torch.randn(size=(b, t, f), dtype=torch.float32)
 
     result, _ = fsmn.forward(inputs)
-    print(result.shape)
-
+    print(f"result.shape: {result.shape}")
+
+    fsmn_export = FSMNExport(model=fsmn)
+
+    cache_list = [
+        torch.zeros(size=(b, hidden_size, (basic_block_lorder - 1) * basic_block_lstride, 1)),
+        torch.zeros(size=(b, hidden_size, (basic_block_lorder - 1) * basic_block_lstride, 1)),
+        torch.zeros(size=(b, hidden_size, (basic_block_lorder - 1) * basic_block_lstride, 1)),
+    ]
+    cache_list = torch.stack(cache_list, dim=0)
+    result, new_cache_list = fsmn_export.forward(inputs, cache_list)
+    print(f"result.shape: {result.shape}")
+    print(f"new_cache_list.shape: {new_cache_list.shape}")
+
+    torch.onnx.export(fsmn_export,
+                      args=(inputs, cache_list),
+                      f="fsmn.onnx",
+                      input_names=["inputs", "cache_list"],
+                      output_names=["outputs", "new_cache_list"],
+                      dynamic_axes={
+                          "inputs": {0: "batch_size"},
+                          "cache_list": {0: "basic_block_layers", 1: "batch_size"},
+                          "outputs": {0: "batch_size"},
+                          "new_cache_list": {0: "basic_block_layers", 1: "batch_size"},
+                      })
+
+    ort_session = ort.InferenceSession("fsmn.onnx")
+    input_feed = {
+        "inputs": inputs.numpy(),
+        "cache_list": cache_list.numpy(),
+    }
+    output_names = [
+        "outputs",
+        "new_cache_list"
+    ]
+    outputs, new_cache_list = ort_session.run(output_names, input_feed)
+    print(f"outputs.shape: {outputs.shape}")
+    print(f"new_cache_list.shape: {new_cache_list.shape}")
     return
 
 
 if __name__ == "__main__":
-    main()
+    main2()

toolbox/torchaudio/models/vad/fsmn_vad/inference_fsmn_vad.py

@@ -18,7 +18,7 @@ torch.set_num_threads(1)
 
 from project_settings import project_path
 from toolbox.torchaudio.models.vad.fsmn_vad.configuration_fsmn_vad import FSMNVadConfig
-from toolbox.torchaudio.models.vad.fsmn_vad.modeling_fsmn_vad import FSMNVadPretrainedModel, MODEL_FILE
+from toolbox.torchaudio.models.vad.fsmn_vad.modeling_fsmn_vad import FSMNVadPretrainedModel, MODEL_FILE, FSMNVadModelExport
 from toolbox.torchaudio.utils.visualization import process_speech_probs, make_visualization
 
 

toolbox/torchaudio/models/vad/fsmn_vad/inference_fsmn_vad_onnx.py

@@ -0,0 +1,168 @@
+#!/usr/bin/python3
+# -*- coding: utf-8 -*-
+import argparse
+import logging
+from pathlib import Path
+import shutil
+import tempfile, time
+from typing import List
+import zipfile
+
+from scipy.io import wavfile
+import numpy as np
+import torch
+import onnxruntime as ort
+
+torch.set_num_threads(1)
+
+from project_settings import project_path
+from toolbox.torchaudio.models.vad.fsmn_vad.configuration_fsmn_vad import FSMNVadConfig
+from toolbox.torchaudio.utils.visualization import process_speech_probs, make_visualization
+
+
+logger = logging.getLogger("toolbox")
+
+
+class InferenceFSMNVadOnnx(object):
+    def __init__(self, pretrained_model_path_or_zip_file: str, device: str = "cpu"):
+        self.pretrained_model_path_or_zip_file = pretrained_model_path_or_zip_file
+        self.device = torch.device(device)
+
+        logger.info(f"loading model; model_file: {self.pretrained_model_path_or_zip_file}")
+        config, ort_session = self.load_models(self.pretrained_model_path_or_zip_file)
+        logger.info(f"model loading completed; model_file: {self.pretrained_model_path_or_zip_file}")
+
+        self.config = config
+        self.ort_session = ort_session
+
+    def load_models(self, model_path: str):
+        model_path = Path(model_path)
+        if model_path.name.endswith(".zip"):
+            with zipfile.ZipFile(model_path.as_posix(), "r") as f_zip:
+                out_root = Path(tempfile.gettempdir()) / "cc_vad"
+                out_root.mkdir(parents=True, exist_ok=True)
+                f_zip.extractall(path=out_root)
+            model_path = out_root / model_path.stem
+
+        config = FSMNVadConfig.from_pretrained(
+            pretrained_model_name_or_path=model_path.as_posix(),
+        )
+        ort_session = ort.InferenceSession(
+            path_or_bytes=(model_path / "model.onnx").as_posix()
+        )
+
+        shutil.rmtree(model_path)
+        return config, ort_session
+
+    def infer(self, signal: np.ndarray) -> np.ndarray:
+        # signal shape: [num_samples,], value between -1 and 1.
+
+        inputs = torch.tensor(signal, dtype=torch.float32)
+        inputs = torch.unsqueeze(inputs, dim=0)
+        inputs = torch.unsqueeze(inputs, dim=0)
+        # inputs shape: [1, 1, num_samples]
+
+        b = 1
+        cache_list = [
+                         torch.zeros(size=(
+                             b, self.config.fsmn_basic_block_hidden_size,
+                             (self.config.fsmn_basic_block_lorder - 1) * self.config.fsmn_basic_block_lstride,
+                             1
+                         )),
+                     ] * self.config.fsmn_basic_block_layers
+        cache_list = torch.stack(cache_list, dim=0)
+
+        input_feed = {
+            "inputs": inputs.numpy(),
+            "cache_list": cache_list.numpy(),
+        }
+        output_names = [
+            "logits", "probs", "lsnr", "new_cache_list"
+        ]
+        logits, probs, lsnr, new_cache_list = self.ort_session.run(output_names, input_feed)
+        # probs shape: [b, t, 1]
+        probs = np.squeeze(probs, axis=-1)
+        # probs shape: [b, t]
+        probs = probs[0]
+
+        # lsnr shape: [b, t, 1]
+        lsnr = np.squeeze(lsnr, axis=-1)
+        # lsnr shape: [b, t]
+        lsnr = lsnr[0]
+
+        result = {
+            "probs": probs,
+            "lsnr": lsnr,
+        }
+        return result
+
+    def post_process(self, probs: List[float]):
+        return
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--wav_file",
+        # default=(project_path / "data/examples/ai_agent/chinese-4.wav").as_posix(),
+        # default=(project_path / "data/examples/ai_agent/chinese-5.wav").as_posix(),
+        # default=(project_path / "data/examples/hado/b556437e-c68b-4f6d-9eed-2977c29db887.wav").as_posix(),
+        # default=(project_path / "data/examples/hado/eae93a33-8ee0-4d86-8f85-cac5116ae6ef.wav").as_posix(),
+        # default=(project_path / "data/examples/speech/active_media_r_0ba69730-66a4-4ecd-8929-ef58f18f4612_2.wav").as_posix(),
+        # default=(project_path / "data/examples/speech/active_media_r_2a2f472b-a0b8-4fd5-b1c4-1aedc5d2ce57_0.wav").as_posix(),
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_w_8b6e28e2-a238-4c8c-b2e3-426b1fca149b_6.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0a56f035-40f6-4530-b852-613f057d718d_6.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0ae70b76-3651-4a71-bc0c-9e1429e4c854_5.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0d483249-57f8-4d45-b4c6-bda82d6816ae_2.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0d952885-5bc2-4633-81b6-e0e809e113f1_2.wav",
+        default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\speech\en-SG\2025-05-19\active_media_r_0ddac777-d986-4a5c-9c7c-ff64be0a463d_11.wav",
+
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_0b8a8e80-52af-423b-8877-03a78b1e6e43_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_0ebffb68-6490-4a8b-8eb6-eb82443d7d75_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_0f6ec933-90df-447b-aca4-6ddc149452ab_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1aac396f-1661-4f26-ab49-1a4879684567_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1aac396f-1661-4f26-ab49-1a4879684567_1.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1aff518b-4749-42fc-adfe-64046f9baeb6_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1b16f2a3-a8c9-4739-9a76-59faf1c64d79_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1b16f2a3-a8c9-4739-9a76-59faf1c64d79_1.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1bb1f22e-9c3a-4aea-b53f-71cc6547a6ee_0.wav",
+        # default=r"D:\Users\tianx\HuggingDatasets\nx_noise\data\noise\en-SG\2025-05-19\active_media_r_1dab161b-2a76-4491-abd1-60dba6172f8d_2.wav",
+        type=str,
+    )
+    args = parser.parse_args()
+    return args
+
+
+SAMPLE_RATE = 8000
+
+
+def main():
+    args = get_args()
+
+    sample_rate, signal = wavfile.read(args.wav_file)
+    if SAMPLE_RATE != sample_rate:
+        raise AssertionError
+    signal = signal / (1 << 15)
+
+    infer = InferenceFSMNVadOnnx(
+        # pretrained_model_path_or_zip_file=(project_path / "trained_models/fsmn-vad-by-webrtcvad-nx-dns3.zip").as_posix(),
+        pretrained_model_path_or_zip_file = (project_path / "trained_models/fsmn-vad-by-webrtcvad-nx2-dns3.zip").as_posix(),
+    )
+    frame_step = infer.config.hop_size
+
+    speech_probs: np.ndarray = infer.infer(signal)
+    speech_probs = speech_probs.tolist()
+
+    speech_probs = process_speech_probs(
+        signal=signal,
+        speech_probs=speech_probs,
+        frame_step=frame_step,
+    )
+
+    # plot
+    make_visualization(signal, speech_probs, SAMPLE_RATE)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/vad/fsmn_vad/modeling_fsmn_vad.py

@@ -20,7 +20,7 @@ from torch.nn import functional as F
 from toolbox.torchaudio.configuration_utils import CONFIG_FILE
 from toolbox.torchaudio.models.vad.fsmn_vad.configuration_fsmn_vad import FSMNVadConfig
 from toolbox.torchaudio.modules.conv_stft import ConvSTFT
-from toolbox.torchaudio.models.vad.fsmn_vad.fsmn_encoder import FSMN
+from toolbox.torchaudio.models.vad.fsmn_vad.fsmn_encoder import FSMN, FSMNExport
 from toolbox.torchaudio.modules.local_snr_target import LocalSnrTarget
 
 
@@ -243,7 +243,45 @@ class FSMNVadPretrainedModel(FSMNVadModel):
         return save_directory
 
 
-def main():
+class FSMNVadModelExport(nn.Module):
+    def __init__(self, model: FSMNVadModel):
+        super(FSMNVadModelExport, self).__init__()
+        self.stft = model.stft
+        self.fsmn_encoder = FSMNExport(model.fsmn_encoder)
+
+        # lsnr
+        self.lsnr_scale = model.lsnr_scale
+        self.lsnr_offset = model.lsnr_offset
+
+    def forward(self,
+                signal: torch.Tensor,
+                cache_list: torch.Tensor,
+                ):
+        # signal shape [b, 1, num_samples]
+
+        mags = self.stft.forward(signal)
+        # mags shape: [b, f, t]
+
+        x = torch.transpose(mags, dim0=1, dim1=2)
+        # x shape: [b, t, f]
+
+        logits, new_cache_list = self.fsmn_encoder.forward(x, cache_list)
+        # logits shape: [b, t, 2]
+
+        splits = torch.split(logits, split_size_or_sections=[1, 1], dim=-1)
+        vad_logits = splits[0]
+        snr_logits = splits[1]
+        # shape: [b, t, 1]
+        vad_probs = F.sigmoid(vad_logits)
+        # vad_probs shape: [b, t, 1]
+
+        lsnr = F.sigmoid(snr_logits) * self.lsnr_scale + self.lsnr_offset
+        # lsnr shape: [b, t, 1]
+
+        return vad_logits, vad_probs, lsnr, new_cache_list
+
+
+def main1():
     config = FSMNVadConfig()
     model = FSMNVadPretrainedModel(config=config)
 
@@ -253,9 +291,62 @@ def main():
     print(f"logits.shape: {logits.shape}")
     print(f"probs.shape: {probs.shape}")
     print(f"lsnr.shape: {lsnr.shape}")
+    return
+
 
+def main2():
+    import onnx
+    import onnxruntime as ort
+
+    config = FSMNVadConfig()
+    model = FSMNVadPretrainedModel(config=config)
+
+    basic_block_layers = config.fsmn_basic_block_layers
+    hidden_size = config.fsmn_basic_block_hidden_size
+    basic_block_lorder = config.fsmn_basic_block_lorder
+    basic_block_lstride = config.fsmn_basic_block_lstride
+
+    model_export = FSMNVadModelExport(model)
+
+    b = 1
+    inputs = torch.randn(size=(b, 1, 16000), dtype=torch.float32)
+    cache_list = [
+        torch.zeros(size=(b, hidden_size, (basic_block_lorder - 1) * basic_block_lstride, 1)),
+    ] * basic_block_layers
+    cache_list = torch.stack(cache_list, dim=0)
+
+    logits, probs, lsnr, new_cache_list = model_export.forward(inputs, cache_list)
+    print(f"logits.shape: {logits.shape}")
+    print(f"new_cache_list.shape: {new_cache_list.shape}")
+
+    torch.onnx.export(model_export,
+                      args=(inputs, cache_list),
+                      f="fsmn_vad.onnx",
+                      input_names=["inputs", "cache_list"],
+                      output_names=["logits", "probs", "lsnr", "new_cache_list"],
+                      dynamic_axes={
+                          "inputs": {0: "batch_size", 2: "num_samples"},
+                          "cache_list": {0: "basic_block_layers", 1: "batch_size"},
+                          "logits": {0: "batch_size"},
+                          "probs": {0: "batch_size"},
+                          "lsnr": {0: "batch_size"},
+                          "new_cache_list": {0: "basic_block_layers", 1: "batch_size"},
+                      })
+
+    ort_session = ort.InferenceSession("fsmn_vad.onnx")
+    input_feed = {
+        "inputs": inputs.numpy(),
+        "cache_list": cache_list.numpy(),
+    }
+    output_names = [
+        "outputs",
+        "new_cache_list"
+    ]
+    outputs, new_cache_list = ort_session.run(output_names, input_feed)
+    print(f"outputs.shape: {outputs.shape}")
+    print(f"new_cache_list.shape: {new_cache_list.shape}")
     return
 
 
 if __name__ == "__main__":
-    main()
+    main2()