File size: 1,751 Bytes
7ef7abb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | 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
|