Skip to content

kirad.feature_extraction_whisper

kirad.feature_extraction_whisper

Model stage: raw audio windows -> Whisper log-mel features, computed on the GPU in torch.

Used only by DAM 1 (mar_2024 model.forward_backbone); DAM 2/3 get log-mel from the HF extractor in KintsugiAudioDataset.

Feature extractor class for Whisper in native torch code. Modified from transformers.models.whisper.feature_extraction_whisper.

WhisperTorchFeatureExtractor

Bases: Module

Torch Whisper log-mel extractor: 80 bins, 25 ms window, 10 ms hop, 16 kHz.

Also applies per-window audio normalization and the IDEAL_LOGMEL_ENERGIES profile. Shapes: [ΣW, 480000] -> [ΣW, 80, 3000].

Constructs a Whisper feature extractor.

Args: feature_size (int, optional, defaults to 80): The feature dimension of the extracted features. hop_length (int, optional, defaults to 160): Length of the overlapping windows for the STFT used to obtain the Mel Frequency coefficients. n_fft (int, optional, defaults to 400): Size of the Fourier transform. dither (float, optional, defaults to 0.0): Adds dithering. In other words, adds a small Gaussian noise to each frame. E.g. use 0.0001 to add dithering with a normal distribution centered around 0.0 with standard deviation 0.0001 (assuming [-1,+1] range of raw_speech). The value 0.0 means no dithering. Dithering has similar effect as spectrogram(mel_floor=...). It reduces the high log_mel_fbank values for signals with hard-zero sections, when VAD cutoff is present in the signal. ideal_logmel_energies (bool, optional, defaults to True): Whether to use ideal_logmel_energies normalization device (str, optional, defaults to 'cpu'): Specifies the device for computation of the log-mel spectrogram of audio signals in the _torch_extract_fbank_features method. (e.g., "cpu", "cuda")

Source code in kirad/feature_extraction_whisper.py 
 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
class WhisperTorchFeatureExtractor(torch.nn.Module):
    r"""Torch Whisper log-mel extractor: 80 bins, 25 ms window, 10 ms hop, 16 kHz.

    Also applies per-window audio normalization and the IDEAL_LOGMEL_ENERGIES profile.
    Shapes: [ΣW, 480000] -> [ΣW, 80, 3000].

    Constructs a Whisper feature extractor.

    Args:
        feature_size (`int`, *optional*, defaults to 80):
            The feature dimension of the extracted features.
        hop_length (`int`, *optional*, defaults to 160):
            Length of the overlapping windows for the STFT used to obtain the Mel Frequency coefficients.
        n_fft (`int`, *optional*, defaults to 400):
            Size of the Fourier transform.
        dither (`float`, *optional*, defaults to 0.0):
            Adds dithering. In other words, adds a small Gaussian noise to each frame.
            E.g. use 0.0001 to add dithering with a normal distribution centered
            around 0.0 with standard deviation 0.0001 (assuming [-1,+1] range of raw_speech).
            The value 0.0 means no dithering.
            Dithering has similar effect as `spectrogram(mel_floor=...)`. It reduces
            the high log_mel_fbank values for signals with hard-zero sections,
            when VAD cutoff is present in the signal.
        ideal_logmel_energies (bool, *optional*, defaults to True):
            Whether to use ideal_logmel_energies normalization
        device (`str`, *optional*, defaults to `'cpu'`):
            Specifies the device for computation of the log-mel spectrogram of audio signals in the
            `_torch_extract_fbank_features` method. (e.g., "cpu", "cuda")
    """

    def __init__(
        self,
        feature_size=80,
        hop_length=160,
        n_fft=400,
        dither=0.0,
        normalize_audio=True,
        ideal_logmel_energies=True,
    ):
        """Precompute Hann window and Slaney mel filters; register the energy profile."""
        super().__init__()
        # Store extractor settings and the Hann window used by the STFT.
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.dither = dither
        self.normalize_audio = normalize_audio
        self.ideal_logmel_energies = ideal_logmel_energies
        self.window = torch.hann_window(self.n_fft)
        # Slaney mel filter bank (201 FFT bins -> 80 mel bins) as a torch tensor.
        self.mel_filters = torch.from_numpy(
            mel_filter_bank(
                num_frequency_bins=1 + n_fft // 2,
                num_mel_filters=feature_size,
                min_frequency=0.0,
                # HARDCODED: 16 kHz input, mel filters up to 8 kHz (Whisper setup)
                max_frequency=8000.0,
                sampling_rate=16000,
                norm="slaney",
                mel_scale="slaney",
            )
        )
        # HARDCODED: fixed per-bin target energies from kirad/constants.py (80 values)
        # To enable model compilation, ideal_logmel_energies should be registered as a buffer or Parameter
        self.register_buffer(
            "ideal_logmel_energies_stats",
            torch.from_numpy(IDEAL_LOGMEL_ENERGIES).float(),
        )

    def _torch_extract_fbank_features(self, waveform: torch.Tensor) -> torch.Tensor:
        """Log10 mel spectrogram, clipped to 8 below the max, scaled (x + 4) / 4 as Whisper.

        Shapes: [ΣW, T] -> [ΣW, 80, T // 160] (last STFT frame dropped).

        Compute the log-mel spectrogram of the audio using PyTorch's GPU-accelerated STFT implementation with batching,
        yielding results similar to cpu computing with 1e-5 tolerance.
        """

        # Note: it would be better to dither the chunked waveform,
        # so overlapping signal does not get the same dithering.
        # But, chunking is happening inside pytorch, so it is here.
        if self.dither != 0.0:
            waveform += self.dither * torch.randn(
                waveform.shape, dtype=waveform.dtype, device=waveform.device
            )

        # Complex STFT of each window (n_fft 400, hop 160).
        stft = torch.stft(
            waveform,
            self.n_fft,
            self.hop_length,
            window=self.window.to(waveform),
            return_complex=True,
        )
        # Power spectrogram; drop the last frame so frames = T // 160.
        magnitudes = stft[..., :-1].abs() ** 2

        # Project power onto the 80 mel filters -> [ΣW, 80, frames].
        mel_spec = self.mel_filters.to(magnitudes).T @ magnitudes

        # Log10 with a floor, then clip to 8 below the max (per window when batched).
        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
        if waveform.dim() == 2:
            max_val = log_spec.max(dim=2, keepdim=True)[0].max(dim=1, keepdim=True)[0]
            # HARDCODED: Whisper's 8 (log10) = 80 dB dynamic range and (x + 4) / 4 scaling
            log_spec = torch.maximum(log_spec, max_val - 8.0)
        else:
            log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
        return (log_spec + 4.0) / 4.0

    def forward(
        self,
        batched_speech: torch.Tensor,
    ) -> torch.Tensor:
        """Normalize each window, compute log-mel, shift bins to IDEAL_LOGMEL_ENERGIES.

        normalize_audio: zero mean, unit variance per window (in float64).
        ideal_logmel_energies: add (target - time-mean) per bin.
        Shapes: [ΣW, 480000] -> [ΣW, 80, 3000], returned in the input dtype.

        Compute mel filter bank features from a batch x samples tensor of audio.
        """
        # Got NaNs when doing feature computations in half or single precision
        batched_speech_double = batched_speech.double()

        # zero-mean and unit-variance normalization
        if self.normalize_audio:
            batched_speech_double = (
                batched_speech_double - batched_speech_double.mean(dim=1, keepdim=True)
            ) / torch.sqrt(batched_speech_double.var(dim=1, keepdim=True) + 1e-7)
        # Log-mel features [ΣW, 80, 3000] (waveform normalized first if enabled).
        features = self._torch_extract_fbank_features(batched_speech_double)

        # Shift each mel bin's time-mean to the IDEAL_LOGMEL_ENERGIES target.
        if self.ideal_logmel_energies:
            features += self.ideal_logmel_energies_stats.to(features)[
                None, :, None
            ] - features.mean(dim=-1, keepdim=True)

        return features.to(batched_speech)

n_fft = n_fft instance-attribute

hop_length = hop_length instance-attribute

dither = dither instance-attribute

normalize_audio = normalize_audio instance-attribute

ideal_logmel_energies = ideal_logmel_energies instance-attribute

window = torch.hann_window(self.n_fft) instance-attribute

mel_filters = torch.from_numpy(mel_filter_bank(num_frequency_bins=1 + n_fft // 2, num_mel_filters=feature_size, min_frequency=0.0, max_frequency=8000.0, sampling_rate=16000, norm='slaney', mel_scale='slaney')) instance-attribute

__init__(feature_size=80, hop_length=160, n_fft=400, dither=0.0, normalize_audio=True, ideal_logmel_energies=True)

Precompute Hann window and Slaney mel filters; register the energy profile.

Source code in kirad/feature_extraction_whisper.py 
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
def __init__(
    self,
    feature_size=80,
    hop_length=160,
    n_fft=400,
    dither=0.0,
    normalize_audio=True,
    ideal_logmel_energies=True,
):
    """Precompute Hann window and Slaney mel filters; register the energy profile."""
    super().__init__()
    # Store extractor settings and the Hann window used by the STFT.
    self.n_fft = n_fft
    self.hop_length = hop_length
    self.dither = dither
    self.normalize_audio = normalize_audio
    self.ideal_logmel_energies = ideal_logmel_energies
    self.window = torch.hann_window(self.n_fft)
    # Slaney mel filter bank (201 FFT bins -> 80 mel bins) as a torch tensor.
    self.mel_filters = torch.from_numpy(
        mel_filter_bank(
            num_frequency_bins=1 + n_fft // 2,
            num_mel_filters=feature_size,
            min_frequency=0.0,
            # HARDCODED: 16 kHz input, mel filters up to 8 kHz (Whisper setup)
            max_frequency=8000.0,
            sampling_rate=16000,
            norm="slaney",
            mel_scale="slaney",
        )
    )
    # HARDCODED: fixed per-bin target energies from kirad/constants.py (80 values)
    # To enable model compilation, ideal_logmel_energies should be registered as a buffer or Parameter
    self.register_buffer(
        "ideal_logmel_energies_stats",
        torch.from_numpy(IDEAL_LOGMEL_ENERGIES).float(),
    )

forward(batched_speech)

Normalize each window, compute log-mel, shift bins to IDEAL_LOGMEL_ENERGIES.

normalize_audio: zero mean, unit variance per window (in float64). ideal_logmel_energies: add (target - time-mean) per bin. Shapes: [ΣW, 480000] -> [ΣW, 80, 3000], returned in the input dtype.

Compute mel filter bank features from a batch x samples tensor of audio.

Source code in kirad/feature_extraction_whisper.py 
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
def forward(
    self,
    batched_speech: torch.Tensor,
) -> torch.Tensor:
    """Normalize each window, compute log-mel, shift bins to IDEAL_LOGMEL_ENERGIES.

    normalize_audio: zero mean, unit variance per window (in float64).
    ideal_logmel_energies: add (target - time-mean) per bin.
    Shapes: [ΣW, 480000] -> [ΣW, 80, 3000], returned in the input dtype.

    Compute mel filter bank features from a batch x samples tensor of audio.
    """
    # Got NaNs when doing feature computations in half or single precision
    batched_speech_double = batched_speech.double()

    # zero-mean and unit-variance normalization
    if self.normalize_audio:
        batched_speech_double = (
            batched_speech_double - batched_speech_double.mean(dim=1, keepdim=True)
        ) / torch.sqrt(batched_speech_double.var(dim=1, keepdim=True) + 1e-7)
    # Log-mel features [ΣW, 80, 3000] (waveform normalized first if enabled).
    features = self._torch_extract_fbank_features(batched_speech_double)

    # Shift each mel bin's time-mean to the IDEAL_LOGMEL_ENERGIES target.
    if self.ideal_logmel_energies:
        features += self.ideal_logmel_energies_stats.to(features)[
            None, :, None
        ] - features.mean(dim=-1, keepdim=True)

    return features.to(batched_speech)