mirror of
https://github.com/SYSTRAN/faster-whisper.git
synced 2026-01-09 21:48:08 -05:00
Remove torch dependency, Faster numpy Feature extraction (#1106)
This commit is contained in:
Mahmoud Ashraf
@@ -1,21 +1,15 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
# Adapted from https://github.com/huggingface/transformers/blob/main/src/transformers/models/whisper/feature_extraction_whisper.py # noqa: E501
|
||||
class FeatureExtractor:
|
||||
def __init__(
|
||||
self,
|
||||
device: str = "auto",
|
||||
feature_size=80,
|
||||
sampling_rate=16000,
|
||||
hop_length=160,
|
||||
chunk_length=30,
|
||||
n_fft=400,
|
||||
):
|
||||
if device == "auto":
|
||||
self.device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
else:
|
||||
self.device = device
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
self.chunk_length = chunk_length
|
||||
@@ -25,24 +19,21 @@ class FeatureExtractor:
|
||||
self.sampling_rate = sampling_rate
|
||||
self.mel_filters = self.get_mel_filters(
|
||||
sampling_rate, n_fft, n_mels=feature_size
|
||||
)
|
||||
).astype("float32")
|
||||
|
||||
@staticmethod
|
||||
def get_mel_filters(sr, n_fft, n_mels=128):
|
||||
"""
|
||||
Implementation of librosa.filters.mel in Pytorch
|
||||
"""
|
||||
# Initialize the weights
|
||||
n_mels = int(n_mels)
|
||||
|
||||
# Center freqs of each FFT bin
|
||||
fftfreqs = torch.fft.rfftfreq(n=n_fft, d=1.0 / sr)
|
||||
fftfreqs = np.fft.rfftfreq(n=n_fft, d=1.0 / sr)
|
||||
|
||||
# 'Center freqs' of mel bands - uniformly spaced between limits
|
||||
min_mel = 0.0
|
||||
max_mel = 45.245640471924965
|
||||
|
||||
mels = torch.linspace(min_mel, max_mel, n_mels + 2)
|
||||
mels = np.linspace(min_mel, max_mel, n_mels + 2)
|
||||
|
||||
# Fill in the linear scale
|
||||
f_min = 0.0
|
||||
@@ -52,30 +43,159 @@ class FeatureExtractor:
|
||||
# And now the nonlinear scale
|
||||
min_log_hz = 1000.0 # beginning of log region (Hz)
|
||||
min_log_mel = (min_log_hz - f_min) / f_sp # same (Mels)
|
||||
logstep = torch.log(torch.tensor(6.4)) / 27.0 # step size for log region
|
||||
logstep = np.log(6.4) / 27.0 # step size for log region
|
||||
|
||||
# If we have vector data, vectorize
|
||||
log_t = mels >= min_log_mel
|
||||
freqs[log_t] = min_log_hz * torch.exp(logstep * (mels[log_t] - min_log_mel))
|
||||
freqs[log_t] = min_log_hz * np.exp(logstep * (mels[log_t] - min_log_mel))
|
||||
|
||||
mel_f = freqs
|
||||
fdiff = np.diff(freqs)
|
||||
ramps = freqs.reshape(-1, 1) - fftfreqs.reshape(1, -1)
|
||||
|
||||
fdiff = torch.diff(mel_f)
|
||||
ramps = mel_f.view(-1, 1) - fftfreqs.view(1, -1)
|
||||
|
||||
lower = -ramps[:-2] / fdiff[:-1].unsqueeze(1)
|
||||
upper = ramps[2:] / fdiff[1:].unsqueeze(1)
|
||||
lower = -ramps[:-2] / np.expand_dims(fdiff[:-1], axis=1)
|
||||
upper = ramps[2:] / np.expand_dims(fdiff[1:], axis=1)
|
||||
|
||||
# Intersect them with each other and zero, vectorized across all i
|
||||
weights = torch.maximum(torch.zeros_like(lower), torch.minimum(lower, upper))
|
||||
weights = np.maximum(np.zeros_like(lower), np.minimum(lower, upper))
|
||||
|
||||
# Slaney-style mel is scaled to be approx constant energy per channel
|
||||
enorm = 2.0 / (mel_f[2 : n_mels + 2] - mel_f[:n_mels])
|
||||
weights *= enorm.unsqueeze(1)
|
||||
enorm = 2.0 / (freqs[2 : n_mels + 2] - freqs[:n_mels])
|
||||
weights *= np.expand_dims(enorm, axis=1)
|
||||
|
||||
return weights
|
||||
|
||||
def __call__(self, waveform, padding=True, chunk_length=None, to_cpu=False):
|
||||
@staticmethod
|
||||
def stft(
|
||||
input_array: np.ndarray,
|
||||
n_fft: int,
|
||||
hop_length: int = None,
|
||||
win_length: int = None,
|
||||
window: np.ndarray = None,
|
||||
center: bool = True,
|
||||
mode: str = "reflect",
|
||||
normalized: bool = False,
|
||||
onesided: bool = None,
|
||||
return_complex: bool = None,
|
||||
):
|
||||
# Default initialization for hop_length and win_length
|
||||
hop_length = hop_length if hop_length is not None else n_fft // 4
|
||||
win_length = win_length if win_length is not None else n_fft
|
||||
input_is_complex = np.iscomplexobj(input_array)
|
||||
|
||||
# Determine if the output should be complex
|
||||
return_complex = (
|
||||
return_complex
|
||||
if return_complex is not None
|
||||
else (input_is_complex or (window is not None and np.iscomplexobj(window)))
|
||||
)
|
||||
|
||||
if not return_complex and return_complex is None:
|
||||
raise ValueError(
|
||||
"stft requires the return_complex parameter for real inputs."
|
||||
)
|
||||
|
||||
# Input checks
|
||||
if not np.issubdtype(input_array.dtype, np.floating) and not input_is_complex:
|
||||
raise ValueError(
|
||||
"stft: expected an array of floating point or complex values,"
|
||||
f" got {input_array.dtype}"
|
||||
)
|
||||
|
||||
if input_array.ndim > 2 or input_array.ndim < 1:
|
||||
raise ValueError(
|
||||
f"stft: expected a 1D or 2D array, but got {input_array.ndim}D array"
|
||||
)
|
||||
|
||||
# Handle 1D input
|
||||
if input_array.ndim == 1:
|
||||
input_array = np.expand_dims(input_array, axis=0)
|
||||
input_array_1d = True
|
||||
else:
|
||||
input_array_1d = False
|
||||
|
||||
# Center padding if required
|
||||
if center:
|
||||
pad_amount = n_fft // 2
|
||||
input_array = np.pad(
|
||||
input_array, ((0, 0), (pad_amount, pad_amount)), mode=mode
|
||||
)
|
||||
|
||||
batch, length = input_array.shape
|
||||
|
||||
# Additional input checks
|
||||
if n_fft <= 0 or n_fft > length:
|
||||
raise ValueError(
|
||||
f"stft: expected 0 < n_fft <= {length}, but got n_fft={n_fft}"
|
||||
)
|
||||
|
||||
if hop_length <= 0:
|
||||
raise ValueError(
|
||||
f"stft: expected hop_length > 0, but got hop_length={hop_length}"
|
||||
)
|
||||
|
||||
if win_length <= 0 or win_length > n_fft:
|
||||
raise ValueError(
|
||||
f"stft: expected 0 < win_length <= n_fft, but got win_length={win_length}"
|
||||
)
|
||||
|
||||
if window is not None:
|
||||
if window.ndim != 1 or window.shape[0] != win_length:
|
||||
raise ValueError(
|
||||
f"stft: expected a 1D window array of size equal to win_length={win_length}, "
|
||||
f"but got window with size {window.shape}"
|
||||
)
|
||||
|
||||
# Handle padding of the window if necessary
|
||||
if win_length < n_fft:
|
||||
left = (n_fft - win_length) // 2
|
||||
window_ = np.zeros(n_fft, dtype=window.dtype)
|
||||
window_[left : left + win_length] = window
|
||||
else:
|
||||
window_ = window
|
||||
|
||||
# Calculate the number of frames
|
||||
n_frames = 1 + (length - n_fft) // hop_length
|
||||
|
||||
# Time to columns
|
||||
input_array = np.lib.stride_tricks.as_strided(
|
||||
input_array,
|
||||
(batch, n_frames, n_fft),
|
||||
(
|
||||
input_array.strides[0],
|
||||
hop_length * input_array.strides[1],
|
||||
input_array.strides[1],
|
||||
),
|
||||
)
|
||||
|
||||
if window_ is not None:
|
||||
input_array = input_array * window_
|
||||
|
||||
# FFT and transpose
|
||||
complex_fft = input_is_complex
|
||||
onesided = onesided if onesided is not None else not complex_fft
|
||||
|
||||
if normalized:
|
||||
norm = "ortho"
|
||||
else:
|
||||
norm = None
|
||||
|
||||
if complex_fft:
|
||||
if onesided:
|
||||
raise ValueError(
|
||||
"Cannot have onesided output if window or input is complex"
|
||||
)
|
||||
output = np.fft.fft(input_array, n=n_fft, axis=-1, norm=norm)
|
||||
else:
|
||||
output = np.fft.rfft(input_array, n=n_fft, axis=-1, norm=norm)
|
||||
|
||||
output = output.transpose((0, 2, 1))
|
||||
|
||||
if input_array_1d:
|
||||
output = output.squeeze(0)
|
||||
|
||||
return output if return_complex else np.real(output)
|
||||
|
||||
def __call__(self, waveform: np.ndarray, padding=160, chunk_length=None):
|
||||
"""
|
||||
Compute the log-Mel spectrogram of the provided audio.
|
||||
"""
|
||||
@@ -84,31 +204,27 @@ class FeatureExtractor:
|
||||
self.n_samples = chunk_length * self.sampling_rate
|
||||
self.nb_max_frames = self.n_samples // self.hop_length
|
||||
|
||||
if waveform.dtype is not torch.float32:
|
||||
waveform = waveform.to(torch.float32)
|
||||
|
||||
waveform = (
|
||||
waveform.to(self.device)
|
||||
if self.device == "cuda" and not waveform.is_cuda
|
||||
else waveform
|
||||
)
|
||||
if waveform.dtype is not np.float32:
|
||||
waveform = waveform.astype(np.float32)
|
||||
|
||||
if padding:
|
||||
waveform = torch.nn.functional.pad(waveform, (0, self.n_samples))
|
||||
waveform = np.pad(waveform, (0, padding))
|
||||
|
||||
window = torch.hann_window(self.n_fft).to(waveform.device)
|
||||
window = np.hanning(self.n_fft + 1)[:-1].astype("float32")
|
||||
|
||||
stft = torch.stft(
|
||||
waveform, self.n_fft, self.hop_length, window=window, return_complex=True
|
||||
)
|
||||
magnitudes = stft[..., :-1].abs() ** 2
|
||||
stft = self.stft(
|
||||
waveform,
|
||||
self.n_fft,
|
||||
self.hop_length,
|
||||
window=window,
|
||||
return_complex=True,
|
||||
).astype("complex64")
|
||||
magnitudes = np.abs(stft[..., :-1]) ** 2
|
||||
|
||||
mel_spec = self.mel_filters.to(waveform.device) @ magnitudes
|
||||
mel_spec = self.mel_filters @ magnitudes
|
||||
|
||||
log_spec = torch.clamp(mel_spec, min=1e-10).log10()
|
||||
log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
|
||||
log_spec = np.log10(np.clip(mel_spec, a_min=1e-10, a_max=None))
|
||||
log_spec = np.maximum(log_spec, log_spec.max() - 8.0)
|
||||
log_spec = (log_spec + 4.0) / 4.0
|
||||
|
||||
# When the model is running on multiple GPUs, the output should be moved
|
||||
# to the CPU since we don't know which GPU will handle the next job.
|
||||
return log_spec.cpu() if to_cpu else log_spec
|
||||
return log_spec
|
||||
|
||||
Reference in New Issue
Block a user