File size: 8,218 Bytes
35d483e | 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 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 | """Torch-native waveform to log-mel feature extraction.
The frontend intentionally avoids torchaudio so the deployable student has one
fewer binary dependency. Exported production models normally accept log-mel
features; keeping this implementation in the repository gives training, demo,
and parity tests one canonical preprocessing contract.
"""
from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass
import torch
import torch.nn.functional as F
from torch import Tensor, nn
@dataclass(frozen=True)
class LogMelConfig:
sample_rate: int = 16_000
n_fft: int = 400
hop_length: int = 160
win_length: int = 400
n_mels: int = 80
f_min: float = 0.0
f_max: float | None = 8_000.0
log_floor: float = 1e-10
normalize: bool = True
mel_scale: str = "htk"
log_scale: str = "standard"
center: bool = False
drop_last_frame: bool = False
pad_side: str = "right"
@classmethod
def from_mapping(cls, values: Mapping[str, object]) -> LogMelConfig:
known = {field.name for field in cls.__dataclass_fields__.values()}
return cls(**{key: value for key, value in values.items() if key in known}) # type: ignore[arg-type]
def _hz_to_mel(freq: Tensor) -> Tensor:
# HTK mel convention. It is stable, simple, and matches common speech
# frontends closely enough for a model trained with the same extractor.
return 2595.0 * torch.log10(1.0 + freq / 700.0)
def _mel_to_hz(mels: Tensor) -> Tensor:
return 700.0 * (torch.pow(10.0, mels / 2595.0) - 1.0)
def _hz_to_slaney_mel(freq: Tensor) -> Tensor:
linear_spacing = 200.0 / 3.0
min_log_hz = 1_000.0
min_log_mel = min_log_hz / linear_spacing
log_step = torch.log(torch.tensor(6.4, dtype=freq.dtype, device=freq.device)) / 27.0
linear = freq / linear_spacing
logarithmic = min_log_mel + torch.log((freq / min_log_hz).clamp_min(1e-12)) / log_step
return torch.where(freq >= min_log_hz, logarithmic, linear)
def _slaney_mel_to_hz(mels: Tensor) -> Tensor:
linear_spacing = 200.0 / 3.0
min_log_hz = 1_000.0
min_log_mel = min_log_hz / linear_spacing
log_step = torch.log(torch.tensor(6.4, dtype=mels.dtype, device=mels.device)) / 27.0
linear = mels * linear_spacing
logarithmic = min_log_hz * torch.exp(log_step * (mels - min_log_mel))
return torch.where(mels >= min_log_mel, logarithmic, linear)
def make_mel_filterbank(config: LogMelConfig) -> Tensor:
"""Create a ``[n_mels, n_fft // 2 + 1]`` triangular filterbank."""
max_hz = float(config.sample_rate / 2 if config.f_max is None else config.f_max)
if not 0.0 <= config.f_min < max_hz <= config.sample_rate / 2:
raise ValueError("expected 0 <= f_min < f_max <= sample_rate / 2")
fft_freqs = torch.linspace(0.0, config.sample_rate / 2, config.n_fft // 2 + 1)
if config.mel_scale == "htk":
to_mel, to_hz = _hz_to_mel, _mel_to_hz
elif config.mel_scale == "slaney":
to_mel, to_hz = _hz_to_slaney_mel, _slaney_mel_to_hz
else:
raise ValueError("mel_scale must be 'htk' or 'slaney'")
mel_edges = torch.linspace(
to_mel(torch.tensor(float(config.f_min))),
to_mel(torch.tensor(max_hz)),
config.n_mels + 2,
)
hz_edges = to_hz(mel_edges)
lower = hz_edges[:-2, None]
center = hz_edges[1:-1, None]
upper = hz_edges[2:, None]
rising = (fft_freqs[None, :] - lower) / (center - lower).clamp_min(1e-12)
falling = (upper - fft_freqs[None, :]) / (upper - center).clamp_min(1e-12)
filters = torch.minimum(rising, falling).clamp_min(0.0)
# Area normalization reduces sensitivity to mel-band width.
enorm = 2.0 / (upper - lower).clamp_min(1e-12)
return filters * enorm
class LogMelFrontend(nn.Module):
"""Convert padded mono waveforms to normalized log-mel features.
Parameters
----------
waveforms:
Float tensor shaped ``[batch, samples]`` (or ``[samples]``).
lengths:
Optional valid sample counts. The returned mask is ``[batch, frames]``.
"""
def __init__(self, config: LogMelConfig | None = None) -> None:
super().__init__()
config = config or LogMelConfig()
self.config = config
if config.pad_side not in {"left", "right"}:
raise ValueError("pad_side must be 'left' or 'right'")
# numpy.hanning in the dependency-light runtime uses a symmetric window.
self.register_buffer(
"window", torch.hann_window(config.win_length, periodic=False), persistent=False
)
self.register_buffer("mel_filters", make_mel_filterbank(config), persistent=True)
def forward(self, waveforms: Tensor, lengths: Tensor | None = None) -> tuple[Tensor, Tensor]:
if waveforms.ndim == 1:
waveforms = waveforms.unsqueeze(0)
if waveforms.ndim != 2:
raise ValueError("waveforms must have shape [batch, samples]")
batch, original_samples = waveforms.shape
if lengths is None:
lengths = torch.full(
(batch,), original_samples, dtype=torch.long, device=waveforms.device
)
else:
lengths = lengths.to(device=waveforms.device, dtype=torch.long).clamp(
min=0, max=original_samples
)
if original_samples < self.config.n_fft:
waveforms = F.pad(waveforms, (0, self.config.n_fft - original_samples))
spectrum = torch.stft(
waveforms,
n_fft=self.config.n_fft,
hop_length=self.config.hop_length,
win_length=self.config.win_length,
window=self.window.to(dtype=waveforms.dtype),
center=self.config.center,
return_complex=True,
)
if self.config.drop_last_frame:
spectrum = spectrum[..., :-1]
power = spectrum.abs().square()
mel = torch.matmul(self.mel_filters.to(dtype=power.dtype), power)
log_mel = torch.log10(mel.clamp_min(self.config.log_floor))
if self.config.log_scale == "whisper":
dynamic_floor = log_mel.amax(dim=(-2, -1), keepdim=True) - 8.0
log_mel = torch.maximum(log_mel, dynamic_floor)
log_mel = (log_mel + 4.0) / 4.0
elif self.config.log_scale != "standard":
raise ValueError("log_scale must be 'standard' or 'whisper'")
if self.config.center:
if self.config.drop_last_frame:
frame_lengths = torch.div(
lengths + self.config.hop_length - 1,
self.config.hop_length,
rounding_mode="floor",
)
else:
frame_lengths = 1 + torch.div(
lengths, self.config.hop_length, rounding_mode="floor"
)
else:
padded_lengths = lengths.clamp_min(self.config.n_fft)
frame_lengths = 1 + torch.div(
padded_lengths - self.config.n_fft,
self.config.hop_length,
rounding_mode="floor",
)
if self.config.drop_last_frame and not self.config.center:
frame_lengths = (frame_lengths - 1).clamp_min(0)
frame_lengths = torch.where(lengths > 0, frame_lengths, torch.zeros_like(frame_lengths))
frame_lengths = frame_lengths.clamp(max=log_mel.shape[-1])
positions = torch.arange(log_mel.shape[-1], device=waveforms.device)
if self.config.pad_side == "left":
mask = positions.unsqueeze(0) >= (log_mel.shape[-1] - frame_lengths).unsqueeze(1)
else:
mask = positions.unsqueeze(0) < frame_lengths.unsqueeze(1)
if self.config.normalize:
valid = mask.unsqueeze(1).to(log_mel.dtype)
denominator = valid.sum(dim=-1, keepdim=True).clamp_min(1.0)
mean = (log_mel * valid).sum(dim=-1, keepdim=True) / denominator
variance = ((log_mel - mean).square() * valid).sum(dim=-1, keepdim=True)
variance = variance / denominator
log_mel = (log_mel - mean) * torch.rsqrt(variance + 1e-5)
log_mel = log_mel * valid
return log_mel, mask
|