update

.gitignore

@@ -23,3 +23,4 @@
 **/*.wav
 **/*.xlsx
 **/*.jsonl
+**/*.onnx

examples/silero_vad_by_webrtcvad/run.sh

@@ -126,13 +126,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
 
@@ -144,7 +142,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/silero_vad_by_webrtcvad/step_5_export_model.py

@@ -0,0 +1,112 @@
+#!/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 onnxruntime as ort
+import torch
+
+from toolbox.torchaudio.models.vad.silero_vad.modeling_silero_vad import SileroVadModel, SileroVadModelExport, SileroVadPretrainedModel
+
+
+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 = SileroVadPretrainedModel.from_pretrained(args.model_dir)
+    model.eval()
+    config = model.config
+
+    model_export = SileroVadModelExport(model)
+
+    encoder_num_layers = config.encoder_num_layers
+    p = (config.encoder_kernel_size - 1) // 2
+    encoder_in_channels = config.encoder_in_channels
+    encoder_hidden_channels = config.encoder_hidden_channels
+
+    decoder_num_layers = config.decoder_num_layers
+    decoder_hidden_size = config.decoder_hidden_size
+
+    b = 1
+    inputs = torch.randn(size=(b, 1, 16000), dtype=torch.float32)
+
+    encoder_in_cache = torch.zeros(size=(b, 2*p, encoder_in_channels), dtype=torch.float32)
+    encoder_hidden_cache_list = [
+        torch.zeros(size=(b, 2*p, encoder_hidden_channels), dtype=torch.float32)
+    ] * encoder_num_layers
+    encoder_hidden_cache_list = torch.stack(encoder_hidden_cache_list, dim=0)
+
+    lstm_hidden_state = [
+        torch.zeros(size=(decoder_num_layers, b, decoder_hidden_size), dtype=torch.float32)
+    ] * 2
+    lstm_hidden_state = torch.stack(lstm_hidden_state, dim=0)
+
+    logits, probs, lsnr, new_encoder_in_cache, new_encoder_hidden_cache_list, new_lstm_hidden_state = model_export.forward(
+        inputs, encoder_in_cache, encoder_hidden_cache_list, lstm_hidden_state
+    )
+
+    torch.onnx.export(model_export,
+                      args=(inputs, encoder_in_cache, encoder_hidden_cache_list, lstm_hidden_state),
+                      f="silero_vad.onnx",
+                      input_names=["inputs", "encoder_in_cache", "encoder_hidden_cache_list", "lstm_hidden_state"],
+                      output_names=[
+                          "logits", "probs", "lsnr",
+                          "new_encoder_in_cache",
+                          "new_encoder_hidden_cache_list",
+                          "new_lstm_hidden_state"
+                      ],
+                      dynamic_axes={
+                          "inputs": {0: "batch_size", 2: "num_samples"},
+                          "encoder_in_cache": {1: "batch_size"},
+                          "encoder_hidden_cache_list": {1: "batch_size"},
+                          "lstm_hidden_state": {2: "batch_size"},
+                          "logits": {0: "batch_size"},
+                          "probs": {0: "batch_size"},
+                          "lsnr": {0: "batch_size"},
+                          "new_encoder_in_cache": {1: "batch_size"},
+                          "new_encoder_hidden_cache_list": {1: "batch_size"},
+                          "new_lstm_hidden_state": {2: "batch_size"},
+                      })
+
+    ort_session = ort.InferenceSession("silero_vad.onnx")
+    input_feed = {
+        "inputs": inputs.numpy(),
+        "encoder_in_cache": encoder_in_cache.numpy(),
+        "encoder_hidden_cache_list": encoder_hidden_cache_list.numpy(),
+        "lstm_hidden_state": lstm_hidden_state.numpy(),
+    }
+    output_names = [
+        "logits", "probs", "lsnr", "new_encoder_in_cache", "new_encoder_hidden_cache_list", "new_lstm_hidden_state"
+    ]
+    logits, probs, lsnr, new_encoder_in_cache, new_encoder_hidden_cache_list, new_lstm_hidden_state = ort_session.run(output_names, input_feed)
+    return
+
+
+if __name__ == "__main__":
+    main()

toolbox/torchaudio/models/vad/silero_vad/configuration_silero_vad.py

@@ -14,6 +14,8 @@ class SileroVadConfig(PretrainedConfig):
                  win_type: str = "hann",
 
                  encoder_in_channels: int = 64,
+                 encoder_hidden_channels: int = 128,
+                 encoder_out_channels: int = 64,
                  encoder_kernel_size: int = 3,
                  encoder_num_layers: int = 3,
 
@@ -52,6 +54,8 @@ class SileroVadConfig(PretrainedConfig):
 
         # encoder
         self.encoder_in_channels = encoder_in_channels
+        self.encoder_hidden_channels = encoder_hidden_channels
+        self.encoder_out_channels = encoder_out_channels
         self.encoder_kernel_size = encoder_kernel_size
         self.encoder_num_layers = encoder_num_layers
 

toolbox/torchaudio/models/vad/silero_vad/modeling_silero_vad.py

@@ -62,6 +62,7 @@ class Encoder(nn.Module):
                  num_layers: int = 3,
                  ):
         super(Encoder, self).__init__()
+        self.num_layers = num_layers
 
         self.layers = nn.ModuleList(modules=[])
         for i in range(num_layers):
@@ -96,23 +97,33 @@ class EncoderExport(nn.Module):
     def __init__(self, model: Encoder):
         super(EncoderExport, self).__init__()
         self.layers = model.layers
+        self.num_layers = model.num_layers
 
-    def forward(self, x: torch.Tensor, cache_list: torch.Tensor):
+    def forward(self, x: torch.Tensor, in_cache: torch.Tensor, hidden_cache_list: torch.Tensor):
         # x shape: [b, t, f]
-        # cache_list shape: [num_layers, b, 2p, f]
+        # in_cache shape: [b, 2p, f1]
+        # hidden_cache_list shape: [num_layers, b, 2p, fi]
 
-        new_cache_list = list()
+        new_in_cache = None
+        new_hidden_cache_list = list()
         for idx, layer in enumerate(self.layers):
-            cache = cache_list[idx]
+            if idx == 0:
+                cache = in_cache
+            else:
+                cache = hidden_cache_list[idx]
+
             x_pad = torch.concat(tensors=[cache, x], dim=1)
             x = layer.forward(x_pad)
 
             _, twop, _ = cache.shape
             new_cache = x_pad[:, -twop:, :]
-            new_cache_list.append(new_cache)
+            if idx == 0:
+                new_in_cache = new_cache
+            else:
+                new_hidden_cache_list.append(new_cache)
 
-        new_cache_list = torch.stack(tensors=new_cache_list, dim=0)
-        return x, new_cache_list
+        new_hidden_cache_list = torch.stack(tensors=new_hidden_cache_list, dim=0)
+        return x, new_in_cache, new_hidden_cache_list
 
 
 class SileroVadModel(nn.Module):
@@ -123,6 +134,8 @@ class SileroVadModel(nn.Module):
                  hop_size: int,
                  win_type: int,
                  encoder_in_channels: int,
+                 encoder_hidden_channels: int,
+                 encoder_out_channels: int,
                  encoder_kernel_size: int,
                  encoder_num_layers: int,
                  decoder_hidden_size: int,
@@ -139,6 +152,8 @@ class SileroVadModel(nn.Module):
         self.win_type = win_type
 
         self.encoder_in_channels = encoder_in_channels
+        self.encoder_hidden_channels = encoder_hidden_channels
+        self.encoder_out_channels = encoder_out_channels
         self.encoder_kernel_size = encoder_kernel_size
         self.encoder_num_layers = encoder_num_layers
 
@@ -180,8 +195,8 @@ class SileroVadModel(nn.Module):
 
         self.encoder = Encoder(
             in_channels=self.encoder_in_channels,
-            hidden_channels=self.decoder_hidden_size,
-            out_channels=self.decoder_hidden_size,
+            hidden_channels=self.encoder_hidden_channels,
+            out_channels=self.encoder_out_channels,
             kernel_size=self.encoder_kernel_size,
             num_layers=self.encoder_num_layers,
         )
@@ -298,6 +313,8 @@ class SileroVadPretrainedModel(SileroVadModel):
             hop_size=config.hop_size,
             win_type=config.win_type,
             encoder_in_channels=config.encoder_in_channels,
+            encoder_hidden_channels=config.encoder_hidden_channels,
+            encoder_out_channels=config.encoder_out_channels,
             encoder_kernel_size=config.encoder_kernel_size,
             encoder_num_layers=config.encoder_num_layers,
             decoder_hidden_size=config.decoder_hidden_size,
@@ -362,10 +379,12 @@ class SileroVadModelExport(nn.Module):
 
     def forward(self,
                 signal: torch.Tensor,
-                encoder_cache_list: torch.Tensor,
+                encoder_in_cache: torch.Tensor,
+                encoder_hidden_cache_list: torch.Tensor,
                 lstm_hidden_state: torch.Tensor,
                 ):
-        # encoder_cache_list shape: [num_layers, b, 2p, f]
+        # encoder_in_cache shape: [b, 2p, f]
+        # encoder_hidden_cache_list shape: [num_layers, b, 2p, f]
         # lstm_hidden_state shape: [2, num_layers, b, h]
 
         # signal shape [b, 1, num_samples]
@@ -382,7 +401,9 @@ class SileroVadModelExport(nn.Module):
         # x = self.tpad.forward(x)
         # x shape: [b, t+p, f']
 
-        x, new_encoder_cache_list = self.encoder.forward(x, cache_list=encoder_cache_list)
+        x, new_encoder_in_cache, new_encoder_hidden_cache_list = self.encoder.forward(
+            x, in_cache=encoder_in_cache, hidden_cache_list=encoder_hidden_cache_list
+        )
         # x shape: [b, t, f']
 
         x, new_lstm_hidden_state = self.lstm.forward(x, (lstm_hidden_state[0], lstm_hidden_state[1]))
@@ -397,7 +418,7 @@ class SileroVadModelExport(nn.Module):
         lsnr = self.lsnr_fc.forward(x) * self.lsnr_scale + self.lsnr_offset
         # lsnr shape: [b, t, 1]
 
-        return logits, probs, lsnr, new_encoder_cache_list, new_lstm_hidden_state
+        return logits, probs, lsnr, new_encoder_in_cache, new_encoder_hidden_cache_list, new_lstm_hidden_state
 
 
 def main1():
@@ -425,6 +446,7 @@ def main2():
     encoder_num_layers = config.encoder_num_layers
     p = (config.encoder_kernel_size - 1) // 2
     encoder_in_channels = config.encoder_in_channels
+    encoder_hidden_channels = config.encoder_hidden_channels
 
     decoder_num_layers = config.decoder_num_layers
     decoder_hidden_size = config.decoder_hidden_size
@@ -432,49 +454,60 @@ def main2():
     b = 1
     inputs = torch.randn(size=(b, 1, 16000), dtype=torch.float32)
 
-    encoder_cache_list = [
-        torch.zeros(size=(b, 2*p, encoder_in_channels), dtype=torch.float32)
+    encoder_in_cache = torch.zeros(size=(b, 2*p, encoder_in_channels), dtype=torch.float32)
+    encoder_hidden_cache_list = [
+        torch.zeros(size=(b, 2*p, encoder_hidden_channels), dtype=torch.float32)
     ] * encoder_num_layers
-    encoder_cache_list = torch.stack(encoder_cache_list, dim=0)
+    encoder_hidden_cache_list = torch.stack(encoder_hidden_cache_list, dim=0)
 
     lstm_hidden_state = [
         torch.zeros(size=(decoder_num_layers, b, decoder_hidden_size), dtype=torch.float32)
     ] * 2
     lstm_hidden_state = torch.stack(lstm_hidden_state, dim=0)
 
-    logits, probs, lsnr, new_encoder_cache_list, new_lstm_hidden_state = model_export.forward(inputs, encoder_cache_list, lstm_hidden_state)
+    logits, probs, lsnr, new_encoder_in_cache, new_encoder_hidden_cache_list, new_lstm_hidden_state = model_export.forward(
+        inputs, encoder_in_cache, encoder_hidden_cache_list, lstm_hidden_state
+    )
     print(f"logits.shape: {logits.shape}")
-    print(f"new_encoder_cache_list.shape: {new_encoder_cache_list.shape}")
+    print(f"new_encoder_in_cache.shape: {new_encoder_in_cache.shape}")
+    print(f"new_encoder_hidden_cache_list.shape: {new_encoder_hidden_cache_list.shape}")
     print(f"new_lstm_hidden_state.shape: {new_lstm_hidden_state.shape}")
 
     torch.onnx.export(model_export,
-                      args=(inputs, encoder_cache_list, lstm_hidden_state),
+                      args=(inputs, encoder_in_cache, encoder_hidden_cache_list, lstm_hidden_state),
                       f="silero_vad.onnx",
-                      input_names=["inputs", "encoder_cache_list", "lstm_hidden_state"],
-                      output_names=["logits", "probs", "lsnr", "new_encoder_cache_list", "new_lstm_hidden_state"],
+                      input_names=["inputs", "encoder_in_cache", "encoder_hidden_cache_list", "lstm_hidden_state"],
+                      output_names=[
+                          "logits", "probs", "lsnr",
+                          "new_encoder_in_cache",
+                          "new_encoder_hidden_cache_list",
+                          "new_lstm_hidden_state"
+                      ],
                       dynamic_axes={
                           "inputs": {0: "batch_size", 2: "num_samples"},
-                          "encoder_cache_list": {1: "batch_size"},
+                          "encoder_in_cache": {1: "batch_size"},
+                          "encoder_hidden_cache_list": {1: "batch_size"},
                           "lstm_hidden_state": {2: "batch_size"},
                           "logits": {0: "batch_size"},
                           "probs": {0: "batch_size"},
                           "lsnr": {0: "batch_size"},
-                          "new_encoder_cache_list": {1: "batch_size"},
+                          "new_encoder_in_cache": {1: "batch_size"},
+                          "new_encoder_hidden_cache_list": {1: "batch_size"},
                           "new_lstm_hidden_state": {2: "batch_size"},
                       })
 
     ort_session = ort.InferenceSession("silero_vad.onnx")
     input_feed = {
         "inputs": inputs.numpy(),
-        "encoder_cache_list": encoder_cache_list.numpy(),
+        "encoder_in_cache": encoder_in_cache.numpy(),
+        "encoder_hidden_cache_list": encoder_hidden_cache_list.numpy(),
         "lstm_hidden_state": lstm_hidden_state.numpy(),
     }
     output_names = [
-        "logits", "probs", "lsnr", "new_encoder_cache_list", "new_lstm_hidden_state"
+        "logits", "probs", "lsnr", "new_encoder_in_cache", "new_encoder_hidden_cache_list", "new_lstm_hidden_state"
     ]
-    logits, probs, lsnr, new_encoder_cache_list, new_lstm_hidden_state = ort_session.run(output_names, input_feed)
+    logits, probs, lsnr, new_encoder_in_cache, new_encoder_hidden_cache_list, new_lstm_hidden_state = ort_session.run(output_names, input_feed)
     print(f"probs.shape: {probs.shape}")
-    print(f"new_encoder_cache_list.shape: {new_encoder_cache_list.shape}")
     return