841 lines
36 KiB
Python
841 lines
36 KiB
Python
import math
|
|
from collections import OrderedDict
|
|
from typing import Dict, List, Optional, Tuple, Union
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from packaging.version import parse as V
|
|
from torch.nn import init
|
|
from torch.nn.parameter import Parameter
|
|
|
|
from models.mossformer_gan_se.fsmn import UniDeepFsmn
|
|
from models.mossformer_gan_se.conv_module import ConvModule
|
|
from models.mossformer_gan_se.mossformer import MossFormer
|
|
from models.mossformer_gan_se.se_layer import SELayer
|
|
from models.mossformer_gan_se.get_layer_from_string import get_layer
|
|
from models.mossformer_gan_se.discriminator import Discriminator
|
|
|
|
# Check if the installed version of PyTorch is 1.9.0 or higher
|
|
is_torch_1_9_plus = V(torch.__version__) >= V("1.9.0")
|
|
|
|
|
|
class MossFormerGAN_SE_16K(nn.Module):
|
|
"""
|
|
MossFormerGAN_SE_16K: A GAN-based speech enhancement model for 16kHz input audio.
|
|
|
|
This model integrates a synchronous attention network (SyncANet) for
|
|
feature extraction. Depending on the mode (train or inference), it may
|
|
also include a discriminator for adversarial training.
|
|
|
|
Args:
|
|
args (Namespace): Arguments containing configuration parameters,
|
|
including 'fft_len' and 'mode'.
|
|
"""
|
|
|
|
def __init__(self, args):
|
|
"""Initializes the MossFormerGAN_SE_16K model."""
|
|
super(MossFormerGAN_SE_16K, self).__init__()
|
|
|
|
# Initialize SyncANet with specified number of channels and features
|
|
self.model = SyncANet(num_channel=64, num_features=args.fft_len // 2 + 1)
|
|
|
|
# Initialize discriminator if in training mode
|
|
if args.mode == 'train':
|
|
self.discriminator = Discriminator(ndf=16)
|
|
else:
|
|
self.discriminator = None
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass of the MossFormerGAN_SE_16K model.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, num_channels, height, width].
|
|
|
|
Returns:
|
|
Tuple[torch.Tensor, torch.Tensor]: Output tensors representing the real and imaginary parts.
|
|
"""
|
|
output_real, output_imag = self.model(x) # Get real and imaginary outputs from the model
|
|
return output_real, output_imag # Return the outputs
|
|
|
|
|
|
class FSMN_Wrap(nn.Module):
|
|
"""
|
|
FSMN_Wrap: A wrapper around the UniDeepFsmn module to facilitate
|
|
integration into the larger model architecture.
|
|
|
|
Args:
|
|
nIn (int): Number of input features.
|
|
nHidden (int): Number of hidden features in the FSMN (default is 128).
|
|
lorder (int): Order of the FSMN (default is 20).
|
|
nOut (int): Number of output features (default is 128).
|
|
"""
|
|
|
|
def __init__(self, nIn, nHidden=128, lorder=20, nOut=128):
|
|
"""Initializes the FSMN_Wrap module with specified parameters."""
|
|
super(FSMN_Wrap, self).__init__()
|
|
|
|
# Initialize the UniDeepFsmn module
|
|
self.fsmn = UniDeepFsmn(nIn, nHidden, lorder, nHidden)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass of the FSMN_Wrap module.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, channels, height, time, 2].
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor reshaped to [batch_size, channels, height, time].
|
|
"""
|
|
# Shape of input x: [b, c, h, T, 2]
|
|
b, c, T, h = x.size()
|
|
|
|
# Permute x to reshape it for FSMN processing: [b, T, h, c]
|
|
x = x.permute(0, 2, 3, 1) # Change dimensions to [b, T, h, c]
|
|
x = torch.reshape(x, (b * T, h, c)) # Reshape to [b*T, h, c]
|
|
|
|
# Pass through the FSMN
|
|
output = self.fsmn(x) # output: [b*T, h, c]
|
|
|
|
# Reshape output back to original dimensions
|
|
output = torch.reshape(output, (b, T, h, c)) # output: [b, T, h, c]
|
|
|
|
return output.permute(0, 3, 1, 2) # Final output shape: [b, c, h, T]
|
|
|
|
class DilatedDenseNet(nn.Module):
|
|
"""
|
|
DilatedDenseNet: A dilated dense network for feature extraction.
|
|
|
|
This network consists of a series of dilated convolutions organized in a dense block structure,
|
|
allowing for efficient feature reuse and capturing multi-scale information.
|
|
|
|
Args:
|
|
depth (int): The number of layers in the dense block (default is 4).
|
|
in_channels (int): The number of input channels for the first layer (default is 64).
|
|
"""
|
|
|
|
def __init__(self, depth=4, in_channels=64):
|
|
"""Initializes the DilatedDenseNet with specified depth and input channels."""
|
|
super(DilatedDenseNet, self).__init__()
|
|
self.depth = depth
|
|
self.in_channels = in_channels
|
|
self.pad = nn.ConstantPad2d((1, 1, 1, 0), value=0.) # Padding for the first layer
|
|
self.twidth = 2 # Temporal width for convolutions
|
|
self.kernel_size = (self.twidth, 3) # Kernel size for convolutions
|
|
|
|
# Initialize dilated convolutions, padding, normalization, and FSMN for each layer
|
|
for i in range(self.depth):
|
|
dil = 2 ** i # Dilation factor for the current layer
|
|
pad_length = self.twidth + (dil - 1) * (self.twidth - 1) - 1 # Calculate padding length
|
|
setattr(self, 'pad{}'.format(i + 1), nn.ConstantPad2d((1, 1, pad_length, 0), value=0.))
|
|
setattr(self, 'conv{}'.format(i + 1),
|
|
nn.Conv2d(self.in_channels * (i + 1), self.in_channels, kernel_size=self.kernel_size,
|
|
dilation=(dil, 1))) # Convolution layer
|
|
setattr(self, 'norm{}'.format(i + 1), nn.InstanceNorm2d(in_channels, affine=True)) # Normalization
|
|
setattr(self, 'prelu{}'.format(i + 1), nn.PReLU(self.in_channels)) # Activation function
|
|
setattr(self, 'fsmn{}'.format(i + 1), FSMN_Wrap(nIn=self.in_channels, nHidden=self.in_channels, lorder=5, nOut=self.in_channels))
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass for the DilatedDenseNet.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, channels, height, width].
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after processing through the dense network.
|
|
"""
|
|
skip = x # Initialize skip connection with input
|
|
for i in range(self.depth):
|
|
# Apply padding, convolution, normalization, activation, and FSMN in sequence
|
|
out = getattr(self, 'pad{}'.format(i + 1))(skip)
|
|
out = getattr(self, 'conv{}'.format(i + 1))(out)
|
|
out = getattr(self, 'norm{}'.format(i + 1))(out)
|
|
out = getattr(self, 'prelu{}'.format(i + 1))(out)
|
|
out = getattr(self, 'fsmn{}'.format(i + 1))(out)
|
|
skip = torch.cat([out, skip], dim=1) # Concatenate outputs for dense connectivity
|
|
return out # Return the final output
|
|
|
|
|
|
class DenseEncoder(nn.Module):
|
|
"""
|
|
DenseEncoder: A dense encoding module for feature extraction from input data.
|
|
|
|
This module consists of a series of convolutional layers followed by a
|
|
dilated dense network for robust feature learning.
|
|
|
|
Args:
|
|
in_channel (int): Number of input channels for the encoder.
|
|
channels (int): Number of output channels for each convolutional layer (default is 64).
|
|
"""
|
|
|
|
def __init__(self, in_channel, channels=64):
|
|
"""Initializes the DenseEncoder with specified input channels and feature size."""
|
|
super(DenseEncoder, self).__init__()
|
|
self.conv_1 = nn.Sequential(
|
|
nn.Conv2d(in_channel, channels, (1, 1), (1, 1)), # Initial convolution layer
|
|
nn.InstanceNorm2d(channels, affine=True), # Normalization layer
|
|
nn.PReLU(channels) # Activation function
|
|
)
|
|
self.dilated_dense = DilatedDenseNet(depth=4, in_channels=channels) # Dilated Dense Network
|
|
self.conv_2 = nn.Sequential(
|
|
nn.Conv2d(channels, channels, (1, 3), (1, 2), padding=(0, 1)), # Second convolution layer
|
|
nn.InstanceNorm2d(channels, affine=True), # Normalization layer
|
|
nn.PReLU(channels) # Activation function
|
|
)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass for the DenseEncoder.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, in_channel, height, width].
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after processing through the encoder.
|
|
"""
|
|
x = self.conv_1(x) # Process through the first convolutional layer
|
|
x = self.dilated_dense(x) # Process through the dilated dense network
|
|
x = self.conv_2(x) # Process through the second convolutional layer
|
|
return x # Return the final output
|
|
|
|
|
|
class SPConvTranspose2d(nn.Module):
|
|
"""
|
|
SPConvTranspose2d: A spatially separable convolution transpose layer.
|
|
|
|
This module implements a transposed convolution operation with spatial separability,
|
|
allowing for efficient upsampling and feature extraction.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels.
|
|
out_channels (int): Number of output channels.
|
|
kernel_size (tuple): Size of the convolution kernel.
|
|
r (int): Upsampling rate (default is 1).
|
|
"""
|
|
|
|
def __init__(self, in_channels, out_channels, kernel_size, r=1):
|
|
"""Initializes the SPConvTranspose2d with specified parameters."""
|
|
super(SPConvTranspose2d, self).__init__()
|
|
self.pad1 = nn.ConstantPad2d((1, 1, 0, 0), value=0.) # Padding for input
|
|
self.out_channels = out_channels # Store number of output channels
|
|
self.conv = nn.Conv2d(in_channels, out_channels * r, kernel_size=kernel_size, stride=(1, 1)) # Convolution layer
|
|
self.r = r # Store the upsampling rate
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass for the SPConvTranspose2d module.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, in_channels, height, width].
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor after transposed convolution operation.
|
|
"""
|
|
x = self.pad1(x) # Apply padding to input
|
|
out = self.conv(x) # Perform convolution operation
|
|
batch_size, nchannels, H, W = out.shape # Get output shape
|
|
out = out.view((batch_size, self.r, nchannels // self.r, H, W)) # Reshape output for separation
|
|
out = out.permute(0, 2, 3, 4, 1) # Rearrange dimensions
|
|
out = out.contiguous().view((batch_size, nchannels // self.r, H, -1)) # Final output shape
|
|
return out # Return the final output
|
|
|
|
class MaskDecoder(nn.Module):
|
|
"""
|
|
MaskDecoder: A decoder module for estimating masks used in audio processing.
|
|
|
|
This module utilizes a dilated dense network to capture features and
|
|
applies sub-pixel convolution to upscale the output. It produces
|
|
a mask that can be applied to the magnitude of audio signals.
|
|
|
|
Args:
|
|
num_features (int): The number of features in the output mask.
|
|
num_channel (int): The number of channels in intermediate layers (default is 64).
|
|
out_channel (int): The number of output channels for the final output mask (default is 1).
|
|
"""
|
|
|
|
def __init__(self, num_features, num_channel=64, out_channel=1):
|
|
"""Initializes the MaskDecoder with specified parameters."""
|
|
super(MaskDecoder, self).__init__()
|
|
self.dense_block = DilatedDenseNet(depth=4, in_channels=num_channel) # Dense feature extraction
|
|
self.sub_pixel = SPConvTranspose2d(num_channel, num_channel, (1, 3), 2) # Sub-pixel convolution for upsampling
|
|
self.conv_1 = nn.Conv2d(num_channel, out_channel, (1, 2)) # Convolution layer to produce mask
|
|
self.norm = nn.InstanceNorm2d(out_channel, affine=True) # Normalization layer
|
|
self.prelu = nn.PReLU(out_channel) # Activation function
|
|
self.final_conv = nn.Conv2d(out_channel, out_channel, (1, 1)) # Final convolution layer
|
|
self.prelu_out = nn.PReLU(num_features, init=-0.25) # Final activation for output mask
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass for the MaskDecoder.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, channels, height, width].
|
|
|
|
Returns:
|
|
torch.Tensor: Output mask tensor after processing through the decoder.
|
|
"""
|
|
x = self.dense_block(x) # Feature extraction using dilated dense block
|
|
x = self.sub_pixel(x) # Upsample the features
|
|
x = self.conv_1(x) # Convolution to estimate the mask
|
|
x = self.prelu(self.norm(x)) # Apply normalization and activation
|
|
x = self.final_conv(x).permute(0, 3, 2, 1).squeeze(-1) # Final convolution and rearrangement
|
|
return self.prelu_out(x).permute(0, 2, 1).unsqueeze(1) # Final output shape
|
|
|
|
|
|
class ComplexDecoder(nn.Module):
|
|
"""
|
|
ComplexDecoder: A decoder module for estimating complex-valued outputs.
|
|
|
|
This module processes features through a dilated dense network and a
|
|
sub-pixel convolution layer to generate two output channels representing
|
|
the real and imaginary parts of the complex output.
|
|
|
|
Args:
|
|
num_channel (int): The number of channels in intermediate layers (default is 64).
|
|
"""
|
|
|
|
def __init__(self, num_channel=64):
|
|
"""Initializes the ComplexDecoder with specified parameters."""
|
|
super(ComplexDecoder, self).__init__()
|
|
self.dense_block = DilatedDenseNet(depth=4, in_channels=num_channel) # Dense feature extraction
|
|
self.sub_pixel = SPConvTranspose2d(num_channel, num_channel, (1, 3), 2) # Sub-pixel convolution for upsampling
|
|
self.prelu = nn.PReLU(num_channel) # Activation function
|
|
self.norm = nn.InstanceNorm2d(num_channel, affine=True) # Normalization layer
|
|
self.conv = nn.Conv2d(num_channel, 2, (1, 2)) # Convolution layer to produce complex outputs
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass for the ComplexDecoder.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, channels, height, width].
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor containing real and imaginary parts.
|
|
"""
|
|
x = self.dense_block(x) # Feature extraction using dilated dense block
|
|
x = self.sub_pixel(x) # Upsample the features
|
|
x = self.prelu(self.norm(x)) # Apply normalization and activation
|
|
x = self.conv(x) # Generate complex output
|
|
return x # Return the output tensor
|
|
|
|
|
|
class SyncANet(nn.Module):
|
|
"""
|
|
SyncANet: A synchronous audio processing network for separating audio signals.
|
|
|
|
This network integrates dense encoding, synchronous attention blocks,
|
|
and separate decoders for estimating masks and complex-valued outputs.
|
|
|
|
Args:
|
|
num_channel (int): The number of channels in the network (default is 64).
|
|
num_features (int): The number of features for the mask decoder (default is 201).
|
|
"""
|
|
|
|
def __init__(self, num_channel=64, num_features=201):
|
|
"""Initializes the SyncANet with specified parameters."""
|
|
super(SyncANet, self).__init__()
|
|
self.dense_encoder = DenseEncoder(in_channel=3, channels=num_channel) # Dense encoder for input
|
|
self.n_layers = 6 # Number of synchronous attention layers
|
|
self.blocks = nn.ModuleList([]) # List to hold attention blocks
|
|
|
|
# Initialize attention blocks
|
|
for _ in range(self.n_layers):
|
|
self.blocks.append(
|
|
SyncANetBlock(
|
|
emb_dim=num_channel,
|
|
emb_ks=2,
|
|
emb_hs=1,
|
|
n_freqs=int(num_features//2)+1,
|
|
hidden_channels=num_channel*2,
|
|
n_head=4,
|
|
approx_qk_dim=512,
|
|
activation='prelu',
|
|
eps=1.0e-5,
|
|
)
|
|
)
|
|
|
|
self.mask_decoder = MaskDecoder(num_features, num_channel=num_channel, out_channel=1) # Mask decoder
|
|
self.complex_decoder = ComplexDecoder(num_channel=num_channel) # Complex decoder
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass for the SyncANet.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, 2, height, width] representing complex signals.
|
|
|
|
Returns:
|
|
list: List containing the real and imaginary parts of the output tensor.
|
|
"""
|
|
out_list = [] # List to store outputs
|
|
mag = torch.sqrt(x[:, 0, :, :]**2 + x[:, 1, :, :]**2).unsqueeze(1) # Calculate magnitude
|
|
noisy_phase = torch.angle(torch.complex(x[:, 0, :, :], x[:, 1, :, :])).unsqueeze(1) # Calculate phase
|
|
x_in = torch.cat([mag, x], dim=1) # Concatenate magnitude and input for processing
|
|
|
|
x = self.dense_encoder(x_in) # Feature extraction using dense encoder
|
|
for ii in range(self.n_layers):
|
|
x = self.blocks[ii](x) # Pass through attention blocks
|
|
|
|
mask = self.mask_decoder(x) # Estimate mask from features
|
|
out_mag = mask * mag # Apply mask to magnitude
|
|
|
|
complex_out = self.complex_decoder(x) # Generate complex output
|
|
mag_real = out_mag * torch.cos(noisy_phase) # Real part of the output
|
|
mag_imag = out_mag * torch.sin(noisy_phase) # Imaginary part of the output
|
|
final_real = mag_real + complex_out[:, 0, :, :].unsqueeze(1) # Final real output
|
|
final_imag = mag_imag + complex_out[:, 1, :, :].unsqueeze(1) # Final imaginary output
|
|
out_list.append(final_real) # Append real output to list
|
|
out_list.append(final_imag) # Append imaginary output to list
|
|
|
|
return out_list # Return list of outputs
|
|
|
|
class FFConvM(nn.Module):
|
|
"""
|
|
FFConvM: A feedforward convolutional module combining linear layers, normalization,
|
|
non-linear activation, and convolution operations.
|
|
|
|
This module processes input tensors through a sequence of transformations, including
|
|
normalization, a linear layer with a SiLU activation, a convolutional operation, and
|
|
dropout for regularization.
|
|
|
|
Args:
|
|
dim_in (int): The number of input features (dimensionality of input).
|
|
dim_out (int): The number of output features (dimensionality of output).
|
|
norm_klass (nn.Module): The normalization class to be applied (default is nn.LayerNorm).
|
|
dropout (float): The dropout probability for regularization (default is 0.1).
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim_in,
|
|
dim_out,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1
|
|
):
|
|
"""Initializes the FFConvM with specified parameters."""
|
|
super().__init__()
|
|
|
|
# Define the sequential model
|
|
self.mdl = nn.Sequential(
|
|
norm_klass(dim_in), # Apply normalization to input
|
|
nn.Linear(dim_in, dim_out), # Linear transformation to dim_out
|
|
nn.SiLU(), # Non-linear activation using SiLU (Sigmoid Linear Unit)
|
|
ConvModule(dim_out), # Convolution operation on the output
|
|
nn.Dropout(dropout) # Dropout layer for regularization
|
|
)
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Defines the forward pass for the FFConvM.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [batch_size, dim_in].
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor of shape [batch_size, dim_out] after processing.
|
|
"""
|
|
output = self.mdl(x) # Pass input through the sequential model
|
|
return output # Return the processed output
|
|
|
|
class SyncANetBlock(nn.Module):
|
|
"""
|
|
SyncANetBlock implements a modified version of the MossFormer (GatedFormer) module,
|
|
inspired by the TF-GridNet architecture (https://arxiv.org/abs/2211.12433).
|
|
It combines gated triple-attention schemes and Finite Short Memory Network (FSMN) modules
|
|
to enhance computational efficiency and overall performance in audio processing tasks.
|
|
|
|
Attributes:
|
|
emb_dim (int): Dimensionality of the embedding.
|
|
emb_ks (int): Kernel size for embeddings.
|
|
emb_hs (int): Stride size for embeddings.
|
|
n_freqs (int): Number of frequency bands.
|
|
hidden_channels (int): Number of hidden channels.
|
|
n_head (int): Number of attention heads.
|
|
approx_qk_dim (int): Approximate dimension for query-key matrices.
|
|
activation (str): Activation function to use.
|
|
eps (float): Small value to avoid division by zero in normalization layers.
|
|
"""
|
|
|
|
def __getitem__(self, key):
|
|
"""
|
|
Allows accessing module attributes using indexing.
|
|
|
|
Args:
|
|
key: Attribute name to retrieve.
|
|
|
|
Returns:
|
|
The requested attribute.
|
|
"""
|
|
return getattr(self, key)
|
|
|
|
def __init__(
|
|
self,
|
|
emb_dim,
|
|
emb_ks,
|
|
emb_hs,
|
|
n_freqs,
|
|
hidden_channels,
|
|
n_head=4,
|
|
approx_qk_dim=512,
|
|
activation="prelu",
|
|
eps=1e-5,
|
|
):
|
|
"""
|
|
Initializes the SyncANetBlock with the specified parameters.
|
|
|
|
Args:
|
|
emb_dim (int): Dimensionality of the embedding.
|
|
emb_ks (int): Kernel size for embeddings.
|
|
emb_hs (int): Stride size for embeddings.
|
|
n_freqs (int): Number of frequency bands.
|
|
hidden_channels (int): Number of hidden channels.
|
|
n_head (int): Number of attention heads. Default is 4.
|
|
approx_qk_dim (int): Approximate dimension for query-key matrices. Default is 512.
|
|
activation (str): Activation function to use. Default is "prelu".
|
|
eps (float): Small value to avoid division by zero in normalization layers. Default is 1e-5.
|
|
"""
|
|
super().__init__()
|
|
|
|
in_channels = emb_dim * emb_ks # Calculate the number of input channels
|
|
|
|
## Intra modules: Modules for internal processing within the block
|
|
self.Fconv = nn.Conv2d(emb_dim, in_channels, kernel_size=(1, emb_ks), stride=(1, 1), groups=emb_dim)
|
|
self.intra_norm = LayerNormalization4D(emb_dim, eps=eps) # Layer normalization
|
|
self.intra_to_u = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_channels,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
self.intra_to_v = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_channels,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
self.intra_rnn = self._build_repeats(in_channels, hidden_channels, 20, hidden_channels, repeats=1) # FSMN layers
|
|
self.intra_mossformer = MossFormer(dim=emb_dim, group_size=n_freqs) # MossFormer module
|
|
|
|
# Linear transformation for intra module output
|
|
self.intra_linear = nn.ConvTranspose1d(
|
|
hidden_channels, emb_dim, emb_ks, stride=emb_hs
|
|
)
|
|
self.intra_se = SELayer(channel=emb_dim, reduction=1) # Squeeze-and-excitation layer
|
|
|
|
## Inter modules: Modules for external processing between blocks
|
|
self.inter_norm = LayerNormalization4D(emb_dim, eps=eps) # Layer normalization
|
|
self.inter_to_u = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_channels,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
self.inter_to_v = FFConvM(
|
|
dim_in=in_channels,
|
|
dim_out=hidden_channels,
|
|
norm_klass=nn.LayerNorm,
|
|
dropout=0.1,
|
|
)
|
|
self.inter_rnn = self._build_repeats(in_channels, hidden_channels, 20, hidden_channels, repeats=1) # FSMN layers
|
|
self.inter_mossformer = MossFormer(dim=emb_dim, group_size=256) # MossFormer module
|
|
|
|
# Linear transformation for inter module output
|
|
self.inter_linear = nn.ConvTranspose1d(
|
|
hidden_channels, emb_dim, emb_ks, stride=emb_hs
|
|
)
|
|
self.inter_se = SELayer(channel=emb_dim, reduction=1) # Squeeze-and-excitation layer
|
|
|
|
# Approximate query-key dimension calculation
|
|
E = math.ceil(approx_qk_dim * 1.0 / n_freqs)
|
|
assert emb_dim % n_head == 0 # Ensure emb_dim is divisible by n_head
|
|
|
|
# Define attention convolution layers for each head
|
|
for ii in range(n_head):
|
|
self.add_module(
|
|
f"attn_conv_Q_{ii}",
|
|
nn.Sequential(
|
|
nn.Conv2d(emb_dim, E, 1),
|
|
get_layer(activation)(),
|
|
LayerNormalization4DCF((E, n_freqs), eps=eps),
|
|
),
|
|
)
|
|
self.add_module(
|
|
f"attn_conv_K_{ii}",
|
|
nn.Sequential(
|
|
nn.Conv2d(emb_dim, E, 1),
|
|
get_layer(activation)(),
|
|
LayerNormalization4DCF((E, n_freqs), eps=eps),
|
|
),
|
|
)
|
|
self.add_module(
|
|
f"attn_conv_V_{ii}",
|
|
nn.Sequential(
|
|
nn.Conv2d(emb_dim, emb_dim // n_head, 1),
|
|
get_layer(activation)(),
|
|
LayerNormalization4DCF((emb_dim // n_head, n_freqs), eps=eps),
|
|
),
|
|
)
|
|
|
|
# Final attention concatenation projection
|
|
self.add_module(
|
|
"attn_concat_proj",
|
|
nn.Sequential(
|
|
nn.Conv2d(emb_dim, emb_dim, 1),
|
|
get_layer(activation)(),
|
|
LayerNormalization4DCF((emb_dim, n_freqs), eps=eps),
|
|
),
|
|
)
|
|
|
|
# Store parameters for further processing
|
|
self.emb_dim = emb_dim
|
|
self.emb_ks = emb_ks
|
|
self.emb_hs = emb_hs
|
|
self.n_head = n_head
|
|
|
|
def _build_repeats(self, in_channels, out_channels, lorder, hidden_size, repeats=1):
|
|
"""
|
|
Constructs a sequence of UniDeepFSMN modules.
|
|
|
|
Args:
|
|
in_channels (int): Number of input channels.
|
|
out_channels (int): Number of output channels.
|
|
lorder (int): Order of the filter.
|
|
hidden_size (int): Hidden size for the FSMN.
|
|
repeats (int): Number of times to repeat the module. Default is 1.
|
|
|
|
Returns:
|
|
nn.Sequential: A sequence of UniDeepFSMN modules.
|
|
"""
|
|
repeats = [
|
|
UniDeepFsmn(in_channels, out_channels, lorder, hidden_size)
|
|
for _ in range(repeats)
|
|
]
|
|
return nn.Sequential(*repeats)
|
|
|
|
def forward(self, x):
|
|
"""Performs a forward pass through the SyncANetBlock.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [B, C, T, Q] where
|
|
B is batch size, C is number of channels,
|
|
T is temporal dimension, and Q is frequency dimension.
|
|
|
|
Returns:
|
|
torch.Tensor: Output tensor of the same shape [B, C, T, Q].
|
|
"""
|
|
B, C, old_T, old_Q = x.shape
|
|
|
|
# Calculate new dimensions for padding
|
|
T = math.ceil((old_T - self.emb_ks) / self.emb_hs) * self.emb_hs + self.emb_ks
|
|
Q = math.ceil((old_Q - self.emb_ks) / self.emb_hs) * self.emb_hs + self.emb_ks
|
|
|
|
# Pad the input tensor to match the new dimensions
|
|
x = F.pad(x, (0, Q - old_Q, 0, T - old_T))
|
|
|
|
# Intra-process
|
|
input_ = x
|
|
intra_rnn = self.intra_norm(input_) # Normalize input for intra-process
|
|
intra_rnn = self.Fconv(intra_rnn) # Apply depthwise convolution
|
|
intra_rnn = (
|
|
intra_rnn.transpose(1, 2).contiguous().view(B * T, C * self.emb_ks, -1)
|
|
) # Reshape for subsequent operations
|
|
|
|
intra_rnn = intra_rnn.transpose(1, 2) # Reshape for processing
|
|
intra_rnn_u = self.intra_to_u(intra_rnn) # Linear transformation
|
|
intra_rnn_v = self.intra_to_v(intra_rnn) # Linear transformation
|
|
intra_rnn_u = self.intra_rnn(intra_rnn_u) # Apply FSMN
|
|
intra_rnn = intra_rnn_v * intra_rnn_u # Element-wise multiplication
|
|
intra_rnn = intra_rnn.transpose(1, 2) # Reshape back
|
|
intra_rnn = self.intra_linear(intra_rnn) # Linear projection
|
|
intra_rnn = intra_rnn.transpose(1, 2) # Reshape for mossformer
|
|
intra_rnn = intra_rnn.view([B, T, Q, C]) # Reshape for mossformer
|
|
intra_rnn = self.intra_mossformer(intra_rnn) # Apply MossFormer
|
|
intra_rnn = intra_rnn.transpose(1, 2) # Reshape back
|
|
intra_rnn = intra_rnn.view([B, T, C, Q]) # Reshape back
|
|
intra_rnn = intra_rnn.transpose(1, 2).contiguous() # Final reshape
|
|
intra_rnn = self.intra_se(intra_rnn) # Squeeze-and-excitation layer
|
|
intra_rnn = intra_rnn + input_ # Residual connection
|
|
|
|
# Inter-process
|
|
input_ = intra_rnn
|
|
inter_rnn = self.inter_norm(input_) # Normalize input for inter-process
|
|
inter_rnn = (
|
|
inter_rnn.permute(0, 3, 1, 2).contiguous().view(B * Q, C, T)
|
|
) # Reshape for processing
|
|
inter_rnn = F.unfold(
|
|
inter_rnn[..., None], (self.emb_ks, 1), stride=(self.emb_hs, 1)
|
|
) # Extract sliding windows
|
|
inter_rnn = inter_rnn.transpose(1, 2) # Reshape for further processing
|
|
inter_rnn_u = self.inter_to_u(inter_rnn) # Linear transformation
|
|
inter_rnn_v = self.inter_to_v(inter_rnn) # Linear transformation
|
|
inter_rnn_u = self.inter_rnn(inter_rnn_u) # Apply FSMN
|
|
inter_rnn = inter_rnn_v * inter_rnn_u # Element-wise multiplication
|
|
inter_rnn = inter_rnn.transpose(1, 2) # Reshape back
|
|
inter_rnn = self.inter_linear(inter_rnn) # Linear projection
|
|
inter_rnn = inter_rnn.transpose(1, 2) # Reshape for mossformer
|
|
inter_rnn = inter_rnn.view([B, Q, T, C]) # Reshape for mossformer
|
|
inter_rnn = self.inter_mossformer(inter_rnn) # Apply MossFormer
|
|
inter_rnn = inter_rnn.transpose(1, 2) # Reshape back
|
|
inter_rnn = inter_rnn.view([B, Q, C, T]) # Final reshape
|
|
inter_rnn = inter_rnn.permute(0, 2, 3, 1).contiguous() # Permute for SE layer
|
|
inter_rnn = self.inter_se(inter_rnn) # Squeeze-and-excitation layer
|
|
inter_rnn = inter_rnn + input_ # Residual connection
|
|
|
|
# Attention mechanism
|
|
inter_rnn = inter_rnn[..., :old_T, :old_Q] # Trim to original shape
|
|
|
|
batch = inter_rnn
|
|
all_Q, all_K, all_V = [], [], []
|
|
|
|
# Compute query, key, and value for each attention head
|
|
for ii in range(self.n_head):
|
|
all_Q.append(self["attn_conv_Q_%d" % ii](batch)) # Query
|
|
all_K.append(self["attn_conv_K_%d" % ii](batch)) # Key
|
|
all_V.append(self["attn_conv_V_%d" % ii](batch)) # Value
|
|
|
|
Q = torch.cat(all_Q, dim=0) # Concatenate all queries
|
|
K = torch.cat(all_K, dim=0) # Concatenate all keys
|
|
V = torch.cat(all_V, dim=0) # Concatenate all values
|
|
|
|
# Reshape for attention calculation
|
|
Q = Q.transpose(1, 2)
|
|
Q = Q.flatten(start_dim=2) # Flatten for attention calculation
|
|
K = K.transpose(1, 2)
|
|
K = K.flatten(start_dim=2) # Flatten for attention calculation
|
|
V = V.transpose(1, 2) # Reshape for attention calculation
|
|
old_shape = V.shape
|
|
V = V.flatten(start_dim=2) # Flatten for attention calculation
|
|
emb_dim = Q.shape[-1]
|
|
|
|
# Compute scaled dot-product attention
|
|
attn_mat = torch.matmul(Q, K.transpose(1, 2)) / (emb_dim**0.5) # Attention matrix
|
|
attn_mat = F.softmax(attn_mat, dim=2) # Softmax over attention scores
|
|
V = torch.matmul(attn_mat, V) # Weighted sum of values
|
|
|
|
V = V.reshape(old_shape) # Reshape back
|
|
V = V.transpose(1, 2) # Final reshaping
|
|
emb_dim = V.shape[1]
|
|
|
|
batch = V.view([self.n_head, B, emb_dim, old_T, -1]) # Reshape for multi-head
|
|
batch = batch.transpose(0, 1) # Permute for batch processing
|
|
batch = batch.contiguous().view(
|
|
[B, self.n_head * emb_dim, old_T, -1]
|
|
) # Final reshape for concatenation
|
|
batch = self["attn_concat_proj"](batch) # Final linear projection
|
|
|
|
# Combine inter-process result with attention output
|
|
out = batch + inter_rnn
|
|
return out # Return the output tensor
|
|
|
|
class LayerNormalization4D(nn.Module):
|
|
"""
|
|
LayerNormalization4D applies layer normalization to 4D tensors
|
|
(e.g., [B, C, T, F]), where B is the batch size, C is the number of channels,
|
|
T is the temporal dimension, and F is the frequency dimension.
|
|
|
|
Attributes:
|
|
gamma (torch.Parameter): Learnable scaling parameter.
|
|
beta (torch.Parameter): Learnable shifting parameter.
|
|
eps (float): Small value for numerical stability during variance calculation.
|
|
"""
|
|
|
|
def __init__(self, input_dimension, eps=1e-5):
|
|
"""
|
|
Initializes the LayerNormalization4D layer.
|
|
|
|
Args:
|
|
input_dimension (int): The number of channels in the input tensor.
|
|
eps (float, optional): Small constant added for numerical stability.
|
|
"""
|
|
super().__init__()
|
|
param_size = [1, input_dimension, 1, 1]
|
|
self.gamma = Parameter(torch.Tensor(*param_size).to(torch.float32)) # Scale parameter
|
|
self.beta = Parameter(torch.Tensor(*param_size).to(torch.float32)) # Shift parameter
|
|
init.ones_(self.gamma) # Initialize gamma to 1
|
|
init.zeros_(self.beta) # Initialize beta to 0
|
|
self.eps = eps # Set the epsilon value
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the layer normalization.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [B, C, T, F].
|
|
|
|
Returns:
|
|
torch.Tensor: Normalized output tensor of the same shape.
|
|
"""
|
|
if x.ndim == 4:
|
|
_, C, _, _ = x.shape # Extract the number of channels
|
|
stat_dim = (1,) # Dimension to compute statistics over
|
|
else:
|
|
raise ValueError("Expect x to have 4 dimensions, but got {}".format(x.ndim))
|
|
|
|
# Compute mean and standard deviation along the specified dimension
|
|
mu_ = x.mean(dim=stat_dim, keepdim=True) # [B, 1, T, F]
|
|
std_ = torch.sqrt(
|
|
x.var(dim=stat_dim, unbiased=False, keepdim=True) + self.eps
|
|
) # [B, 1, T, F]
|
|
|
|
# Normalize the input tensor and apply learnable parameters
|
|
x_hat = ((x - mu_) / std_) * self.gamma + self.beta # [B, C, T, F]
|
|
return x_hat
|
|
|
|
class LayerNormalization4DCF(nn.Module):
|
|
"""
|
|
LayerNormalization4DCF applies layer normalization to 4D tensors
|
|
(e.g., [B, C, T, F]) specifically designed for DCF (Dynamic Channel Frequency) inputs.
|
|
|
|
Attributes:
|
|
gamma (torch.Parameter): Learnable scaling parameter.
|
|
beta (torch.Parameter): Learnable shifting parameter.
|
|
eps (float): Small value for numerical stability during variance calculation.
|
|
"""
|
|
|
|
def __init__(self, input_dimension, eps=1e-5):
|
|
"""
|
|
Initializes the LayerNormalization4DCF layer.
|
|
|
|
Args:
|
|
input_dimension (tuple): A tuple containing the dimensions of the input tensor
|
|
(number of channels, frequency dimension).
|
|
eps (float, optional): Small constant added for numerical stability.
|
|
"""
|
|
super().__init__()
|
|
assert len(input_dimension) == 2, "Input dimension must be a tuple of length 2."
|
|
param_size = [1, input_dimension[0], 1, input_dimension[1]] # Shape based on input dimensions
|
|
self.gamma = Parameter(torch.Tensor(*param_size).to(torch.float32)) # Scale parameter
|
|
self.beta = Parameter(torch.Tensor(*param_size).to(torch.float32)) # Shift parameter
|
|
init.ones_(self.gamma) # Initialize gamma to 1
|
|
init.zeros_(self.beta) # Initialize beta to 0
|
|
self.eps = eps # Set the epsilon value
|
|
|
|
def forward(self, x):
|
|
"""
|
|
Forward pass for the layer normalization.
|
|
|
|
Args:
|
|
x (torch.Tensor): Input tensor of shape [B, C, T, F].
|
|
|
|
Returns:
|
|
torch.Tensor: Normalized output tensor of the same shape.
|
|
"""
|
|
if x.ndim == 4:
|
|
stat_dim = (1, 3) # Dimensions to compute statistics over
|
|
else:
|
|
raise ValueError("Expect x to have 4 dimensions, but got {}".format(x.ndim))
|
|
|
|
# Compute mean and standard deviation along the specified dimensions
|
|
mu_ = x.mean(dim=stat_dim, keepdim=True) # [B, 1, T, 1]
|
|
std_ = torch.sqrt(
|
|
x.var(dim=stat_dim, unbiased=False, keepdim=True) + self.eps
|
|
) # [B, 1, T, F]
|
|
|
|
# Normalize the input tensor and apply learnable parameters
|
|
x_hat = ((x - mu_) / std_) * self.gamma + self.beta # [B, C, T, F]
|
|
return x_hat
|