Conformer Assembly (#8)

Add Conformer Assembly

---------

Co-authored-by: Judyxujj <[email protected]>
Co-authored-by: Albert Zeyer <[email protected]>
Co-authored-by: pzheng <[email protected]>
Co-authored-by: Mohammad Zeineldeen <[email protected]>
Co-authored-by: Nick Rossenbach <[email protected]>
Co-authored-by: vieting <[email protected]>
Co-authored-by: michelwi <[email protected]>

Author: 8 people

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