update

examples/silero_vad_by_webrtcvad/yaml/config.yaml

@@ -8,6 +8,7 @@ hop_size: 80
 win_type: hann
 
 # model
+conv_channels: 32
 hidden_size: 80
 kernel_size:
   - 3

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

@@ -13,8 +13,8 @@ class SileroVadConfig(PretrainedConfig):
                  hop_size: int = 80,
                  win_type: str = "hann",
 
-                 in_channels: int = 64,
-                 hidden_size: int = 128,
+                 conv_channels: int = 32,
+                 hidden_size: int = 80,
                  kernel_size: Tuple[int, int] = (3, 3),
 
                  n_frame: int = 3,
@@ -48,7 +48,7 @@ class SileroVadConfig(PretrainedConfig):
         self.win_type = win_type
 
         # encoder
-        self.in_channels = in_channels
+        self.conv_channels = conv_channels
         self.hidden_size = hidden_size
         self.kernel_size = kernel_size
 

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

@@ -132,6 +132,7 @@ class CausalConv2d(nn.Module):
 
 class CausalEncoder(nn.Module):
     def __init__(self,
+                 conv_channels: int,
                  kernel_size: Tuple[int, int] = (3, 3),
                  num_layers: int = 3,
                  ):
@@ -139,7 +140,7 @@ class CausalEncoder(nn.Module):
         self.layers: List[CausalConv2d] = nn.ModuleList(modules=[
             CausalConv2d(
                 in_channels=1,
-                out_channels=1,
+                out_channels=conv_channels,
                 kernel_size=kernel_size,
                 bias=False,
                 separable=True,
@@ -147,8 +148,8 @@ class CausalEncoder(nn.Module):
             )
             if i == 0 else
             CausalConv2d(
-                in_channels=1,
-                out_channels=1,
+                in_channels=conv_channels,
+                out_channels=conv_channels,
                 kernel_size=kernel_size,
                 bias=False,
                 separable=True,
@@ -160,7 +161,7 @@ class CausalEncoder(nn.Module):
     def forward(self, x: torch.Tensor, cache_list: List[torch.Tensor] = None):
         # x shape: [b, t, f]
         x = torch.unsqueeze(x, dim=1)
-        # x shape: [b, c, t, f]
+        # x shape: [b, 1, t, f]
 
         new_cache_list = list()
         for idx, layer in enumerate(self.layers):
@@ -169,8 +170,11 @@ class CausalEncoder(nn.Module):
             new_cache_list.append(new_cache)
 
         # x shape: [b, c, t, f]
-        x = torch.squeeze(x, dim=1)
-        # x shape: [b, t, f]
+        x = x.permute(0, 2, 1, 3)
+        # x shape: [b, t, c, f]
+        b, t, c, f = x.shape
+        x = torch.reshape(x, shape=(b, t, c*f))
+        # x shape: [b, t, c*f]
         return x, new_cache_list
 
 
@@ -182,6 +186,7 @@ class SileroVadModel(nn.Module):
                  hop_size: int,
                  win_type: int,
 
+                 conv_channels: int,
                  hidden_size: int,
                  kernel_size: Tuple[int, int],
 
@@ -197,6 +202,7 @@ class SileroVadModel(nn.Module):
         self.hop_size = hop_size
         self.win_type = win_type
 
+        self.conv_channels = conv_channels
         self.hidden_size = hidden_size
         self.kernel_size = kernel_size
 
@@ -229,11 +235,12 @@ class SileroVadModel(nn.Module):
         )
 
         self.encoder = CausalEncoder(
+            conv_channels=conv_channels,
             kernel_size=(3, 3),
         )
 
         self.lstm = nn.LSTM(
-            input_size=self.hidden_size,
+            input_size=self.conv_channels * self.hidden_size,
             hidden_size=self.hidden_size,
             bidirectional=False,
             batch_first=True
@@ -338,6 +345,7 @@ class SileroVadPretrainedModel(SileroVadModel):
             win_size=config.win_size,
             hop_size=config.hop_size,
             win_type=config.win_type,
+            conv_channels=config.conv_channels,
             hidden_size=config.hidden_size,
             kernel_size=config.kernel_size,
             n_frame=config.n_frame,

toolbox/torchaudio/models/vad/silero_vad/yaml/config.yaml

@@ -8,6 +8,7 @@ hop_size: 80
 win_type: hann
 
 # model
+conv_channels: 32
 hidden_size: 80
 kernel_size:
   - 3