Merge branch 'main' into zeineldeen_att_decoder

Author: mmz33

i6_models/__init__.py

@@ -0,0 +1 @@
+

i6_models/assemblies/conformer/__init__.py

@@ -0,0 +1 @@
+from .conformer_v1 import *

i6_models/assemblies/conformer/conformer_v1.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+__all__ = ["ConformerBlockV1Config", "ConformerEncoderV1Config", "ConformerBlockV1", "ConformerEncoderV1"]
+
+import torch
+from torch import nn
+from dataclasses import dataclass
+from typing import Tuple
+
+from i6_models.config import ModelConfiguration, ModuleFactoryV1
+from i6_models.parts.conformer import (
+    ConformerConvolutionV1,
+    ConformerConvolutionV1Config,
+    ConformerMHSAV1,
+    ConformerMHSAV1Config,
+    ConformerPositionwiseFeedForwardV1,
+    ConformerPositionwiseFeedForwardV1Config,
+)
+
+
+@dataclass
+class ConformerBlockV1Config(ModelConfiguration):
+    """
+    Attributes:
+        ff_cfg: Configuration for ConformerPositionwiseFeedForwardV1
+        mhsa_cfg: Configuration for ConformerMHSAV1
+        conv_cfg: Configuration for ConformerConvolutionV1
+    """
+
+    # nested configurations
+    ff_cfg: ConformerPositionwiseFeedForwardV1Config
+    mhsa_cfg: ConformerMHSAV1Config
+    conv_cfg: ConformerConvolutionV1Config
+
+
+class ConformerBlockV1(nn.Module):
+    """
+    Conformer block module
+    """
+
+    def __init__(self, cfg: ConformerBlockV1Config):
+        """
+        :param cfg: conformer block configuration with subunits for the different conformer parts
+        """
+        super().__init__()
+        self.ff1 = ConformerPositionwiseFeedForwardV1(cfg=cfg.ff_cfg)
+        self.mhsa = ConformerMHSAV1(cfg=cfg.mhsa_cfg)
+        self.conv = ConformerConvolutionV1(model_cfg=cfg.conv_cfg)
+        self.ff2 = ConformerPositionwiseFeedForwardV1(cfg=cfg.ff_cfg)
+        self.final_layer_norm = torch.nn.LayerNorm(cfg.ff_cfg.input_dim)
+
+    def forward(self, x: torch.Tensor, /, sequence_mask: torch.Tensor) -> torch.Tensor:
+        """
+        :param x: input tensor of shape [B, T, F]
+        :param sequence_mask: mask tensor where 0 defines positions within the sequence and 1 outside, shape: [B, T]
+        :return: torch.Tensor of shape [B, T, F]
+        """
+        x = 0.5 * self.ff1(x) + x  #  [B, T, F]
+        x = self.mhsa(x, sequence_mask) + x  #  [B, T, F]
+        x = self.conv(x) + x  #  [B, T, F]
+        x = 0.5 * self.ff2(x) + x  #  [B, T, F]
+        x = self.final_layer_norm(x)  #  [B, T, F]
+        return x
+
+
+@dataclass
+class ConformerEncoderV1Config(ModelConfiguration):
+    """
+    Attributes:
+        num_layers: Number of conformer layers in the conformer encoder
+        frontend: A pair of ConformerFrontend and corresponding config
+        block_cfg: Configuration for ConformerBlockV1
+    """
+
+    num_layers: int
+
+    # nested configurations
+    frontend: ModuleFactoryV1
+    block_cfg: ConformerBlockV1Config
+
+
+class ConformerEncoderV1(nn.Module):
+    """
+    Implementation of the convolution-augmented Transformer (short Conformer), as in the original publication.
+    The model consists of a frontend and a stack of N conformer blocks.
+    C.f. https://arxiv.org/pdf/2005.08100.pdf
+    """
+
+    def __init__(self, cfg: ConformerEncoderV1Config):
+        """
+        :param cfg: conformer encoder configuration with subunits for frontend and conformer blocks
+        """
+        super().__init__()
+
+        self.frontend = cfg.frontend()
+        self.module_list = torch.nn.ModuleList([ConformerBlockV1(cfg.block_cfg) for _ in range(cfg.num_layers)])
+
+    def forward(self, data_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        :param data_tensor: input tensor of shape [B, T', F]
+        :param sequence_mask: mask tensor where 0 defines positions within the sequence and 1 outside, shape: [B, T']
+        :return: (output, out_seq_mask)
+            where output is torch.Tensor of shape [B, T, F'],
+            out_seq_mask is a torch.Tensor of shape [B, T]
+
+        F: input feature dim, F': internal and output feature dim
+        T': data time dim, T: down-sampled time dim (internal time dim)
+        """
+        x, sequence_mask = self.frontend(data_tensor, sequence_mask)  # [B, T, F']
+        for module in self.module_list:
+            x = module(x, sequence_mask)  # [B, T, F']
+
+        return x, sequence_mask

i6_models/parts/conformer/mhsa.py

@@ -2,10 +2,10 @@
 
 __all__ = ["ConformerMHSAV1", "ConformerMHSAV1Config"]
 from dataclasses import dataclass
-from typing import Optional
 import torch
 
 from i6_models.config import ModelConfiguration
+from i6_models.util import compat
 
 
 @dataclass
@@ -43,17 +43,19 @@ def __init__(self, cfg: ConformerMHSAV1Config):
         )
         self.dropout = cfg.dropout
 
-    def forward(self, input_tensor: torch.Tensor, key_padding_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+    def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> torch.Tensor:
         """
         Apply layer norm and multi-head self attention and dropout
-        :param Optional[torch.Tensor] key_padding_mask: could be a binary or float mask of shape (B, T)
+
+        :param input_tensor: Input to the self attention of shape (B, T, F)
+        :param sequence_mask: bool mask of shape (B, T), True signals within sequence, False outside, will be inverted to match the torch.nn.MultiheadAttention module
         which will be applied/added to dot product, used to mask padded key positions out
         """
-
+        inv_sequence_mask = compat.logical_not(sequence_mask)
         output_tensor = self.layernorm(input_tensor)  # [B,T,F]
 
         output_tensor, _ = self.mhsa(
-            output_tensor, output_tensor, output_tensor, key_padding_mask=key_padding_mask, need_weights=False
+            output_tensor, output_tensor, output_tensor, key_padding_mask=inv_sequence_mask, need_weights=False
         )  # [B,T,F]
         output_tensor = torch.nn.functional.dropout(output_tensor, p=self.dropout, training=self.training)  # [B,T,F]
 

i6_models/parts/frontend/README.md

@@ -0,0 +1,12 @@
+# Different front-ends for acoustic encoders
+
+### Contributing
+
+If you want to add your own front-end:
+
+- Normally two classes are required. A config class and a model class
+- `Config` class inherits from `ModelConfiguration`
+- `Model` class inherits from `nn.Module` from `torch`
+- `forward(tensor: torch.Tensor, sequence_mask: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]`
+- `sequence_mask` is a boolean tensor where `True` means is inside the sequence and `False` is masked.
+- Please add tests

i6_models/parts/frontend/__init__.py

@@ -0,0 +1,48 @@
+from typing import Tuple, Union
+
+import torch
+from torch import nn
+from torch.nn import functional
+
+
+def get_same_padding(input_size: Union[int, Tuple[int, ...]]) -> Union[int, Tuple[int, ...]]:
+    """
+    get padding in order to not reduce the time dimension
+
+    :param input_size:
+    :return:
+    """
+    if isinstance(input_size, int):
+        return (input_size - 1) // 2
+    elif isinstance(input_size, tuple):
+        return tuple((s - 1) // 2 for s in input_size)
+    else:
+        raise TypeError(f"unexpected size type {type(input_size)}")
+
+
+def mask_pool(seq_mask: torch.Tensor, *, kernel_size: int, stride: int, padding: int) -> torch.Tensor:
+    """
+    apply strides to the masking
+
+    :param seq_mask: [B,T]
+    :param kernel_size:
+    :param stride:
+    :param padding:
+    :return: [B,T'] using maxpool
+    """
+    if stride == 1 and 2 * padding == kernel_size - 1:
+        return seq_mask
+
+    seq_mask = seq_mask.float()
+    seq_mask = torch.unsqueeze(seq_mask, 1)  # [B,1,T]
+    seq_mask = nn.functional.max_pool1d(seq_mask, kernel_size, stride, padding)  # [B,1,T']
+    seq_mask = torch.squeeze(seq_mask, 1)  # [B,T']
+    seq_mask = seq_mask.bool()
+    return seq_mask
+
+
+def calculate_output_dim(in_dim: int, *, filter_size: int, stride: int, padding: int) -> int:
+    def ceildiv(a: int, b: int):
+        return -(-a // b)
+
+    return ceildiv(in_dim + 2 * padding - (filter_size - 1) * 1, stride)