|
from typing import List, Tuple
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from models.scnet_unofficial.utils import create_intervals
|
|
|
|
|
|
class Downsample(nn.Module):
|
|
"""
|
|
Downsample class implements a module for downsampling input tensors using 2D convolution.
|
|
|
|
Args:
|
|
- input_dim (int): Dimensionality of the input channels.
|
|
- output_dim (int): Dimensionality of the output channels.
|
|
- stride (int): Stride value for the convolution operation.
|
|
|
|
Shapes:
|
|
- Input: (B, C_in, F, T) where
|
|
B is batch size,
|
|
C_in is the number of input channels,
|
|
F is the frequency dimension,
|
|
T is the time dimension.
|
|
- Output: (B, C_out, F // stride, T) where
|
|
B is batch size,
|
|
C_out is the number of output channels,
|
|
F // stride is the downsampled frequency dimension.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
output_dim: int,
|
|
stride: int,
|
|
):
|
|
"""
|
|
Initializes Downsample with input dimension, output dimension, and stride.
|
|
"""
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(input_dim, output_dim, 1, (stride, 1))
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Performs forward pass through the Downsample module.
|
|
|
|
Args:
|
|
- x (torch.Tensor): Input tensor of shape (B, C_in, F, T).
|
|
|
|
Returns:
|
|
- torch.Tensor: Downsampled tensor of shape (B, C_out, F // stride, T).
|
|
"""
|
|
return self.conv(x)
|
|
|
|
|
|
class ConvolutionModule(nn.Module):
|
|
"""
|
|
ConvolutionModule class implements a module with a sequence of convolutional layers similar to Conformer.
|
|
|
|
Args:
|
|
- input_dim (int): Dimensionality of the input features.
|
|
- hidden_dim (int): Dimensionality of the hidden features.
|
|
- kernel_sizes (List[int]): List of kernel sizes for the convolutional layers.
|
|
- bias (bool, optional): If True, adds a learnable bias to the output. Default is False.
|
|
|
|
Shapes:
|
|
- Input: (B, T, D) where
|
|
B is batch size,
|
|
T is sequence length,
|
|
D is input dimensionality.
|
|
- Output: (B, T, D) where
|
|
B is batch size,
|
|
T is sequence length,
|
|
D is input dimensionality.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
hidden_dim: int,
|
|
kernel_sizes: List[int],
|
|
bias: bool = False,
|
|
) -> None:
|
|
"""
|
|
Initializes ConvolutionModule with input dimension, hidden dimension, kernel sizes, and bias.
|
|
"""
|
|
super().__init__()
|
|
self.sequential = nn.Sequential(
|
|
nn.GroupNorm(num_groups=1, num_channels=input_dim),
|
|
nn.Conv1d(
|
|
input_dim,
|
|
2 * hidden_dim,
|
|
kernel_sizes[0],
|
|
stride=1,
|
|
padding=(kernel_sizes[0] - 1) // 2,
|
|
bias=bias,
|
|
),
|
|
nn.GLU(dim=1),
|
|
nn.Conv1d(
|
|
hidden_dim,
|
|
hidden_dim,
|
|
kernel_sizes[1],
|
|
stride=1,
|
|
padding=(kernel_sizes[1] - 1) // 2,
|
|
groups=hidden_dim,
|
|
bias=bias,
|
|
),
|
|
nn.GroupNorm(num_groups=1, num_channels=hidden_dim),
|
|
nn.SiLU(),
|
|
nn.Conv1d(
|
|
hidden_dim,
|
|
input_dim,
|
|
kernel_sizes[2],
|
|
stride=1,
|
|
padding=(kernel_sizes[2] - 1) // 2,
|
|
bias=bias,
|
|
),
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Performs forward pass through the ConvolutionModule.
|
|
|
|
Args:
|
|
- x (torch.Tensor): Input tensor of shape (B, T, D).
|
|
|
|
Returns:
|
|
- torch.Tensor: Output tensor of shape (B, T, D).
|
|
"""
|
|
x = x.transpose(1, 2)
|
|
x = x + self.sequential(x)
|
|
x = x.transpose(1, 2)
|
|
return x
|
|
|
|
|
|
class SDLayer(nn.Module):
|
|
"""
|
|
SDLayer class implements a subband decomposition layer with downsampling and convolutional modules.
|
|
|
|
Args:
|
|
- subband_interval (Tuple[float, float]): Tuple representing the frequency interval for subband decomposition.
|
|
- input_dim (int): Dimensionality of the input channels.
|
|
- output_dim (int): Dimensionality of the output channels after downsampling.
|
|
- downsample_stride (int): Stride value for the downsampling operation.
|
|
- n_conv_modules (int): Number of convolutional modules.
|
|
- kernel_sizes (List[int]): List of kernel sizes for the convolutional layers.
|
|
- bias (bool, optional): If True, adds a learnable bias to the convolutional layers. Default is True.
|
|
|
|
Shapes:
|
|
- Input: (B, Fi, T, Ci) where
|
|
B is batch size,
|
|
Fi is the number of input subbands,
|
|
T is sequence length, and
|
|
Ci is the number of input channels.
|
|
- Output: (B, Fi+1, T, Ci+1) where
|
|
B is batch size,
|
|
Fi+1 is the number of output subbands,
|
|
T is sequence length,
|
|
Ci+1 is the number of output channels.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
subband_interval: Tuple[float, float],
|
|
input_dim: int,
|
|
output_dim: int,
|
|
downsample_stride: int,
|
|
n_conv_modules: int,
|
|
kernel_sizes: List[int],
|
|
bias: bool = True,
|
|
):
|
|
"""
|
|
Initializes SDLayer with subband interval, input dimension,
|
|
output dimension, downsample stride, number of convolutional modules, kernel sizes, and bias.
|
|
"""
|
|
super().__init__()
|
|
self.subband_interval = subband_interval
|
|
self.downsample = Downsample(input_dim, output_dim, downsample_stride)
|
|
self.activation = nn.GELU()
|
|
conv_modules = [
|
|
ConvolutionModule(
|
|
input_dim=output_dim,
|
|
hidden_dim=output_dim // 4,
|
|
kernel_sizes=kernel_sizes,
|
|
bias=bias,
|
|
)
|
|
for _ in range(n_conv_modules)
|
|
]
|
|
self.conv_modules = nn.Sequential(*conv_modules)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Performs forward pass through the SDLayer.
|
|
|
|
Args:
|
|
- x (torch.Tensor): Input tensor of shape (B, Fi, T, Ci).
|
|
|
|
Returns:
|
|
- torch.Tensor: Output tensor of shape (B, Fi+1, T, Ci+1).
|
|
"""
|
|
B, F, T, C = x.shape
|
|
x = x[:, int(self.subband_interval[0] * F) : int(self.subband_interval[1] * F)]
|
|
x = x.permute(0, 3, 1, 2)
|
|
x = self.downsample(x)
|
|
x = self.activation(x)
|
|
x = x.permute(0, 2, 3, 1)
|
|
|
|
B, F, T, C = x.shape
|
|
x = x.reshape((B * F), T, C)
|
|
x = self.conv_modules(x)
|
|
x = x.reshape(B, F, T, C)
|
|
|
|
return x
|
|
|
|
|
|
class SDBlock(nn.Module):
|
|
"""
|
|
SDBlock class implements a block with subband decomposition layers and global convolution.
|
|
|
|
Args:
|
|
- input_dim (int): Dimensionality of the input channels.
|
|
- output_dim (int): Dimensionality of the output channels.
|
|
- bandsplit_ratios (List[float]): List of ratios for splitting the frequency bands.
|
|
- downsample_strides (List[int]): List of stride values for downsampling in each subband layer.
|
|
- n_conv_modules (List[int]): List specifying the number of convolutional modules in each subband layer.
|
|
- kernel_sizes (List[int], optional): List of kernel sizes for the convolutional layers. Default is None.
|
|
|
|
Shapes:
|
|
- Input: (B, Fi, T, Ci) where
|
|
B is batch size,
|
|
Fi is the number of input subbands,
|
|
T is sequence length,
|
|
Ci is the number of input channels.
|
|
- Output: (B, Fi+1, T, Ci+1) where
|
|
B is batch size,
|
|
Fi+1 is the number of output subbands,
|
|
T is sequence length,
|
|
Ci+1 is the number of output channels.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
input_dim: int,
|
|
output_dim: int,
|
|
bandsplit_ratios: List[float],
|
|
downsample_strides: List[int],
|
|
n_conv_modules: List[int],
|
|
kernel_sizes: List[int] = None,
|
|
):
|
|
"""
|
|
Initializes SDBlock with input dimension, output dimension, band split ratios, downsample strides, number of convolutional modules, and kernel sizes.
|
|
"""
|
|
super().__init__()
|
|
if kernel_sizes is None:
|
|
kernel_sizes = [3, 3, 1]
|
|
assert sum(bandsplit_ratios) == 1, "The split ratios must sum up to 1."
|
|
subband_intervals = create_intervals(bandsplit_ratios)
|
|
self.sd_layers = nn.ModuleList(
|
|
SDLayer(
|
|
input_dim=input_dim,
|
|
output_dim=output_dim,
|
|
subband_interval=sbi,
|
|
downsample_stride=dss,
|
|
n_conv_modules=ncm,
|
|
kernel_sizes=kernel_sizes,
|
|
)
|
|
for sbi, dss, ncm in zip(
|
|
subband_intervals, downsample_strides, n_conv_modules
|
|
)
|
|
)
|
|
self.global_conv2d = nn.Conv2d(output_dim, output_dim, 1, 1)
|
|
|
|
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
|
"""
|
|
Performs forward pass through the SDBlock.
|
|
|
|
Args:
|
|
- x (torch.Tensor): Input tensor of shape (B, Fi, T, Ci).
|
|
|
|
Returns:
|
|
- Tuple[torch.Tensor, torch.Tensor]: Output tensor and skip connection tensor.
|
|
"""
|
|
x_skip = torch.concat([layer(x) for layer in self.sd_layers], dim=1)
|
|
x = self.global_conv2d(x_skip.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
|
|
return x, x_skip
|
|
|