871 lines
33 KiB
Python
871 lines
33 KiB
Python
"""
|
|
Implementation of the MossFormer2 block
|
|
This source code is rewritten by Shengkui Zhao based on https://github.com/lucidrains/FLASH-pytorch
|
|
"""
|
|
|
|
import math
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn, einsum
|
|
from torchinfo import summary
|
|
from einops import rearrange
|
|
from rotary_embedding_torch import RotaryEmbedding
|
|
|
|
from models.mossformer2_ss.conv_module import ConvModule, GLU, FFConvM_Dilated
|
|
from models.mossformer2_ss.fsmn import UniDeepFsmn, UniDeepFsmn_dilated
|
|
from models.mossformer2_ss.layer_norm import CLayerNorm, GLayerNorm, GlobLayerNorm, ILayerNorm
|
|
|
|
# Functions
|
|
|
|
def identity(t, *args, **kwargs):
|
|
"""Identity function, returns the input tensor unchanged."""
|
|
return t
|
|
|
|
def append_dims(x, num_dims):
|
|
"""Appends extra dimensions to the input tensor `x`."""
|
|
if num_dims <= 0:
|
|
return x
|
|
return x.view(*x.shape, *((1,) * num_dims))
|
|
|
|
def exists(val):
|
|
"""Checks if a value exists (is not None)."""
|
|
return val is not None
|
|
|
|
def default(val, d):
|
|
"""Returns the value if it exists, otherwise returns the default `d`."""
|
|
return val if exists(val) else d
|
|
|
|
def padding_to_multiple_of(n, mult):
|
|
"""Returns the padding required to make `n` a multiple of `mult`."""
|
|
remainder = n % mult
|
|
if remainder == 0:
|
|
return 0
|
|
return mult - remainder
|
|
|
|
# ScaleNorm Layer
|
|
|
|
class ScaleNorm(nn.Module):
|
|
"""
|
|
ScaleNorm Layer: A variant of LayerNorm that scales the input tensor
|
|
by a factor proportional to the inverse square root of the dimension.
|
|
|
|
Args:
|
|
dim (int): Dimensionality of the input.
|
|
eps (float): A small value to avoid division by zero.
|
|
"""
|
|
def __init__(self, dim, eps=1e-5):
|
|
super().__init__()
|
|
self.scale = dim ** -0.5 # Scaling factor
|
|
self.eps = eps # Epsilon for numerical stability
|
|
self.g = nn.Parameter(torch.ones(1)) # Trainable scaling parameter
|
|
|
|
def forward(self, x):
|
|
norm = torch.norm(x, dim=-1, keepdim=True) * self.scale # Compute norm
|
|
return x / norm.clamp(min=self.eps) * self.g # Scale input by norm
|
|
|
|
# Absolute Positional Encodings
|
|
|
|
class ScaledSinuEmbedding(nn.Module):
|
|
"""
|
|
Scaled Sinusoidal Embedding: Generates sinusoidal positional encodings
|
|
that are scaled by a learnable parameter.
|
|
|
|
Args:
|
|
dim (int): Dimensionality of the embedding.
|
|
"""
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.scale = nn.Parameter(torch.ones(1)) # Learnable scaling parameter
|
|
inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim)) # Inverse frequency for sinusoidal embeddings
|
|
self.register_buffer('inv_freq', inv_freq) # Save as a non-trainable buffer
|
|
|
|
def forward(self, x):
|
|
n, device = x.shape[1], x.device
|
|
t = torch.arange(n, device=device).type_as(self.inv_freq) # Time indices
|
|
sinu = einsum('i, j -> i j', t, self.inv_freq) # Sinusoidal function
|
|
emb = torch.cat((sinu.sin(), sinu.cos()), dim=-1) # Concatenate sine and cosine embeddings
|
|
return emb * self.scale # Scale the embeddings
|
|
|
|
# Offset-Scale Layer
|
|
|
|
class OffsetScale(nn.Module):
|
|
"""
|
|
OffsetScale: Applies an element-wise affine transformation (scaling and offset)
|
|
to the input tensor.
|
|
|
|
Args:
|
|
dim (int): Dimensionality of the input.
|
|
heads (int): Number of heads for multi-head operations.
|
|
"""
|
|
def __init__(self, dim, heads=1):
|
|
super().__init__()
|
|
self.gamma = nn.Parameter(torch.ones(heads, dim)) # Learnable scaling parameter
|
|
self.beta = nn.Parameter(torch.zeros(heads, dim)) # Learnable bias parameter
|
|
nn.init.normal_(self.gamma, std=0.02) # Initialize gamma with a small standard deviation
|
|
|
|
def forward(self, x):
|
|
# Apply scaling and offset, then split along the head dimension
|
|
out = einsum('... d, h d -> ... h d', x, self.gamma) + self.beta
|
|
return out.unbind(dim=-2) # Unbind (split) the tensor along the head dimension
|
|
|
|
# Feedforward Convolution Module
|
|
|
|
class FFConvM(nn.Module):
|
|
"""
|
|
Feedforward Convolution Module: A feedforward network with normalization,
|
|
linear projection, and convolution for processing sequential data.
|
|
|
|
Args:
|
|
dim_in (int): Input dimensionality.
|
|
dim_out (int): Output dimensionality.
|
|
norm_klass (class): Normalization layer class (e.g., LayerNorm).
|
|
dropout (float): Dropout rate.
|
|
"""
|
|
def __init__(self, dim_in, dim_out, norm_klass=nn.LayerNorm, dropout=0.1):
|
|
super().__init__()
|
|
self.mdl = nn.Sequential(
|
|
norm_klass(dim_in), # Apply normalization
|
|
nn.Linear(dim_in, dim_out), # Linear transformation
|
|
nn.SiLU(), # SiLU activation function
|
|
ConvModule(dim_out), # Apply convolution module
|
|
nn.Dropout(dropout) # Apply dropout for regularization
|
|
)
|
|
|
|
def forward(self, x):
|
|
output = self.mdl(x) # Forward pass through the module
|
|
return output
|
|
|
|
class FFM(nn.Module):
|
|
"""
|
|
Feedforward Module (FFM): A basic feedforward network that consists of
|
|
normalization, linear projection, activation, and dropout for regularization.
|
|
|
|
Args:
|
|
dim_in (int): Input dimensionality.
|
|
dim_out (int): Output dimensionality.
|
|
norm_klass (class): Normalization layer class, default is LayerNorm.
|
|
dropout (float): Dropout rate for regularization, default is 0.1.
|
|
"""
|
|
def __init__(self, dim_in, dim_out, norm_klass=nn.LayerNorm, dropout=0.1):
|
|
super().__init__()
|
|
# Define a sequential feedforward network with normalization, linear projection, and activation
|
|
self.mdl = nn.Sequential(
|
|
norm_klass(dim_in), # Apply normalization to stabilize learning
|
|
nn.Linear(dim_in, dim_out), # Linear transformation to project input to output dimensionality
|
|
nn.SiLU(), # SiLU activation function for non-linearity
|
|
nn.Dropout(dropout) # Dropout for regularization to prevent overfitting
|
|
)
|
|
|
|
def forward(self, x):
|
|
"""Forward pass through the feedforward network."""
|
|
output = self.mdl(x) # Apply the feedforward module to the input
|
|
return output
|
|
|
|
class FLASH_ShareA_FFConvM(nn.Module):
|
|
"""
|
|
FLASH_ShareA_FFConvM: A block that combines feedforward convolutional modules (FFConvM)
|
|
with a specialized attention mechanism to process sequences in groups and
|
|
perform efficient attention calculations.
|
|
|
|
This module includes both quadratic and linear attention mechanisms,
|
|
with optional token shifting for better performance in causal settings.
|
|
It also supports rotary positional embeddings and flexible normalization.
|
|
|
|
Args:
|
|
dim (int): The input and output dimensionality of the model.
|
|
group_size (int): The size of groups used for attention calculations. Default is 256.
|
|
query_key_dim (int): Dimensionality of the query and key vectors. Default is 128.
|
|
expansion_factor (float): Factor to expand the dimensionality in the hidden layer. Default is 1.0.
|
|
causal (bool): Whether to use causal attention (for autoregressive tasks). Default is False.
|
|
dropout (float): Dropout rate for regularization. Default is 0.1.
|
|
rotary_pos_emb (RotaryEmbedding, optional): Positional embedding using rotary encoding. Default is None.
|
|
norm_klass (class): Normalization class, defaults to LayerNorm.
|
|
shift_tokens (bool): Whether to shift tokens before attention for performance boost. Default is True.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
*,
|
|
dim,
|
|
group_size=256,
|
|
query_key_dim=128,
|
|
expansion_factor=1.,
|
|
causal=False,
|
|
dropout=0.1,
|
|
rotary_pos_emb=None,
|
|
norm_klass=nn.LayerNorm,
|
|
shift_tokens=True
|
|
):
|
|
super().__init__()
|
|
hidden_dim = int(dim * expansion_factor)
|
|
self.group_size = group_size
|
|
self.causal = causal
|
|
self.shift_tokens = shift_tokens
|
|
|
|
# Rotary positional embeddings
|
|
self.rotary_pos_emb = rotary_pos_emb
|
|
|
|
# Dropout layer for regularization
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
# Input projections
|
|
self.to_hidden = FFConvM( # FFConvM for value and gating
|
|
dim_in=dim,
|
|
dim_out=hidden_dim,
|
|
norm_klass=norm_klass,
|
|
dropout=dropout,
|
|
)
|
|
self.to_qk = FFConvM( # FFConvM for query and key
|
|
dim_in=dim,
|
|
dim_out=query_key_dim,
|
|
norm_klass=norm_klass,
|
|
dropout=dropout,
|
|
)
|
|
|
|
# Scaling and offset for attention
|
|
self.qk_offset_scale = OffsetScale(query_key_dim, heads=4)
|
|
|
|
# Output projection
|
|
self.to_out = FFConvM( # FFConvM to combine and produce final output
|
|
dim_in=dim*2,
|
|
dim_out=dim,
|
|
norm_klass=norm_klass,
|
|
dropout=dropout,
|
|
)
|
|
|
|
# Sigmoid gate activation
|
|
self.gateActivate = nn.Sigmoid()
|
|
|
|
def forward(self, x, *, mask=None):
|
|
"""
|
|
Forward pass for the block.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch, sequence length, dim).
|
|
mask (Tensor, optional): Attention mask. Default is None.
|
|
|
|
Returns:
|
|
Tensor: Output tensor after attention and feedforward operations.
|
|
"""
|
|
|
|
# Save input as residual for skip connection
|
|
residual = x
|
|
|
|
# Optional token shifting
|
|
if self.shift_tokens:
|
|
x_shift, x_pass = x.chunk(2, dim=-1)
|
|
x_shift = F.pad(x_shift, (0, 0, 1, -1), value=0.) # Shift tokens
|
|
x = torch.cat((x_shift, x_pass), dim=-1)
|
|
|
|
# Projections for value and gating
|
|
v, u = self.to_hidden(x).chunk(2, dim=-1) # Split into two branches: v and u
|
|
qk = self.to_qk(x) # Query and key projections
|
|
|
|
# Offset and scale for queries and keys
|
|
quad_q, lin_q, quad_k, lin_k = self.qk_offset_scale(qk)
|
|
|
|
# Calculate attention output
|
|
att_v, att_u = self.cal_attention(x, quad_q, lin_q, quad_k, lin_k, v, u, mask)
|
|
|
|
# Gated interaction between attention outputs and gating mechanism
|
|
out = (att_u * v) * self.gateActivate(att_v * u)
|
|
|
|
# Residual connection and output projection
|
|
x = residual + self.to_out(out)
|
|
return x
|
|
|
|
def cal_attention(self, x, quad_q, lin_q, quad_k, lin_k, v, u, mask=None):
|
|
"""
|
|
Computes attention using quadratic and linear mechanisms.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch, sequence length, dim).
|
|
quad_q (Tensor): Quadratic query.
|
|
lin_q (Tensor): Linear query.
|
|
quad_k (Tensor): Quadratic key.
|
|
lin_k (Tensor): Linear key.
|
|
v (Tensor): Value tensor.
|
|
u (Tensor): Gating tensor.
|
|
mask (Tensor, optional): Attention mask. Default is None.
|
|
|
|
Returns:
|
|
Tuple[Tensor]: Attention outputs for value and gating.
|
|
"""
|
|
b, n, device, g = x.shape[0], x.shape[-2], x.device, self.group_size
|
|
|
|
# Mask for linear attention (if provided)
|
|
if exists(mask):
|
|
lin_mask = rearrange(mask, '... -> ... 1')
|
|
lin_k = lin_k.masked_fill(~lin_mask, 0.)
|
|
|
|
# Rotary embeddings for queries and keys (if provided)
|
|
if exists(self.rotary_pos_emb):
|
|
quad_q, lin_q, quad_k, lin_k = map(self.rotary_pos_emb.rotate_queries_or_keys, (quad_q, lin_q, quad_k, lin_k))
|
|
|
|
# Padding to match group size
|
|
padding = padding_to_multiple_of(n, g)
|
|
if padding > 0:
|
|
quad_q, quad_k, lin_q, lin_k, v, u = map(lambda t: F.pad(t, (0, 0, 0, padding), value=0.), (quad_q, quad_k, lin_q, lin_k, v, u))
|
|
|
|
mask = default(mask, torch.ones((b, n), device=device, dtype=torch.bool))
|
|
mask = F.pad(mask, (0, padding), value=False)
|
|
|
|
# Group inputs along sequence
|
|
quad_q, quad_k, lin_q, lin_k, v, u = map(lambda t: rearrange(t, 'b (g n) d -> b g n d', n=self.group_size), (quad_q, quad_k, lin_q, lin_k, v, u))
|
|
|
|
if exists(mask):
|
|
mask = rearrange(mask, 'b (g j) -> b g 1 j', j=g)
|
|
|
|
# Quadratic attention
|
|
sim = einsum('... i d, ... j d -> ... i j', quad_q, quad_k) / g
|
|
attn = F.relu(sim) ** 2
|
|
attn = self.dropout(attn)
|
|
|
|
if exists(mask):
|
|
attn = attn.masked_fill(~mask, 0.)
|
|
|
|
if self.causal:
|
|
causal_mask = torch.ones((g, g), dtype=torch.bool, device=device).triu(1)
|
|
attn = attn.masked_fill(causal_mask, 0.)
|
|
|
|
quad_out_v = einsum('... i j, ... j d -> ... i d', attn, v)
|
|
quad_out_u = einsum('... i j, ... j d -> ... i d', attn, u)
|
|
|
|
# Linear attention (with cumulative sum for causal mode)
|
|
if self.causal:
|
|
lin_kv = einsum('b g n d, b g n e -> b g d e', lin_k, v) / g
|
|
lin_kv = lin_kv.cumsum(dim=1)
|
|
lin_kv = F.pad(lin_kv, (0, 0, 0, 0, 1, -1), value=0.)
|
|
lin_out_v = einsum('b g d e, b g n d -> b g n e', lin_kv, lin_q)
|
|
|
|
lin_ku = einsum('b g n d, b g n e -> b g d e', lin_k, u) / g
|
|
lin_ku = lin_ku.cumsum(dim=1)
|
|
lin_ku = F.pad(lin_ku, (0, 0, 0, 0, 1, -1), value=0.)
|
|
lin_out_u = einsum('b g d e, b g n d -> b g n e', lin_ku, lin_q)
|
|
else:
|
|
lin_kv = einsum('b g n d, b g n e -> b d e', lin_k, v) / n
|
|
lin_out_v = einsum('b g n d, b d e -> b g n e', lin_q, lin_kv)
|
|
|
|
lin_ku = einsum('b g n d, b g n e -> b d e', lin_k, u) / n
|
|
lin_out_u = einsum('b g n d, b d e -> b g n e', lin_q, lin_ku)
|
|
|
|
# Fold groups back into full sequence
|
|
return map(lambda t: rearrange(t, 'b g n d -> b (g n) d')[:, :n], (quad_out_v + lin_out_v, quad_out_u + lin_out_u))
|
|
|
|
class Gated_FSMN(nn.Module):
|
|
"""
|
|
Gated_FSMN: A gated feedforward sequential memory network (FSMN) block that combines
|
|
the outputs of two feedforward convolutional modules (FFConvM) to enhance sequence modeling.
|
|
This module applies gated interactions between the outputs of FSMN and a second FFConvM block.
|
|
|
|
The FSMN is useful for capturing long-term dependencies in sequential data while
|
|
the gating mechanism regulates the influence of FSMN outputs.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels.
|
|
out_channels (int): Number of output channels.
|
|
lorder (int): Filter length or order for FSMN.
|
|
hidden_size (int): Size of the hidden layers used within the FSMN and FFConvM.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
out_channels,
|
|
lorder,
|
|
hidden_size
|
|
):
|
|
super().__init__()
|
|
|
|
# FFConvM block for 'u' branch
|
|
self.to_u = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_size,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
|
|
# FFConvM block for 'v' branch
|
|
self.to_v = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_size,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
|
|
# Unidirectional FSMN (UniDeepFsmn) for processing 'u' branch
|
|
self.fsmn = UniDeepFsmn(in_channels, out_channels, lorder, hidden_size)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the Gated_FSMN block.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch, sequence length, in_channels).
|
|
|
|
Returns:
|
|
Tensor: Output tensor after applying gated FSMN and feedforward operations.
|
|
"""
|
|
input = x # Save original input for skip connection
|
|
|
|
# Process input through FFConvM for both 'u' and 'v' branches
|
|
x_u = self.to_u(x)
|
|
x_v = self.to_v(x)
|
|
|
|
# Apply FSMN to the 'u' branch
|
|
x_u = self.fsmn(x_u)
|
|
|
|
# Gated interaction between 'u' and 'v' branches, followed by skip connection
|
|
x = x_v * x_u + input
|
|
return x
|
|
|
|
class Gated_FSMN_Block(nn.Module):
|
|
"""
|
|
Gated-FSMN Block: A sequential block that combines a Gated Feedforward Sequential Memory Network (FSMN)
|
|
with normalization and convolutional layers for enhanced feature learning. This block applies gating
|
|
mechanisms on sequential data to capture long-term dependencies, while maintaining efficient processing.
|
|
|
|
Args:
|
|
dim (int): Number of input channels.
|
|
inner_channels (int, optional): Number of channels used in the inner layers. Defaults to 256.
|
|
group_size (int, optional): Size of the groups in sequential processing. Defaults to 256.
|
|
norm_type (str, optional): Type of normalization to use ('scalenorm' or 'layernorm'). Defaults to 'scalenorm'.
|
|
"""
|
|
def __init__(self,
|
|
dim,
|
|
inner_channels=256,
|
|
group_size=256,
|
|
norm_type='scalenorm'):
|
|
super(Gated_FSMN_Block, self).__init__()
|
|
|
|
# Select the normalization method based on 'norm_type'
|
|
if norm_type == 'scalenorm':
|
|
norm_klass = ScaleNorm
|
|
elif norm_type == 'layernorm':
|
|
norm_klass = nn.LayerNorm
|
|
|
|
self.group_size = group_size
|
|
|
|
# First 1D convolution layer to project input to 'inner_channels' dimension
|
|
self.conv1 = nn.Sequential(
|
|
nn.Conv1d(dim, inner_channels, kernel_size=1), # Pointwise convolution
|
|
nn.PReLU(), # Parametric ReLU activation
|
|
)
|
|
|
|
# First layer normalization (using CLayerNorm for channel-wise normalization)
|
|
self.norm1 = CLayerNorm(inner_channels)
|
|
|
|
# Gated FSMN for long-term sequential modeling with gating mechanism
|
|
self.gated_fsmn = Gated_FSMN(inner_channels, inner_channels, lorder=20, hidden_size=inner_channels)
|
|
|
|
# Second layer normalization (channel-wise) after FSMN
|
|
self.norm2 = CLayerNorm(inner_channels)
|
|
|
|
# Second 1D convolution layer to project output back to 'dim' dimension
|
|
self.conv2 = nn.Conv1d(inner_channels, dim, kernel_size=1)
|
|
|
|
def forward(self, input):
|
|
"""
|
|
Forward pass through the Gated-FSMN Block.
|
|
|
|
Args:
|
|
input (Tensor): Input tensor of shape (batch_size, seq_len, dim).
|
|
|
|
Returns:
|
|
Tensor: Output tensor of shape (batch_size, seq_len, dim).
|
|
"""
|
|
|
|
# Apply first 1D convolution and activation (transpose to match Conv1d format)
|
|
conv1 = self.conv1(input.transpose(2, 1))
|
|
|
|
# Normalize the output of the first convolution
|
|
norm1 = self.norm1(conv1)
|
|
|
|
# Apply the Gated FSMN block to the normalized output (transpose to match FSMN format)
|
|
seq_out = self.gated_fsmn(norm1.transpose(2, 1))
|
|
|
|
# Normalize the output of FSMN (transpose back to match Conv1d format)
|
|
norm2 = self.norm2(seq_out.transpose(2, 1))
|
|
|
|
# Apply second 1D convolution
|
|
conv2 = self.conv2(norm2)
|
|
|
|
# Add the input (skip connection) and return the result (transpose back to original format)
|
|
return conv2.transpose(2, 1) + input
|
|
|
|
import torch.nn as nn
|
|
|
|
class Gated_FSMN_dilated(nn.Module):
|
|
"""
|
|
Gated FSMN (Finite State Machine Network) with dilated convolutions.
|
|
|
|
This module implements a gated mechanism using two parallel feedforward
|
|
convolutions to generate the input for a dilated FSMN. The gated outputs
|
|
are combined to enhance the input features, allowing for better speech
|
|
enhancement performance.
|
|
|
|
Attributes:
|
|
to_u (FFConvM): Feedforward convolution module for input transformation
|
|
to the u-gate.
|
|
to_v (FFConvM): Feedforward convolution module for input transformation
|
|
to the v-gate.
|
|
fsmn (UniDeepFsmn_dilated): The dilated FSMN for processing the u-gate
|
|
output.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
lorder: int,
|
|
hidden_size: int
|
|
):
|
|
"""
|
|
Initializes the Gated_FSMN_dilated module.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels (features).
|
|
out_channels (int): Number of output channels (features).
|
|
lorder (int): Order of the FSMN.
|
|
hidden_size (int): Number of hidden units in the feedforward layers.
|
|
"""
|
|
super().__init__()
|
|
|
|
# Feedforward convolution for the u-gate
|
|
self.to_u = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_size,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
|
|
# Feedforward convolution for the v-gate
|
|
self.to_v = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_size,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
|
|
# Initialize the dilated FSMN
|
|
self.fsmn = UniDeepFsmn_dilated(in_channels, out_channels, lorder, hidden_size)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Forward pass through the Gated FSMN module.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape (batch_size, in_channels, seq_length).
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after processing through the gated FSMN.
|
|
"""
|
|
input = x # Store the original input for later use
|
|
x_u = self.to_u(x) # Process input through u-gate
|
|
x_v = self.to_v(x) # Process input through v-gate
|
|
|
|
x_u = self.fsmn(x_u) # Apply FSMN to u-gate output
|
|
|
|
# Combine the outputs from u-gate and v-gate with the original input
|
|
x = x_v * x_u + input # Gated output with residual connection
|
|
|
|
return x # Return the final output tensor
|
|
|
|
import torch.nn as nn
|
|
|
|
class Gated_FSMN_Block_Dilated(nn.Module):
|
|
"""
|
|
Gated FSMN (Finite State Machine Network) block with dilated convolutions.
|
|
|
|
This module implements a Gated FSMN block that utilizes dilated convolutions
|
|
for feature extraction and gating mechanisms to enhance speech processing.
|
|
The architecture consists of convolutional layers followed by normalization
|
|
and a gated FSMN for robust feature extraction.
|
|
|
|
Attributes:
|
|
group_size (int): Size of the groups for normalization.
|
|
conv1 (nn.Sequential): Initial 1D convolutional layer followed by
|
|
PReLU activation.
|
|
norm1 (CLayerNorm): First normalization layer.
|
|
gated_fsmn (Gated_FSMN_dilated): Gated FSMN module for processing.
|
|
norm2 (CLayerNorm): Second normalization layer.
|
|
conv2 (nn.Conv1d): Final 1D convolutional layer to map features back
|
|
to the original dimension.
|
|
"""
|
|
|
|
def __init__(self,
|
|
dim: int,
|
|
inner_channels: int = 256,
|
|
group_size: int = 256,
|
|
norm_type: str = 'scalenorm',
|
|
):
|
|
"""
|
|
Initializes the Gated_FSMN_Block_Dilated module.
|
|
|
|
Args:
|
|
dim (int): The number of input channels (features).
|
|
inner_channels (int): The number of channels in the inner layers.
|
|
group_size (int): Size of the groups for normalization.
|
|
norm_type (str): Type of normalization to use ('scalenorm' or 'layernorm').
|
|
"""
|
|
super(Gated_FSMN_Block_Dilated, self).__init__()
|
|
|
|
# Set normalization class based on the specified type
|
|
if norm_type == 'scalenorm':
|
|
norm_klass = ScaleNorm # Use ScaleNorm
|
|
elif norm_type == 'layernorm':
|
|
norm_klass = nn.LayerNorm # Use LayerNorm
|
|
|
|
self.group_size = group_size
|
|
|
|
# Initial convolution layer with PReLU activation
|
|
self.conv1 = nn.Sequential(
|
|
nn.Conv1d(dim, inner_channels, kernel_size=1),
|
|
nn.PReLU(),
|
|
)
|
|
|
|
self.norm1 = CLayerNorm(inner_channels) # First normalization layer
|
|
|
|
# Gated FSMN block with dilated convolutions
|
|
self.gated_fsmn = Gated_FSMN_dilated(inner_channels, inner_channels, lorder=20, hidden_size=inner_channels)
|
|
|
|
self.norm2 = CLayerNorm(inner_channels) # Second normalization layer
|
|
self.conv2 = nn.Conv1d(inner_channels, dim, kernel_size=1) # Output convolution layer
|
|
|
|
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Forward pass through the Gated FSMN block.
|
|
|
|
Args:
|
|
input (torch.Tensor): Input tensor of shape (batch_size, seq_length, dim).
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after processing through the Gated FSMN block.
|
|
"""
|
|
# Apply the first convolution and transpose for correct dimensions
|
|
conv1 = self.conv1(input.transpose(2, 1))
|
|
norm1 = self.norm1(conv1) # Apply normalization after first convolution
|
|
|
|
# Process through the gated FSMN and transpose back to original dimensions
|
|
seq_out = self.gated_fsmn(norm1.transpose(2, 1))
|
|
norm2 = self.norm2(seq_out.transpose(2, 1)) # Apply normalization after gated FSMN
|
|
|
|
# Apply the second convolution and return the residual connection
|
|
conv2 = self.conv2(norm2) # Final convolution
|
|
return conv2.transpose(2, 1) + input # Return output with residual connection
|
|
|
|
|
|
class MossformerBlock_GFSMN(nn.Module):
|
|
"""
|
|
Mossformer2 Block with Gated FSMN: A module that integrates single-head gated attention mechanisms
|
|
with Gated Feedforward Sequential Memory Networks (FSMNs) to enhance feature representation
|
|
in sequential data. This block employs multiple layers of attention and gated mechanisms
|
|
for improved learning capabilities.
|
|
|
|
Args:
|
|
dim (int): Number of input channels.
|
|
depth (int): Number of layers to stack in the block.
|
|
group_size (int, optional): Size of the groups for sequential processing. Defaults to 256.
|
|
query_key_dim (int, optional): Dimension for query and key projections in attention. Defaults to 128.
|
|
expansion_factor (float, optional): Factor to expand dimensions in the feedforward layers. Defaults to 4.0.
|
|
causal (bool, optional): Whether to apply causal masking in attention. Defaults to False.
|
|
attn_dropout (float, optional): Dropout rate for attention layers. Defaults to 0.1.
|
|
norm_type (str, optional): Type of normalization to use ('scalenorm' or 'layernorm'). Defaults to 'scalenorm'.
|
|
shift_tokens (bool, optional): Whether to apply token shifting. Defaults to True.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
*,
|
|
dim,
|
|
depth,
|
|
group_size=256,
|
|
query_key_dim=128,
|
|
expansion_factor=4.0,
|
|
causal=False,
|
|
attn_dropout=0.1,
|
|
norm_type='scalenorm',
|
|
shift_tokens=True
|
|
):
|
|
super().__init__()
|
|
|
|
# Assert valid normalization type
|
|
assert norm_type in ('scalenorm', 'layernorm'), 'norm_type must be one of scalenorm or layernorm'
|
|
|
|
# Choose normalization class based on the specified type
|
|
if norm_type == 'scalenorm':
|
|
norm_klass = ScaleNorm
|
|
elif norm_type == 'layernorm':
|
|
norm_klass = nn.LayerNorm
|
|
|
|
self.group_size = group_size
|
|
|
|
# Initialize rotary positional embeddings with a maximum dimension of 32
|
|
rotary_pos_emb = RotaryEmbedding(dim=min(32, query_key_dim))
|
|
|
|
# Create a list of Gated FSMN blocks for each layer
|
|
self.fsmn = nn.ModuleList([Gated_FSMN_Block_Dilated(dim) for _ in range(depth)])
|
|
|
|
# Create a list of FLASH attention layers
|
|
self.layers = nn.ModuleList([
|
|
FLASH_ShareA_FFConvM(
|
|
dim=dim,
|
|
group_size=group_size,
|
|
query_key_dim=query_key_dim,
|
|
expansion_factor=expansion_factor,
|
|
causal=causal,
|
|
dropout=attn_dropout,
|
|
rotary_pos_emb=rotary_pos_emb,
|
|
norm_klass=norm_klass,
|
|
shift_tokens=shift_tokens
|
|
) for _ in range(depth)
|
|
])
|
|
|
|
def _build_repeats(self, in_channels, out_channels, lorder, hidden_size, repeats=1):
|
|
"""
|
|
Build a sequential block of UniDeep FSMNs.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels for the FSMN.
|
|
out_channels (int): Number of output channels for the FSMN.
|
|
lorder (int): Order for the FSMN.
|
|
hidden_size (int): Hidden size for the FSMN.
|
|
repeats (int, optional): Number of repetitions of the FSMN block. Defaults to 1.
|
|
|
|
Returns:
|
|
nn.Sequential: A sequential module containing the specified number of UniDeep FSMNs.
|
|
"""
|
|
repeats = [
|
|
UniDeepFsmn(in_channels, out_channels, lorder, hidden_size)
|
|
for i in range(repeats)
|
|
]
|
|
return nn.Sequential(*repeats)
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
*,
|
|
mask=None
|
|
):
|
|
"""
|
|
Forward pass through the Mossformer Block.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch_size, seq_len, dim).
|
|
mask (Tensor, optional): Attention mask to apply. Defaults to None.
|
|
|
|
Returns:
|
|
Tensor: Output tensor after passing through all layers, of shape (batch_size, seq_len, dim).
|
|
"""
|
|
ii = 0
|
|
# Iterate through all FLASH attention layers and Gated FSMN blocks
|
|
for flash in self.layers:
|
|
x = flash(x, mask=mask) # Apply FLASH attention layer
|
|
x = self.fsmn[ii](x) # Apply corresponding Gated FSMN block
|
|
ii += 1 # Increment index for the Gated FSMN block
|
|
|
|
return x # Return the final output after all layers
|
|
|
|
class MossformerBlock(nn.Module):
|
|
"""
|
|
Mossformer Block: A module that employs a series of signle-head gated attention layers to process
|
|
sequential data. This block is designed for flexibility in feature dimension, depth,
|
|
and normalization techniques, making it suitable for various tasks in deep learning.
|
|
|
|
Args:
|
|
dim (int): Number of input channels (features).
|
|
depth (int): Number of layers in the block.
|
|
group_size (int, optional): Size of the groups for processing. Defaults to 256.
|
|
query_key_dim (int, optional): Dimension for query and key projections in attention. Defaults to 128.
|
|
expansion_factor (float, optional): Factor to expand the dimensionality in feedforward layers. Defaults to 4.0.
|
|
causal (bool, optional): Whether to apply causal masking in attention. Defaults to False.
|
|
attn_dropout (float, optional): Dropout rate applied to attention layers. Defaults to 0.1.
|
|
norm_type (str, optional): Type of normalization to apply ('scalenorm' or 'layernorm'). Defaults to 'scalenorm'.
|
|
shift_tokens (bool, optional): Whether to apply token shifting. Defaults to True.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
*,
|
|
dim,
|
|
depth,
|
|
group_size=256,
|
|
query_key_dim=128,
|
|
expansion_factor=4.0,
|
|
causal=False,
|
|
attn_dropout=0.1,
|
|
norm_type='scalenorm',
|
|
shift_tokens=True
|
|
):
|
|
super().__init__()
|
|
|
|
# Assert valid normalization type
|
|
assert norm_type in ('scalenorm', 'layernorm'), 'norm_type must be one of scalenorm or layernorm'
|
|
|
|
# Select normalization class based on the specified type
|
|
if norm_type == 'scalenorm':
|
|
norm_klass = ScaleNorm
|
|
elif norm_type == 'layernorm':
|
|
norm_klass = nn.LayerNorm
|
|
|
|
self.group_size = group_size
|
|
|
|
# Initialize rotary positional embeddings, limiting max dimension to 32
|
|
rotary_pos_emb = RotaryEmbedding(dim=min(32, query_key_dim))
|
|
|
|
# Create a list of FLASH attention layers for the specified depth
|
|
self.layers = nn.ModuleList([
|
|
FLASH_ShareA_FFConvM(
|
|
dim=dim,
|
|
group_size=group_size,
|
|
query_key_dim=query_key_dim,
|
|
expansion_factor=expansion_factor,
|
|
causal=causal,
|
|
dropout=attn_dropout,
|
|
rotary_pos_emb=rotary_pos_emb,
|
|
norm_klass=norm_klass,
|
|
shift_tokens=shift_tokens
|
|
) for _ in range(depth)
|
|
])
|
|
|
|
def _build_repeats(self, in_channels, out_channels, lorder, hidden_size, repeats=1):
|
|
"""
|
|
Build a sequential block of UniDeep FSMNs.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels for the FSMN.
|
|
out_channels (int): Number of output channels for the FSMN.
|
|
lorder (int): Order for the FSMN.
|
|
hidden_size (int): Hidden size for the FSMN.
|
|
repeats (int, optional): Number of repetitions of the FSMN block. Defaults to 1.
|
|
|
|
Returns:
|
|
nn.Sequential: A sequential module containing the specified number of UniDeep FSMNs.
|
|
"""
|
|
repeats = [
|
|
UniDeepFsmn(in_channels, out_channels, lorder, hidden_size)
|
|
for i in range(repeats)
|
|
]
|
|
return nn.Sequential(*repeats)
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
*,
|
|
mask=None
|
|
):
|
|
"""
|
|
Forward pass through the Mossformer Block.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch_size, seq_len, dim).
|
|
mask (Tensor, optional): Attention mask to apply. Defaults to None.
|
|
|
|
Returns:
|
|
Tensor: Output tensor after passing through all layers, of shape (batch_size, seq_len, dim).
|
|
"""
|
|
ii = 0
|
|
# Iterate through all FLASH attention layers and apply them to the input
|
|
for flash in self.layers:
|
|
x = flash(x, mask=mask) # Apply FLASH attention layer
|
|
ii += 1 # Increment layer index
|
|
|
|
return x # Return the final output after processing through all layers
|
|
|