Add MHSA module, Conformer block and encoder with relative PE

i6_models/assemblies/conformer/__init__.py

@@ -1,2 +1,3 @@
 from .conformer_v1 import *
 from .conformer_v2 import *
+from .conformer_rel_pos_v1 import *

i6_models/assemblies/conformer/conformer_rel_pos_v1.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+__all__ = [
+    "ConformerRelPosBlockV1Config",
+    "ConformerRelPosEncoderV1Config",
+    "ConformerRelPosBlockV1",
+    "ConformerRelPosEncoderV1",
+]
+
+import torch
+from torch import nn
+from dataclasses import dataclass, field
+from typing import List
+
+from i6_models.config import ModelConfiguration, ModuleFactoryV1
+from i6_models.parts.conformer import (
+    ConformerConvolutionV1,
+    ConformerConvolutionV1Config,
+    ConformerMHSARelPosV1,
+    ConformerMHSARelPosV1Config,
+    ConformerPositionwiseFeedForwardV1,
+    ConformerPositionwiseFeedForwardV1Config,
+)
+from i6_models.assemblies.conformer import ConformerEncoderV2
+
+
+@dataclass
+class ConformerRelPosBlockV1Config(ModelConfiguration):
+    """
+    Attributes:
+        ff_cfg: Configuration for ConformerPositionwiseFeedForwardV1
+        mhsa_cfg: Configuration for ConformerMHSARelPosV1
+        conv_cfg: Configuration for ConformerConvolutionV1
+        modules: List of modules to use for ConformerRelPosBlockV1,
+            "ff" for feed forward module, "mhsa" for multi-head self attention module, "conv" for conv module
+        scales: List of scales to apply to the module outputs before the residual connection
+    """
+
+    # nested configurations
+    ff_cfg: ConformerPositionwiseFeedForwardV1Config
+    mhsa_cfg: ConformerMHSARelPosV1Config
+    conv_cfg: ConformerConvolutionV1Config
+    modules: List[str] = field(default_factory=lambda: ["ff", "mhsa", "conv", "ff"])
+    scales: List[float] = field(default_factory=lambda: [0.5, 1.0, 1.0, 0.5])
+
+    def __post__init__(self):
+        super().__post_init__()
+        assert len(self.modules) == len(self.scales), "modules and scales must have same length"
+        for module_name in self.modules:
+            assert module_name in ["ff", "mhsa", "conv"], "module not supported"
+
+
+class ConformerRelPosBlockV1(nn.Module):
+    """
+    Conformer block module, modifications compared to ConformerBlockV1:
+    - uses ConfomerMHSARelPosV1 as MHSA module
+    - enable constructing the block with self-defined module_list as ConformerBlockV2
+    """
+
+    def __init__(self, cfg: ConformerRelPosBlockV1Config):
+        """
+        :param cfg: conformer block configuration with subunits for the different conformer parts
+        """
+        super().__init__()
+
+        modules = []
+        for module_name in cfg.modules:
+            if module_name == "ff":
+                modules.append(ConformerPositionwiseFeedForwardV1(cfg=cfg.ff_cfg))
+            elif module_name == "mhsa":
+                modules.append(ConformerMHSARelPosV1(cfg=cfg.mhsa_cfg))
+            elif module_name == "conv":
+                modules.append(ConformerConvolutionV1(model_cfg=cfg.conv_cfg))
+            else:
+                raise NotImplementedError
+
+        self.module_list = nn.ModuleList(modules)
+        self.scales = cfg.scales
+        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 1 defines positions within the sequence and 0 outside, shape: [B, T]
+        :return: torch.Tensor of shape [B, T, F]
+        """
+        for scale, module in zip(self.scales, self.module_list):
+            if isinstance(module, ConformerMHSARelPosV1):
+                x = scale * module(x, sequence_mask) + x
+            else:
+                x = scale * module(x) + x
+        x = self.final_layer_norm(x)  #  [B, T, F]
+        return x
+
+
+@dataclass
+class ConformerRelPosEncoderV1Config(ModelConfiguration):
+    """
+    Attributes:
+        num_layers: Number of conformer layers in the conformer encoder
+        frontend: A pair of ConformerFrontend and corresponding config
+        block_cfg: Configuration for ConformerRelPosBlockV1
+    """
+
+    num_layers: int
+
+    # nested configurations
+    frontend: ModuleFactoryV1
+    block_cfg: ConformerRelPosBlockV1Config
+
+
+class ConformerRelPosEncoderV1(ConformerEncoderV2):
+    """
+    Modifications compared to ConformerEncoderV2:
+    - uses multi-headed self-attention with Shaw's relative positional encoding
+    """
+
+    def __init__(self, cfg: ConformerRelPosEncoderV1Config):
+        """
+        :param cfg: conformer encoder configuration with subunits for frontend and conformer blocks
+        """
+        super().__init__(cfg)
+
+        self.frontend = cfg.frontend()
+        self.module_list = torch.nn.ModuleList([ConformerRelPosBlockV1(cfg.block_cfg) for _ in range(cfg.num_layers)])

i6_models/parts/conformer/__init__.py

@@ -1,4 +1,5 @@
 from .convolution import *
 from .feedforward import *
 from .mhsa import *
+from .mhsa_rel_pos import *
 from .norm import *

i6_models/parts/conformer/convolution.py

@@ -20,13 +20,15 @@ class ConformerConvolutionV1Config(ModelConfiguration):
         dropout: dropout probability
         activation: activation function applied after normalization
         norm: normalization layer with input of shape [N,C,T]
+        broadcast_dropout: whether to broadcast dropout on the feature axis to time axis
     """
 
     channels: int
     kernel_size: int
     dropout: float
     activation: Union[nn.Module, Callable[[torch.Tensor], torch.Tensor]]
     norm: Union[nn.Module, Callable[[torch.Tensor], torch.Tensor]]
+    broadcast_dropout: bool = False
 
     def check_valid(self):
         assert self.kernel_size % 2 == 1, "ConformerConvolutionV1 only supports odd kernel sizes"
@@ -62,7 +64,8 @@ def __init__(self, model_cfg: ConformerConvolutionV1Config):
         self.pointwise_conv2 = nn.Linear(in_features=model_cfg.channels, out_features=model_cfg.channels)
         self.layer_norm = nn.LayerNorm(model_cfg.channels)
         self.norm = deepcopy(model_cfg.norm)
-        self.dropout = nn.Dropout(model_cfg.dropout)
+        self.dropout = nn.Dropout1d(model_cfg.dropout) if model_cfg.broadcast_dropout else nn.Dropout(model_cfg.dropout)
+        self.broadcast_dropout = model_cfg.broadcast_dropout
         self.activation = model_cfg.activation
 
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
@@ -84,4 +87,9 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         tensor = self.activation(tensor)
         tensor = self.pointwise_conv2(tensor)
 
-        return self.dropout(tensor)
+        if self.broadcast_dropout:
+            tensor = self.dropout(tensor.transpose(1, 2)).transpose(1, 2)
+        else:
+            tensor = self.dropout(tensor)
+
+        return tensor

i6_models/parts/conformer/feedforward.py

@@ -18,12 +18,14 @@ class ConformerPositionwiseFeedForwardV1Config(ModelConfiguration):
         input_dim: input dimension
         hidden_dim: hidden dimension (normally set to 4*input_dim as suggested by the paper)
         dropout: dropout probability
+        broadcast_dropout: whether to broadcast dropout on the feature axis to time axis
         activation: activation function
     """
 
     input_dim: int
     hidden_dim: int
     dropout: float
+    broadcast_dropout: bool = False
     activation: Callable[[torch.Tensor], torch.Tensor] = nn.functional.silu
 
 
@@ -40,6 +42,7 @@ def __init__(self, cfg: ConformerPositionwiseFeedForwardV1Config):
         self.activation = cfg.activation
         self.linear_out = nn.Linear(in_features=cfg.hidden_dim, out_features=cfg.input_dim, bias=True)
         self.dropout = cfg.dropout
+        self.broadcast_dropout = cfg.broadcast_dropout
 
     def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         """
@@ -49,7 +52,17 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
         tensor = self.layer_norm(tensor)
         tensor = self.linear_ff(tensor)  # [B,T,F]
         tensor = self.activation(tensor)  # [B,T,F]
-        tensor = nn.functional.dropout(tensor, p=self.dropout, training=self.training)  # [B,T,F]
-        tensor = self.linear_out(tensor)  # [B,T,F]
-        tensor = nn.functional.dropout(tensor, p=self.dropout, training=self.training)  # [B,T,F]
+
+        if self.broadcast_dropout:
+            tensor = nn.functional.dropout1d(tensor.transpose(1, 2), p=self.dropout, training=self.training).transpose(
+                1, 2
+            )
+            tensor = self.linear_out(tensor)
+            tensor = nn.functional.dropout1d(tensor.transpose(1, 2), p=self.dropout, training=self.training).transpose(
+                1, 2
+            )
+        else:
+            tensor = nn.functional.dropout(tensor, p=self.dropout, training=self.training)  # [B,T,F]
+            tensor = self.linear_out(tensor)  # [B,T,F]
+            tensor = nn.functional.dropout(tensor, p=self.dropout, training=self.training)  # [B,T,F]
         return tensor

i6_models/parts/conformer/mhsa.py

@@ -16,12 +16,14 @@ class ConformerMHSAV1Config(ModelConfiguration):
         num_att_heads: number of attention heads
         att_weights_dropout: attention weights dropout
         dropout: multi-headed self attention output dropout
+        broadcast_dropout: whether to broadcast dropout on the feature axis to time axis
     """
 
     input_dim: int
     num_att_heads: int
     att_weights_dropout: float
     dropout: float
+    broadcast_dropout: bool = False
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -42,6 +44,7 @@ def __init__(self, cfg: ConformerMHSAV1Config):
             cfg.input_dim, cfg.num_att_heads, dropout=cfg.att_weights_dropout, batch_first=True
         )
         self.dropout = cfg.dropout
+        self.broadcast_dropout = cfg.broadcast_dropout
 
     def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> torch.Tensor:
         """
@@ -57,6 +60,14 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to
         output_tensor, _ = self.mhsa(
             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]
+
+        if self.broadcast_dropout:
+            output_tensor = torch.nn.functional.dropout1d(
+                output_tensor.transpose(1, 2), p=self.dropout, training=self.training
+            ).transpose(1, 2)
+        else:
+            output_tensor = torch.nn.functional.dropout(
+                output_tensor, p=self.dropout, training=self.training
+            )  # [B,T,F]
 
         return output_tensor