Upload modeling_pyannote.py

modeling_pyannote.py

@@ -0,0 +1,162 @@
+import torch
+import torch.nn as nn
+
+from models.pyannote.layers import SincNet
+
+from asteroid_filterbanks.enc_dec import Filterbank, Encoder
+from asteroid_filterbanks.param_sinc_fb import ParamSincFB
+
+
+
+class SincNet(nn.Module):
+    """Filtering and convolutional part of Pyannote
+
+    Arguments
+    ---------
+    n_filters : list, int
+        List consist of number of each convolution kernel
+    stride_ : in 
+        Stride of ParamSincFB fliltering.
+
+    Returns
+    -------
+    Sincnet model: class
+
+    """
+    
+    def __init__(self, 
+                 n_filters = [80,60,60],
+                 stride_ = 10,
+                 ):
+        super(SincNet,self).__init__()
+        
+
+        sincnet_list = nn.ModuleList(
+            [
+                nn.InstanceNorm1d(1, eps=1e-05, momentum=0.1, affine=True, track_running_stats=False),
+                Encoder(ParamSincFB(n_filters=n_filters[0], kernel_size=251, stride=stride_)),
+                nn.MaxPool1d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False),
+                nn.InstanceNorm1d(n_filters[0], eps=1e-05, momentum=0.1, affine=True, track_running_stats=False),
+            ]
+        )
+        for counter in range(len(n_filters) - 1):
+            sincnet_list.append(nn.Conv1d(n_filters[counter], n_filters[counter+1], kernel_size=(5,), stride=(1,)))
+            sincnet_list.append(nn.MaxPool1d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False))
+            sincnet_list.append(nn.InstanceNorm1d(n_filters[counter+1], eps=1e-05, momentum=0.1, affine=True, track_running_stats=False))
+
+        self.sincnet_layer = nn.Sequential(*sincnet_list)
+
+    def forward(self, x):
+        """This method should implement forwarding operation in the SincNet model.
+
+        Arguments
+        ---------
+        x : float (Tensor)
+            The input of SincNet model.
+
+        Returns
+        -------
+        out : float (Tensor)
+            The output of SincNet model.
+        """
+        out = self.sincnet_layer(x)
+        return out
+
+
+
+class PyanNet(nn.Module):
+    """Pyannote model
+
+    Arguments
+    ---------
+    model_config : dict, str
+        consist of model parameters
+
+    Returns
+    -------
+    Pyannote model: class
+
+    """
+    def __init__(self,
+                 model_config,
+                 ):
+        super(PyanNet,self).__init__()
+
+        self.model_config = model_config
+
+        sincnet_filters = model_config["sincnet_filters"]
+        sincnet_stride = model_config["sincnet_stride"]
+        linear_blocks = model_config["linear_blocks"]
+
+        self.sincnet = SincNet(n_filters=sincnet_filters, stride_ = sincnet_stride)
+
+        if model_config["sequence_type"] == "lstm":
+            self.sequence_blocks = nn.LSTM(sincnet_filters[-1],
+                                           model_config["sequence_neuron"],
+                                           num_layers=model_config["sequence_nlayers"],
+                                           batch_first=True,
+                                           dropout=model_config["sequence_drop_out"],
+                                           bidirectional=model_config["sequence_bidirectional"],
+                                           )
+        elif model_config["sequence_type"] == "gru":
+            self.sequence_blocks = nn.GRU(sincnet_filters[-1],
+                                          model_config["sequence_neuron"],
+                                          num_layers=model_config["sequence_nlayers"],
+                                          batch_first=True,
+                                          dropout=model_config["sequence_drop_out"],
+                                          bidirectional=model_config["sequence_bidirectional"],
+                                          )
+        elif model_config["sequence_type"] == "attention":
+            self.sequence_blocks = nn.TransformerEncoderLayer(d_model=sincnet_filters[-1],
+                                                              dim_feedforward=model_config["sequence_neuron"],
+                                                              nhead=model_config["sequence_nlayers"],
+                                                              batch_first=True,
+                                                              dropout=model_config["sequence_drop_out"])
+        else:
+            raise ValueError("Model type is not valid!!!")
+
+
+        if model_config["sequence_bidirectional"]:
+            last_sequence_block = model_config["sequence_neuron"] * 2
+        else:
+            last_sequence_block = model_config["sequence_neuron"]
+
+
+        linear_blocks = [last_sequence_block] + linear_blocks
+        linears_list = nn.ModuleList()
+        for counter in range(len(linear_blocks) - 1):
+            linears_list.append(
+                nn.Linear(
+                    in_features=linear_blocks[counter],
+                    out_features=linear_blocks[counter+1],
+                    bias=True,
+                )
+            )
+        linears_list.append(nn.Sigmoid())
+        self.linears = nn.Sequential(*linears_list)
+
+
+    def forward(self, x):
+        """This method should implement forwarding operation in the Pyannote model.
+
+        Arguments
+        ---------
+        x : float (Tensor)
+            The input of Pyannote model.
+
+        Returns
+        -------
+        out : float (Tensor)
+            The output of Pyannote model.
+        """
+        x = torch.unsqueeze(x, 1)
+        x = self.sincnet(x)
+        x = x.permute(0,2,1)
+
+        if self.model_config["sequence_type"] == "attention":
+            x = self.sequence_blocks(x)
+        else:
+            x = self.sequence_blocks(x)[0]
+
+        out = self.linears(x)
+        return out