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/mhsa_rel_pos.py

@@ -0,0 +1,170 @@
+from __future__ import annotations
+
+__all__ = ["ConformerMHSARelPosV1", "ConformerMHSARelPosV1Config"]
+
+import math
+from dataclasses import dataclass
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from i6_models.config import ModelConfiguration
+from i6_models.util import compat
+
+
+@dataclass
+class ConformerMHSARelPosV1Config(ModelConfiguration):
+    """
+    Attributes:
+        input_dim: input dim and total dimension for query/key and value projections, should be divisible by `num_att_heads`
+        num_att_heads: number of attention heads
+        rel_pos_clip: maximal relative postion for embedding
+        att_weights_dropout: attention weights dropout
+        dropout: multi-headed self attention output dropout
+    """
+
+    input_dim: int
+    num_att_heads: int
+    rel_pos_clip: int
+    att_weights_dropout: float
+    dropout: float
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        assert self.input_dim % self.num_att_heads == 0, "input_dim must be divisible by num_att_heads"
+
+
+class ConformerMHSARelPosV1(nn.Module):
+    """
+    Conformer multi-headed self-attention module with relative positional encoding proposed by Shaw et al. (cf. https://arxiv.org/abs/1803.02155)
+    """
+
+    def __init__(self, cfg: ConformerMHSARelPosV1Config):
+
+        super().__init__()
+
+        self.layernorm = nn.LayerNorm(cfg.input_dim)
+
+        self.embed_dim = cfg.input_dim
+        self.num_heads = cfg.num_att_heads
+        self.embed_dim_per_head = self.embed_dim // self.num_heads
+
+        self.rel_pos_clip = cfg.rel_pos_clip
+
+        self.att_weights_dropout = cfg.att_weights_dropout
+
+        assert self.embed_dim % self.num_heads == 0, "embed_dim must be divisible by num_heads"
+
+        # projection matrices
+        self.q_proj_weight = nn.parameter.Parameter(torch.empty((self.embed_dim, self.embed_dim)))
+        self.k_proj_weight = nn.parameter.Parameter(torch.empty((self.embed_dim, self.embed_dim)))
+        self.v_proj_weight = nn.parameter.Parameter(torch.empty((self.embed_dim, self.embed_dim)))
+
+        self.in_proj_bias = nn.parameter.Parameter(torch.empty(3 * self.embed_dim))
+
+        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True)
+
+        if self.rel_pos_clip:
+            self.rel_pos_embeddings = nn.parameter.Parameter(
+                torch.empty(self.rel_pos_clip * 2 + 1, self.embed_dim // self.num_heads)
+            )
+        else:
+            self.register_parameter("rel_pos_embeddings", None)
+
+        self.dropout = cfg.dropout
+
+        self._reset_parameters()  # initialize parameters
+
+    def _reset_parameters(self):
+        nn.init.xavier_uniform_(self.q_proj_weight)
+        nn.init.xavier_uniform_(self.k_proj_weight)
+        nn.init.xavier_uniform_(self.v_proj_weight)
+
+        if self.rel_pos_clip:
+            nn.init.normal_(self.rel_pos_embeddings)
+
+        nn.init.constant_(self.in_proj_bias, 0.0)
+
+    def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> torch.Tensor:
+        """
+        Apply layer norm and multi-head self attention and dropout
+
+        :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
+        """
+        output_tensor = self.layernorm(input_tensor)  # [B,T,F]
+
+        time_dim_size = output_tensor.shape[1]
+        batch_dim_size = output_tensor.shape[0]
+
+        # attention mask
+        inv_sequence_mask = compat.logical_not(sequence_mask)  # [B, T]
+        mask = (
+            torch.zeros_like(inv_sequence_mask, dtype=input_tensor.dtype)
+            .masked_fill(inv_sequence_mask, float("-inf"))
+            .view(batch_dim_size, 1, 1, time_dim_size)
+            .expand(-1, self.num_heads, -1, -1)
+        )  # [B, #heads, 1, T']
+
+        # query, key and value sequences
+        bias_k, bias_q, bias_v = self.in_proj_bias.chunk(3)
+
+        query_seq = F.linear(output_tensor, self.q_proj_weight, bias_q)  # [B, T, #heads * F']
+        key_seq = F.linear(output_tensor, self.k_proj_weight, bias_k)
+        value_seq = F.linear(output_tensor, self.v_proj_weight, bias_v)
+
+        q1 = query_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head).transpose(
+            1, 2
+        )  # [B, #heads, T, F']
+        k_t = key_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head).permute(
+            0, 2, 3, 1
+        )  # [B, #heads, F', T']
+        # attention between query and key sequences
+        attn1 = torch.matmul(q1, k_t)  # [B, #heads, T, T']i
+
+        if self.rel_pos_clip:
+            q2 = (
+                query_seq.transpose(0, 1)
+                .contiguous()
+                .view(time_dim_size, batch_dim_size * self.num_heads, self.embed_dim_per_head)
+            )  # [T, B*#heads, F']
+
+            range_vec_q = torch.arange(time_dim_size, device=input_tensor.device)
+            range_vec_k = torch.arange(time_dim_size, device=input_tensor.device)
+
+            distance_mat = range_vec_k[None, :] - range_vec_q[:, None]
+            distance_mat_clipped = torch.clamp(distance_mat, -self.rel_pos_clip, self.rel_pos_clip)
+
+            final_mat = distance_mat_clipped + self.rel_pos_clip
+            # relative positional embeddings
+            rel_pos_embeddings = self.rel_pos_embeddings[final_mat]  # [T, T', F']
+
+            # attention between query sequence and relative positional embeddings
+            attn2 = torch.matmul(q2, rel_pos_embeddings.transpose(1, 2)).transpose(0, 1)  # [B*#heads, T, T']
+            attn2 = attn2.contiguous().view(
+                batch_dim_size, self.num_heads, time_dim_size, time_dim_size
+            )  # [B, #heads, T, T']
+
+            attn = (attn1 + attn2 + mask) * (math.sqrt(1.0 / float(self.embed_dim_per_head)))  # [B, #heads, T, T']
+        else:
+            attn = (attn1 + mask) * (math.sqrt(1.0 / float(self.embed_dim_per_head)))  # [B, #heads, T, T']
+
+        # softmax and dropout
+        attn_output_weights = F.dropout(
+            F.softmax(attn, dim=-1), p=self.att_weights_dropout, training=self.training
+        )  # [B, #heads, T, T']
+
+        # sequence of weighted sums over value sequence
+        v = value_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head).transpose(
+            1, 2
+        )  # [B, #heads, T', F']
+        attn_output = (
+            torch.matmul(attn_output_weights, v).transpose(1, 2).contiguous().view(batch_dim_size, -1, self.embed_dim)
+        )  # [B, T, F]
+
+        output_tensor = self.out_proj(attn_output)
+
+        output_tensor = F.dropout(output_tensor, p=self.dropout, training=self.training)  # [B,T,F]
+
+        return output_tensor