Update model & README.md

Author: Soohwan Kim

README.md

@@ -62,22 +62,25 @@ batch_size, sequence_length, dim = 3, 12345, 80
 cuda = torch.cuda.is_available()  
 device = torch.device('cuda' if cuda else 'cpu')
 
+criterion = nn.CTCLoss()
+
 inputs = torch.rand(batch_size, sequence_length, dim).to(device)
 input_lengths = torch.IntTensor([12345, 12300, 12000])
 targets = torch.LongTensor([[1, 3, 3, 3, 3, 3, 4, 5, 6, 2],
                             [1, 3, 3, 3, 3, 3, 4, 5, 2, 0],
                             [1, 3, 3, 3, 3, 3, 4, 2, 0, 0]]).to(device)
 target_lengths = torch.LongTensor([9, 8, 7])
 
-model = nn.DataParallel(Conformer(num_classes=10, input_dim=dim, 
-                                  encoder_dim=32, num_encoder_layers=3, 
-                                  decoder_dim=32)).to(device)
+model = Conformer(num_classes=10, 
+                  input_dim=dim, 
+                  encoder_dim=32, 
+                  num_encoder_layers=3)
 
 # Forward propagate
-outputs = model(inputs, input_lengths, targets, target_lengths)
+outputs, output_lengths = model(inputs, input_lengths)
 
-# Recognize input speech
-outputs = model.module.recognize(inputs, input_lengths)
+# Calculate CTC Loss
+loss = criterion(outputs.transpose(0, 1), targets, output_lengths, target_lengths)
 ```
 
 ## Troubleshoots and Contributing

conformer/decoder.py

@@ -15,8 +15,8 @@
 import torch
 import torch.nn as nn
 from torch import Tensor
+from typing import Tuple
 
-from conformer.decoder import DecoderRNNT
 from conformer.encoder import ConformerEncoder
 from conformer.modules import Linear
 
@@ -31,17 +31,13 @@ class Conformer(nn.Module):
         num_classes (int): Number of classification classes
         input_dim (int, optional): Dimension of input vector
         encoder_dim (int, optional): Dimension of conformer encoder
-        decoder_dim (int, optional): Dimension of conformer decoder
         num_encoder_layers (int, optional): Number of conformer blocks
-        num_decoder_layers (int, optional): Number of decoder layers
-        decoder_rnn_type (str, optional): type of RNN cell
         num_attention_heads (int, optional): Number of attention heads
         feed_forward_expansion_factor (int, optional): Expansion factor of feed forward module
         conv_expansion_factor (int, optional): Expansion factor of conformer convolution module
         feed_forward_dropout_p (float, optional): Probability of feed forward module dropout
         attention_dropout_p (float, optional): Probability of attention module dropout
         conv_dropout_p (float, optional): Probability of conformer convolution module dropout
-        decoder_dropout_p (float, optional): Probability of conformer decoder dropout
         conv_kernel_size (int or tuple, optional): Size of the convolving kernel
         half_step_residual (bool): Flag indication whether to use half step residual or not
 
@@ -58,20 +54,16 @@ def __init__(
             num_classes: int,
             input_dim: int = 80,
             encoder_dim: int = 512,
-            decoder_dim: int = 640,
             num_encoder_layers: int = 17,
-            num_decoder_layers: int = 1,
             num_attention_heads: int = 8,
             feed_forward_expansion_factor: int = 4,
             conv_expansion_factor: int = 2,
             input_dropout_p: float = 0.1,
             feed_forward_dropout_p: float = 0.1,
             attention_dropout_p: float = 0.1,
             conv_dropout_p: float = 0.1,
-            decoder_dropout_p: float = 0.1,
             conv_kernel_size: int = 31,
             half_step_residual: bool = True,
-            decoder_rnn_type: str = "lstm",
     ) -> None:
         super(Conformer, self).__init__()
         self.encoder = ConformerEncoder(
@@ -88,137 +80,27 @@ def __init__(
             conv_kernel_size=conv_kernel_size,
             half_step_residual=half_step_residual,
         )
-        self.decoder = DecoderRNNT(
-            num_classes=num_classes,
-            hidden_state_dim=decoder_dim,
-            output_dim=encoder_dim,
-            num_layers=num_decoder_layers,
-            rnn_type=decoder_rnn_type,
-            dropout_p=decoder_dropout_p,
-        )
         self.fc = Linear(encoder_dim << 1, num_classes, bias=False)
 
-    def set_encoder(self, encoder):
-        """ Setter for encoder """
-        self.encoder = encoder
-
-    def set_decoder(self, decoder):
-        """ Setter for decoder """
-        self.decoder = decoder
-
     def count_parameters(self) -> int:
         """ Count parameters of encoder """
-        num_encoder_parameters = self.encoder.count_parameters()
-        num_decoder_parameters = self.decoder.count_parameters()
-        return num_encoder_parameters + num_decoder_parameters
+        return self.encoder.count_parameters()
 
     def update_dropout(self, dropout_p) -> None:
         """ Update dropout probability of model """
         self.encoder.update_dropout(dropout_p)
-        self.decoder.update_dropout(dropout_p)
-
-    def joint(self, encoder_outputs: Tensor, decoder_outputs: Tensor) -> Tensor:
-        """
-        Joint `encoder_outputs` and `decoder_outputs`.
-
-        Args:
-            encoder_outputs (torch.FloatTensor): A output sequence of encoder. `FloatTensor` of size
-                ``(batch, seq_length, dimension)``
-            decoder_outputs (torch.FloatTensor): A output sequence of decoder. `FloatTensor` of size
-                ``(batch, seq_length, dimension)``
-
-        Returns:
-            * outputs (torch.FloatTensor): outputs of joint `encoder_outputs` and `decoder_outputs`..
-        """
-        if encoder_outputs.dim() == 3 and decoder_outputs.dim() == 3:
-            input_length = encoder_outputs.size(1)
-            target_length = decoder_outputs.size(1)
 
-            encoder_outputs = encoder_outputs.unsqueeze(2)
-            decoder_outputs = decoder_outputs.unsqueeze(1)
-
-            encoder_outputs = encoder_outputs.repeat([1, 1, target_length, 1])
-            decoder_outputs = decoder_outputs.repeat([1, input_length, 1, 1])
-
-        outputs = torch.cat((encoder_outputs, decoder_outputs), dim=-1)
-        outputs = self.fc(outputs)
-
-        return outputs
-
-    def forward(
-            self,
-            inputs: Tensor,
-            input_lengths: Tensor,
-            targets: Tensor,
-            target_lengths: Tensor
-    ) -> Tensor:
+    def forward(self, inputs: Tensor, input_lengths: Tensor) -> Tuple[Tensor, Tensor]:
         """
         Forward propagate a `inputs` and `targets` pair for training.
 
         Args:
             inputs (torch.FloatTensor): A input sequence passed to encoder. Typically for inputs this will be a padded
                 `FloatTensor` of size ``(batch, seq_length, dimension)``.
             input_lengths (torch.LongTensor): The length of input tensor. ``(batch)``
-            targets (torch.LongTensr): A target sequence passed to decoder. `IntTensor` of size ``(batch, seq_length)``
-            target_lengths (torch.LongTensor): The length of target tensor. ``(batch)``
 
         Returns:
             * predictions (torch.FloatTensor): Result of model predictions.
         """
-        encoder_outputs, _ = self.encoder(inputs, input_lengths)
-        decoder_outputs, _ = self.decoder(targets, target_lengths)
-        outputs = self.joint(encoder_outputs, decoder_outputs)
-        return outputs
-
-    @torch.no_grad()
-    def decode(self, encoder_output: Tensor, max_length: int) -> Tensor:
-        """
-        Decode `encoder_outputs`.
-
-        Args:
-            encoder_output (torch.FloatTensor): A output sequence of encoder. `FloatTensor` of size
-                ``(seq_length, dimension)``
-            max_length (int): max decoding time step
-
-        Returns:
-            * predicted_log_probs (torch.FloatTensor): Log probability of model predictions.
-        """
-        pred_tokens, hidden_state = list(), None
-        decoder_input = encoder_output.new_tensor([[self.decoder.sos_id]], dtype=torch.long)
-
-        for t in range(max_length):
-            decoder_output, hidden_state = self.decoder(decoder_input, hidden_states=hidden_state)
-            step_output = self.joint(encoder_output[t].view(-1), decoder_output.view(-1))
-            step_output = step_output.softmax(dim=0)
-            pred_token = step_output.argmax(dim=0)
-            pred_token = int(pred_token.item())
-            pred_tokens.append(pred_token)
-            decoder_input = step_output.new_tensor([[pred_token]], dtype=torch.long)
-
-        return torch.LongTensor(pred_tokens)
-
-    @torch.no_grad()
-    def recognize(self, inputs: Tensor, input_lengths: Tensor):
-        """
-        Recognize input speech. This method consists of the forward of the encoder and the decode() of the decoder.
-
-        Args:
-            inputs (torch.FloatTensor): A input sequence passed to encoder. Typically for inputs this will be a padded
-                `FloatTensor` of size ``(batch, seq_length, dimension)``.
-            input_lengths (torch.LongTensor): The length of input tensor. ``(batch)``
-
-        Returns:
-            * predictions (torch.FloatTensor): Result of model predictions.
-        """
-        outputs = list()
-
-        encoder_outputs, output_lengths = self.encoder(inputs, input_lengths)
-        max_length = encoder_outputs.size(1)
-
-        for encoder_output in encoder_outputs:
-            decoded_seq = self.decode(encoder_output, max_length)
-            outputs.append(decoded_seq)
-
-        outputs = torch.stack(outputs, dim=1).transpose(0, 1)
-
-        return outputs
+        encoder_outputs, encoder_output_lengths = self.encoder(inputs, input_lengths)
+        return encoder_outputs, encoder_output_lengths