wip: commit before going on vacation :)

1 files changed, 150 insertions, 0 deletions

diff --git a/swr2_asr/model_deep_speech.py b/swr2_asr/model_deep_speech.py
new file mode 100644
index 0000000..ea0b667
--- /dev/null
+++ b/swr2_asr/model_deep_speech.py
@@ -0,0 +1,150 @@
+from torch import nn
+import torch.nn.functional as F
+
+
+class CNNLayerNorm(nn.Module):
+    """Layer normalization built for cnns input"""
+
+    def __init__(self, n_feats: int):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(n_feats)
+
+    def forward(self, data):
+        """x (batch, channel, feature, time)"""
+        data = data.transpose(2, 3).contiguous()  # (batch, channel, time, feature)
+        data = self.layer_norm(data)
+        return data.transpose(2, 3).contiguous()  # (batch, channel, feature, time)
+
+
+class ResidualCNN(nn.Module):
+    """Residual CNN inspired by https://arxiv.org/pdf/1603.05027.pdf"""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel: int,
+        stride: int,
+        dropout: float,
+        n_feats: int,
+    ):
+        super().__init__()
+
+        self.cnn1 = nn.Conv2d(
+            in_channels, out_channels, kernel, stride, padding=kernel // 2
+        )
+        self.cnn2 = nn.Conv2d(
+            out_channels,
+            out_channels,
+            kernel,
+            stride,
+            padding=kernel // 2,
+        )
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.layer_norm1 = CNNLayerNorm(n_feats)
+        self.layer_norm2 = CNNLayerNorm(n_feats)
+
+    def forward(self, data):
+        """x (batch, channel, feature, time)"""
+        residual = data  # (batch, channel, feature, time)
+        data = self.layer_norm1(data)
+        data = F.gelu(data)
+        data = self.dropout1(data)
+        data = self.cnn1(data)
+        data = self.layer_norm2(data)
+        data = F.gelu(data)
+        data = self.dropout2(data)
+        data = self.cnn2(data)
+        data += residual
+        return data  # (batch, channel, feature, time)
+
+
+class BidirectionalGRU(nn.Module):
+    """BIdirectional GRU with Layer Normalization and Dropout"""
+
+    def __init__(
+        self,
+        rnn_dim: int,
+        hidden_size: int,
+        dropout: float,
+        batch_first: bool,
+    ):
+        super().__init__()
+
+        self.bi_gru = nn.GRU(
+            input_size=rnn_dim,
+            hidden_size=hidden_size,
+            num_layers=1,
+            batch_first=batch_first,
+            bidirectional=True,
+        )
+        self.layer_norm = nn.LayerNorm(rnn_dim)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, data):
+        """data (batch, time, feature)"""
+        data = self.layer_norm(data)
+        data = F.gelu(data)
+        data = self.dropout(data)
+        data, _ = self.bi_gru(data)
+        return data
+
+
+class SpeechRecognitionModel(nn.Module):
+    """Speech Recognition Model Inspired by DeepSpeech 2"""
+
+    def __init__(
+        self,
+        n_cnn_layers: int,
+        n_rnn_layers: int,
+        rnn_dim: int,
+        n_class: int,
+        n_feats: int,
+        stride: int = 2,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        n_feats //= 2
+        self.cnn = nn.Conv2d(1, 32, 3, stride=stride, padding=3 // 2)
+        # n residual cnn layers with filter size of 32
+        self.rescnn_layers = nn.Sequential(
+            *[
+                ResidualCNN(
+                    32, 32, kernel=3, stride=1, dropout=dropout, n_feats=n_feats
+                )
+                for _ in range(n_cnn_layers)
+            ]
+        )
+        self.fully_connected = nn.Linear(n_feats * 32, rnn_dim)
+        self.birnn_layers = nn.Sequential(
+            *[
+                BidirectionalGRU(
+                    rnn_dim=rnn_dim if i == 0 else rnn_dim * 2,
+                    hidden_size=rnn_dim,
+                    dropout=dropout,
+                    batch_first=i == 0,
+                )
+                for i in range(n_rnn_layers)
+            ]
+        )
+        self.classifier = nn.Sequential(
+            nn.Linear(rnn_dim * 2, rnn_dim),  # birnn returns rnn_dim*2
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(rnn_dim, n_class),
+        )
+
+    def forward(self, data):
+        """data (batch, channel, feature, time)"""
+        data = self.cnn(data)
+        data = self.rescnn_layers(data)
+        sizes = data.size()
+        data = data.view(
+            sizes[0], sizes[1] * sizes[2], sizes[3]
+        )  # (batch, feature, time)
+        data = data.transpose(1, 2)  # (batch, time, feature)
+        data = self.fully_connected(data)
+        data = self.birnn_layers(data)
+        data = self.classifier(data)
+        return data