736 lines
26 KiB
Python
736 lines
26 KiB
Python
"""
|
|
This source code is modified 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 einops import rearrange
|
|
from rotary_embedding_torch import RotaryEmbedding
|
|
from models.mossformer2_se.conv_module import ConvModule, GLU, FFConvM_Dilated
|
|
from models.mossformer2_se.fsmn import UniDeepFsmn, UniDeepFsmn_dilated
|
|
from torchinfo import summary
|
|
from models.mossformer2_se.layer_norm import CLayerNorm, GLayerNorm, GlobLayerNorm, ILayerNorm
|
|
|
|
# Helper functions
|
|
|
|
def identity(t, *args, **kwargs):
|
|
"""
|
|
Returns the input tensor unchanged.
|
|
|
|
Args:
|
|
t (torch.Tensor): Input tensor.
|
|
*args: Additional arguments (ignored).
|
|
**kwargs: Additional keyword arguments (ignored).
|
|
|
|
Returns:
|
|
torch.Tensor: The input tensor.
|
|
"""
|
|
return t
|
|
|
|
def append_dims(x, num_dims):
|
|
"""
|
|
Adds additional dimensions to the input tensor.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor.
|
|
num_dims (int): Number of dimensions to append.
|
|
|
|
Returns:
|
|
torch.Tensor: Tensor with appended dimensions.
|
|
"""
|
|
if num_dims <= 0:
|
|
return x
|
|
return x.view(*x.shape, *((1,) * num_dims)) # Reshape to append dimensions
|
|
|
|
def exists(val):
|
|
"""
|
|
Checks if a value exists (is not None).
|
|
|
|
Args:
|
|
val: The value to check.
|
|
|
|
Returns:
|
|
bool: True if value exists, False otherwise.
|
|
"""
|
|
return val is not None
|
|
|
|
def default(val, d):
|
|
"""
|
|
Returns a default value if the given value does not exist.
|
|
|
|
Args:
|
|
val: The value to check.
|
|
d: Default value to return if val does not exist.
|
|
|
|
Returns:
|
|
The original value if it exists, otherwise the default value.
|
|
"""
|
|
return val if exists(val) else d
|
|
|
|
def padding_to_multiple_of(n, mult):
|
|
"""
|
|
Calculates the amount of padding needed to make a number a multiple of another.
|
|
|
|
Args:
|
|
n (int): The number to pad.
|
|
mult (int): The multiple to match.
|
|
|
|
Returns:
|
|
int: The padding amount required to make n a multiple of mult.
|
|
"""
|
|
remainder = n % mult
|
|
if remainder == 0:
|
|
return 0
|
|
return mult - remainder # Return the required padding
|
|
|
|
# Scale Normalization class
|
|
|
|
class ScaleNorm(nn.Module):
|
|
"""
|
|
ScaleNorm implements a scaled normalization technique for neural network layers.
|
|
|
|
Attributes:
|
|
dim (int): Dimension of the input features.
|
|
eps (float): Small value to prevent division by zero.
|
|
g (nn.Parameter): Learnable parameter for scaling.
|
|
"""
|
|
|
|
def __init__(self, dim, eps=1e-5):
|
|
super().__init__()
|
|
self.scale = dim ** -0.5 # Calculate scale factor
|
|
self.eps = eps # Set epsilon
|
|
self.g = nn.Parameter(torch.ones(1)) # Initialize scaling parameter
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the ScaleNorm layer.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor.
|
|
|
|
Returns:
|
|
torch.Tensor: Scaled and normalized output tensor.
|
|
"""
|
|
norm = torch.norm(x, dim=-1, keepdim=True) * self.scale # Compute norm
|
|
return x / norm.clamp(min=self.eps) * self.g # Normalize and scale
|
|
|
|
# Absolute positional encodings class
|
|
|
|
class ScaledSinuEmbedding(nn.Module):
|
|
"""
|
|
ScaledSinuEmbedding provides sinusoidal positional encodings for inputs.
|
|
|
|
Attributes:
|
|
scale (nn.Parameter): Learnable scale factor for the embeddings.
|
|
inv_freq (torch.Tensor): Inverse frequency used for sine and cosine calculations.
|
|
"""
|
|
|
|
def __init__(self, dim):
|
|
super().__init__()
|
|
self.scale = nn.Parameter(torch.ones(1,)) # Initialize scale
|
|
inv_freq = 1. / (10000 ** (torch.arange(0, dim, 2).float() / dim)) # Calculate inverse frequency
|
|
self.register_buffer('inv_freq', inv_freq) # Register as a buffer
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the ScaledSinuEmbedding layer.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape (batch_size, sequence_length).
|
|
|
|
Returns:
|
|
torch.Tensor: Positional encoding tensor of shape (batch_size, sequence_length, dim).
|
|
"""
|
|
n, device = x.shape[1], x.device # Extract sequence length and device
|
|
t = torch.arange(n, device=device).type_as(self.inv_freq) # Create time steps
|
|
sinu = einsum('i , j -> i j', t, self.inv_freq) # Calculate sine and cosine embeddings
|
|
emb = torch.cat((sinu.sin(), sinu.cos()), dim=-1) # Concatenate sine and cosine embeddings
|
|
return emb * self.scale # Scale the embeddings
|
|
|
|
class OffsetScale(nn.Module):
|
|
"""
|
|
OffsetScale applies learned offsets and scales to the input tensor.
|
|
|
|
Attributes:
|
|
gamma (nn.Parameter): Learnable scale parameter for each head.
|
|
beta (nn.Parameter): Learnable offset parameter for each head.
|
|
"""
|
|
|
|
def __init__(self, dim, heads=1):
|
|
super().__init__()
|
|
self.gamma = nn.Parameter(torch.ones(heads, dim)) # Initialize scale parameters
|
|
self.beta = nn.Parameter(torch.zeros(heads, dim)) # Initialize offset parameters
|
|
nn.init.normal_(self.gamma, std=0.02) # Normal initialization for gamma
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the OffsetScale layer.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor.
|
|
|
|
Returns:
|
|
List[torch.Tensor]: A list of tensors with applied offsets and scales for each head.
|
|
"""
|
|
out = einsum('... d, h d -> ... h d', x, self.gamma) + self.beta # Apply scaling and offsets
|
|
return out.unbind(dim=-2) # Unbind heads into a list
|
|
|
|
# Feed-Forward Convolutional Module
|
|
|
|
class FFConvM(nn.Module):
|
|
"""
|
|
FFConvM is a feed-forward convolutional module with normalization and dropout.
|
|
|
|
Attributes:
|
|
dim_in (int): Input dimension of the features.
|
|
dim_out (int): Output dimension after processing.
|
|
norm_klass (nn.Module): Normalization class to be used.
|
|
dropout (float): Dropout probability.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim_in,
|
|
dim_out,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1
|
|
):
|
|
super().__init__()
|
|
self.mdl = nn.Sequential(
|
|
norm_klass(dim_in), # Normalize input
|
|
nn.Linear(dim_in, dim_out), # Linear transformation
|
|
nn.SiLU(), # Activation function
|
|
ConvModule(dim_out), # Convolution module
|
|
nn.Dropout(dropout) # Apply dropout
|
|
)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the FFConvM module.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor.
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after processing.
|
|
"""
|
|
output = self.mdl(x) # Pass through the model
|
|
return output
|
|
|
|
class FFM(nn.Module):
|
|
"""
|
|
FFM is a feed-forward module with normalization and dropout.
|
|
|
|
Attributes:
|
|
dim_in (int): Input dimension of the features.
|
|
dim_out (int): Output dimension after processing.
|
|
norm_klass (nn.Module): Normalization class to be used.
|
|
dropout (float): Dropout probability.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim_in,
|
|
dim_out,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1
|
|
):
|
|
super().__init__()
|
|
self.mdl = nn.Sequential(
|
|
norm_klass(dim_in), # Normalize input
|
|
nn.Linear(dim_in, dim_out), # Linear transformation
|
|
nn.SiLU(), # Activation function
|
|
nn.Dropout(dropout) # Apply dropout
|
|
)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the FFM module.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor.
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after processing.
|
|
"""
|
|
output = self.mdl(x) # Pass through the model
|
|
return output
|
|
|
|
class FLASH_ShareA_FFConvM(nn.Module):
|
|
"""
|
|
Fast Shared Dual Attention Mechanism with feed-forward convolutional blocks.
|
|
Published in paper: "MossFormer: Pushing the Performance Limit of Monaural Speech Separation
|
|
using Gated Single-Head Transformer with Convolution-Augmented Joint Self-Attentions", ICASSP 2023.
|
|
(https://arxiv.org/abs/2302.11824)
|
|
|
|
Args:
|
|
dim (int): Input dimension.
|
|
group_size (int, optional): Size of groups for processing. Defaults to 256.
|
|
query_key_dim (int, optional): Dimension of the query and key. Defaults to 128.
|
|
expansion_factor (float, optional): Factor to expand the hidden dimension. Defaults to 1.
|
|
causal (bool, optional): Whether to use causal masking. Defaults to False.
|
|
dropout (float, optional): Dropout rate. Defaults to 0.1.
|
|
rotary_pos_emb (optional): Rotary positional embeddings for attention. Defaults to None.
|
|
norm_klass (callable, optional): Normalization class to use. Defaults to nn.LayerNorm.
|
|
shift_tokens (bool, optional): Whether to shift tokens for attention calculation. Defaults to 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
|
|
|
|
# Initialize positional embeddings, dropout, and projections
|
|
self.rotary_pos_emb = rotary_pos_emb
|
|
self.dropout = nn.Dropout(dropout)
|
|
|
|
# Feed-forward layers
|
|
self.to_hidden = FFConvM(
|
|
dim_in=dim,
|
|
dim_out=hidden_dim,
|
|
norm_klass=norm_klass,
|
|
dropout=dropout,
|
|
)
|
|
self.to_qk = FFConvM(
|
|
dim_in=dim,
|
|
dim_out=query_key_dim,
|
|
norm_klass=norm_klass,
|
|
dropout=dropout,
|
|
)
|
|
|
|
# Offset and scale for query and key
|
|
self.qk_offset_scale = OffsetScale(query_key_dim, heads=4)
|
|
|
|
self.to_out = FFConvM(
|
|
dim_in=dim * 2,
|
|
dim_out=dim,
|
|
norm_klass=norm_klass,
|
|
dropout=dropout,
|
|
)
|
|
|
|
self.gateActivate = nn.Sigmoid()
|
|
|
|
def forward(self, x, *, mask=None):
|
|
"""
|
|
Forward pass for FLASH layer.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch, seq_len, features).
|
|
mask (Tensor, optional): Mask for attention. Defaults to None.
|
|
|
|
Returns:
|
|
Tensor: Output tensor after applying attention and projections.
|
|
"""
|
|
|
|
# Pre-normalization step
|
|
normed_x = x
|
|
residual = x # Save residual for skip connection
|
|
|
|
# Token shifting if enabled
|
|
if self.shift_tokens:
|
|
x_shift, x_pass = normed_x.chunk(2, dim=-1)
|
|
x_shift = F.pad(x_shift, (0, 0, 1, -1), value=0.)
|
|
normed_x = torch.cat((x_shift, x_pass), dim=-1)
|
|
|
|
# Initial projections
|
|
v, u = self.to_hidden(normed_x).chunk(2, dim=-1)
|
|
qk = self.to_qk(normed_x)
|
|
|
|
# Offset and scale
|
|
quad_q, lin_q, quad_k, lin_k = self.qk_offset_scale(qk)
|
|
att_v, att_u = self.cal_attention(x, quad_q, lin_q, quad_k, lin_k, v, u)
|
|
|
|
# Output calculation with gating
|
|
out = (att_u * v) * self.gateActivate(att_v * u)
|
|
x = x + self.to_out(out) # Residual connection
|
|
return x
|
|
|
|
def cal_attention(self, x, quad_q, lin_q, quad_k, lin_k, v, u, mask=None):
|
|
"""
|
|
Calculate attention output using quadratic and linear attention mechanisms.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch, seq_len, features).
|
|
quad_q (Tensor): Quadratic query representation.
|
|
lin_q (Tensor): Linear query representation.
|
|
quad_k (Tensor): Quadratic key representation.
|
|
lin_k (Tensor): Linear key representation.
|
|
v (Tensor): Value representation.
|
|
u (Tensor): Additional value representation.
|
|
mask (Tensor, optional): Mask for attention. Defaults to None.
|
|
|
|
Returns:
|
|
Tuple[Tensor, Tensor]: Attention outputs for v and u.
|
|
"""
|
|
b, n, device, g = x.shape[0], x.shape[-2], x.device, self.group_size
|
|
|
|
# Apply mask to linear keys if provided
|
|
if exists(mask):
|
|
lin_mask = rearrange(mask, '... -> ... 1')
|
|
lin_k = lin_k.masked_fill(~lin_mask, 0.)
|
|
|
|
# Rotate queries and keys with rotary positional embeddings
|
|
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 for group processing
|
|
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 along sequence for attention
|
|
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)
|
|
|
|
# Calculate quadratic attention output
|
|
sim = einsum('... i d, ... j d -> ... i j', quad_q, quad_k) / g
|
|
attn = F.relu(sim) ** 2 # ReLU activation
|
|
attn = self.dropout(attn)
|
|
|
|
# Apply mask to attention if provided
|
|
if exists(mask):
|
|
attn = attn.masked_fill(~mask, 0.)
|
|
|
|
# Apply causal mask if needed
|
|
if self.causal:
|
|
causal_mask = torch.ones((g, g), dtype=torch.bool, device=device).triu(1)
|
|
attn = attn.masked_fill(causal_mask, 0.)
|
|
|
|
# Calculate output from attention
|
|
quad_out_v = einsum('... i j, ... j d -> ... i d', attn, v)
|
|
quad_out_u = einsum('... i j, ... j d -> ... i d', attn, u)
|
|
|
|
# Calculate linear attention output
|
|
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) # Cumulative sum for linear attention
|
|
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) # Cumulative sum for linear attention
|
|
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)
|
|
|
|
# Reshape and remove padding from outputs
|
|
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 Frequency Selective Memory Network (FSMN) class.
|
|
|
|
This class implements a gated FSMN that combines two feedforward
|
|
convolutional networks with a frequency selective memory module.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels.
|
|
out_channels (int): Number of output channels.
|
|
lorder (int): Order of the filter for FSMN.
|
|
hidden_size (int): Number of hidden units in the network.
|
|
"""
|
|
def __init__(self, in_channels, out_channels, lorder, hidden_size):
|
|
super().__init__()
|
|
# Feedforward network for the first branch (u)
|
|
self.to_u = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_size,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
# Feedforward network for the second branch (v)
|
|
self.to_v = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_size,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
# Frequency selective memory network
|
|
self.fsmn = UniDeepFsmn(in_channels, out_channels, lorder, hidden_size)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the Gated FSMN.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch_size, in_channels, sequence_length).
|
|
|
|
Returns:
|
|
Tensor: Output tensor after applying gated FSMN operations.
|
|
"""
|
|
input = x
|
|
x_u = self.to_u(x) # Process input through the first branch
|
|
x_v = self.to_v(x) # Process input through the second branch
|
|
x_u = self.fsmn(x_u) # Apply FSMN to the output of the first branch
|
|
x = x_v * x_u + input # Combine outputs with the original input
|
|
return x
|
|
|
|
|
|
class Gated_FSMN_Block(nn.Module):
|
|
"""
|
|
A 1-D convolutional block that incorporates a gated FSMN.
|
|
|
|
This block consists of two convolutional layers, followed by a
|
|
gated FSMN and normalization layers.
|
|
|
|
Args:
|
|
dim (int): Dimensionality of the input.
|
|
inner_channels (int): 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').
|
|
"""
|
|
def __init__(self, dim, inner_channels=256, group_size=256, norm_type='scalenorm'):
|
|
super(Gated_FSMN_Block, self).__init__()
|
|
# Choose normalization class based on the provided type
|
|
if norm_type == 'scalenorm':
|
|
norm_klass = ScaleNorm
|
|
elif norm_type == 'layernorm':
|
|
norm_klass = nn.LayerNorm
|
|
|
|
self.group_size = group_size
|
|
|
|
# First convolutional layer with PReLU activation
|
|
self.conv1 = nn.Sequential(
|
|
nn.Conv1d(dim, inner_channels, kernel_size=1),
|
|
nn.PReLU(),
|
|
)
|
|
self.norm1 = CLayerNorm(inner_channels) # Normalization after first convolution
|
|
self.gated_fsmn = Gated_FSMN(inner_channels, inner_channels, lorder=20, hidden_size=inner_channels) # Gated FSMN layer
|
|
self.norm2 = CLayerNorm(inner_channels) # Normalization after FSMN
|
|
self.conv2 = nn.Conv1d(inner_channels, dim, kernel_size=1) # Final convolutional layer
|
|
|
|
def forward(self, input):
|
|
"""
|
|
Forward pass for the Gated FSMN Block.
|
|
|
|
Args:
|
|
input (Tensor): Input tensor of shape (batch_size, dim, sequence_length).
|
|
|
|
Returns:
|
|
Tensor: Output tensor after processing through the block.
|
|
"""
|
|
conv1 = self.conv1(input.transpose(2, 1)) # Apply first convolution
|
|
norm1 = self.norm1(conv1) # Apply normalization
|
|
seq_out = self.gated_fsmn(norm1.transpose(2, 1)) # Apply gated FSMN
|
|
norm2 = self.norm2(seq_out.transpose(2, 1)) # Apply second normalization
|
|
conv2 = self.conv2(norm2) # Apply final convolution
|
|
return conv2.transpose(2, 1) + input # Residual connection
|
|
|
|
|
|
class MossformerBlock_GFSMN(nn.Module):
|
|
"""
|
|
Mossformer Block with Gated FSMN.
|
|
|
|
This block combines attention mechanisms and gated FSMN layers
|
|
to process input sequences.
|
|
|
|
Args:
|
|
dim (int): Dimensionality of the input.
|
|
depth (int): Number of layers in the block.
|
|
group_size (int): Size of the groups for normalization.
|
|
query_key_dim (int): Dimension of the query and key in attention.
|
|
expansion_factor (float): Expansion factor for feedforward layers.
|
|
causal (bool): If True, enables causal attention.
|
|
attn_dropout (float): Dropout rate for attention layers.
|
|
norm_type (str): Type of normalization to use ('scalenorm' or 'layernorm').
|
|
shift_tokens (bool): If True, shifts tokens in the attention layer.
|
|
"""
|
|
def __init__(self, *, dim, depth, group_size=256, query_key_dim=128, expansion_factor=4., causal=False, attn_dropout=0.1, norm_type='scalenorm', shift_tokens=True):
|
|
super().__init__()
|
|
assert norm_type in ('scalenorm', 'layernorm'), 'norm_type must be one of scalenorm or layernorm'
|
|
|
|
if norm_type == 'scalenorm':
|
|
norm_klass = ScaleNorm
|
|
elif norm_type == 'layernorm':
|
|
norm_klass = nn.LayerNorm
|
|
|
|
self.group_size = group_size
|
|
|
|
# Rotary positional embedding for attention
|
|
rotary_pos_emb = RotaryEmbedding(dim=min(32, query_key_dim))
|
|
|
|
# Create a list of Gated FSMN blocks
|
|
self.fsmn = nn.ModuleList([Gated_FSMN_Block(dim) for _ in range(depth)])
|
|
|
|
# Create a list of attention layers using FLASH_ShareA_FFConvM
|
|
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):
|
|
"""
|
|
Builds repeated UniDeep FSMN layers.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels.
|
|
out_channels (int): Number of output channels.
|
|
lorder (int): Order of the filter for FSMN.
|
|
hidden_size (int): Number of hidden units.
|
|
repeats (int): Number of repetitions.
|
|
|
|
Returns:
|
|
Sequential: A sequential container with repeated layers.
|
|
"""
|
|
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 for the Mossformer Block with Gated FSMN.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch_size, dim, sequence_length).
|
|
mask (Tensor, optional): Mask tensor for attention operations.
|
|
|
|
Returns:
|
|
Tensor: Output tensor after processing through the block.
|
|
"""
|
|
ii = 0
|
|
for flash in self.layers: # Process through each layer
|
|
x = flash(x, mask=mask)
|
|
x = self.fsmn[ii](x) # Apply corresponding Gated FSMN block
|
|
ii += 1
|
|
|
|
return x
|
|
|
|
|
|
class MossformerBlock(nn.Module):
|
|
"""
|
|
Mossformer Block with attention mechanisms.
|
|
|
|
This block is designed to process input sequences using attention
|
|
layers and incorporates rotary positional embeddings. It allows
|
|
for configurable normalization types and can handle causal
|
|
attention.
|
|
|
|
Args:
|
|
dim (int): Dimensionality of the input.
|
|
depth (int): Number of attention layers in the block.
|
|
group_size (int, optional): Size of groups for normalization. Default is 256.
|
|
query_key_dim (int, optional): Dimension of the query and key in attention. Default is 128.
|
|
expansion_factor (float, optional): Expansion factor for feedforward layers. Default is 4.
|
|
causal (bool, optional): If True, enables causal attention. Default is False.
|
|
attn_dropout (float, optional): Dropout rate for attention layers. Default is 0.1.
|
|
norm_type (str, optional): Type of normalization to use ('scalenorm' or 'layernorm'). Default is 'scalenorm'.
|
|
shift_tokens (bool, optional): If True, shifts tokens in the attention layer. Default is 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__()
|
|
|
|
# Ensure normalization type is valid
|
|
assert norm_type in ('scalenorm', 'layernorm'), 'norm_type must be one of scalenorm or layernorm'
|
|
|
|
# Select normalization class based on the provided type
|
|
if norm_type == 'scalenorm':
|
|
norm_klass = ScaleNorm
|
|
elif norm_type == 'layernorm':
|
|
norm_klass = nn.LayerNorm
|
|
|
|
self.group_size = group_size # Group size for normalization
|
|
|
|
# Rotary positional embedding for attention
|
|
rotary_pos_emb = RotaryEmbedding(dim=min(32, query_key_dim))
|
|
# Max rotary embedding dimensions of 32, partial Rotary embeddings, from Wang et al - GPT-J
|
|
|
|
# Create a list of attention layers using FLASH_ShareA_FFConvM
|
|
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):
|
|
"""
|
|
Builds repeated UniDeep FSMN layers.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels.
|
|
out_channels (int): Number of output channels.
|
|
lorder (int): Order of the filter for FSMN.
|
|
hidden_size (int): Number of hidden units.
|
|
repeats (int, optional): Number of repetitions. Default is 1.
|
|
|
|
Returns:
|
|
Sequential: A sequential container with repeated layers.
|
|
"""
|
|
repeats = [
|
|
UniDeepFsmn(in_channels, out_channels, lorder, hidden_size)
|
|
for _ in range(repeats)
|
|
]
|
|
return nn.Sequential(*repeats)
|
|
|
|
def forward(self, x, *, mask=None):
|
|
"""
|
|
Forward pass for the Mossformer Block.
|
|
|
|
Args:
|
|
x (Tensor): Input tensor of shape (batch_size, dim, sequence_length).
|
|
mask (Tensor, optional): Mask tensor for attention operations.
|
|
|
|
Returns:
|
|
Tensor: Output tensor after processing through the block.
|
|
"""
|
|
# Process input through each attention layer
|
|
for flash in self.layers:
|
|
x = flash(x, mask=mask) # Apply attention layer with optional mask
|
|
|
|
return x # Return the final output tensor
|