|
from typing import Tuple |
|
import torch |
|
from torch import nn |
|
import torch.nn.functional as F |
|
|
|
|
|
class _Conv1DLayer(nn.Module): |
|
def __init__( |
|
self, |
|
in_channels: int, |
|
out_channels: int, |
|
kernel_size: int, |
|
stride: int, |
|
): |
|
""" |
|
Args: |
|
x (Tensor): The ouput. Shape: (batch, in_channels, in_frames) |
|
length (Tensor): The valid length of each sample. Shape: (batch) |
|
|
|
Returns: |
|
x (Tensor): The output. Shape: (batch, out_channels, out_frames) |
|
length (Tensor): The valid length of each sample. Shape: (batch) |
|
""" |
|
super().__init__() |
|
|
|
self.kernel_size = kernel_size |
|
self.stride = stride |
|
|
|
self.conv = nn.Conv1d( |
|
in_channels=in_channels, |
|
out_channels=out_channels, |
|
stride=stride, |
|
kernel_size=kernel_size, |
|
bias=False, |
|
) |
|
|
|
self.layernorm = nn.LayerNorm(out_channels) |
|
|
|
def forward(self, x: torch.Tensor, length: torch.Tensor): |
|
x = self.conv(x) |
|
x = x.transpose_(1, 2) |
|
x = self.layernorm(x) |
|
x = x.transpose_(1, 2) |
|
x = F.gelu(x) |
|
|
|
length = (length - self.kernel_size) // self.stride + 1 |
|
length = length.clamp_min_(min=0) |
|
return x, length |
|
|
|
|
|
class FeatureExtractor(nn.Module): |
|
def __init__(self, config): |
|
""" |
|
Extracts features from the waveform. |
|
|
|
Args: |
|
waveforms (Tensor): The waveform to extract features from. Shape: (batch, wavelength) |
|
wavelength (Tensor): The valid length of each waveform. Shape: (batch) |
|
|
|
Returns: |
|
features (Tensor): The extracted features. Shape: (batch, num_frames, num_channels) |
|
num_frames (Tensor): The valid length of each feature. Shape: (batch) |
|
""" |
|
super().__init__() |
|
|
|
num_channels = config.num_channels |
|
kernel_sizes = config.kernel_sizes |
|
strides = config.strides |
|
|
|
assert ( |
|
len(num_channels) == len(kernel_sizes) == len(strides) |
|
), "The number of layers must be the same for all parameters" |
|
|
|
self.conv_layers = nn.ModuleList( |
|
( |
|
_Conv1DLayer( |
|
in_channels=1, |
|
out_channels=num_channels[0], |
|
kernel_size=kernel_sizes[0], |
|
stride=strides[0], |
|
), |
|
) |
|
) |
|
|
|
for i in range(1, len(num_channels)): |
|
self.conv_layers.append( |
|
_Conv1DLayer( |
|
in_channels=num_channels[i - 1], |
|
out_channels=num_channels[i], |
|
kernel_size=kernel_sizes[i], |
|
stride=strides[i], |
|
) |
|
) |
|
|
|
def forward(self, waveforms: torch.Tensor, wavelength: torch.Tensor): |
|
features = waveforms.unsqueeze(1) |
|
|
|
for conv_layer in self.conv_layers: |
|
features, wavelength = conv_layer(features, wavelength) |
|
|
|
|
|
features = features.transpose(1, 2) |
|
return features, wavelength |
|
|
