| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| from nnAudio import features |
|
|
|
|
| class MelSpectrogram(nn.Module): |
| def __init__( |
| self, |
| sample_rate: int = 16000, |
| n_ftt: int = 2048, |
| n_mels: int = 512, |
| hop_length: int = 128, |
| ): |
| """ |
| Melspectrogram transformation layer, supports on-the-fly processing on GPU. |
| |
| Attributes: |
| sample_rate: The sampling rate for the input audio. |
| n_ftt: The window size for the STFT. |
| n_mels: The number of Mel filter banks. |
| hop_length: The hop (or stride) size. |
| """ |
| super().__init__() |
| self.transform = features.MelSpectrogram( |
| sr=sample_rate, |
| n_fft=n_ftt, |
| n_mels=n_mels, |
| hop_length=hop_length, |
| center=True, |
| fmin=0, |
| fmax=sample_rate // 2, |
| pad_mode="constant", |
| ) |
|
|
| def forward(self, samples: torch.tensor) -> torch.tensor: |
| """ |
| Convert a batch of audio frames into a batch of Mel spectrogram frames. |
| |
| For each item in the batch: |
| 1. pad left and right ends of audio by n_fft // 2. |
| 2. run STFT with window size of |n_ftt| and stride of |hop_length|. |
| 3. convert result into mel-scale. |
| 4. therefore, n_frames = n_samples // hop_length + 1. |
| |
| Args: |
| samples: Audio time-series (batch size, n_samples). |
| |
| Returns: |
| A batch of Mel spectrograms of size (batch size, n_frames, n_mels). |
| """ |
| spectrogram = self.transform(samples) |
| spectrogram = spectrogram.permute(0, 2, 1) |
| return spectrogram |
|
|
