From 041c9419d27a00ad8b946f602e164ff1ac332361 Mon Sep 17 00:00:00 2001 From: "google-labs-jules[bot]" <161369871+google-labs-jules[bot]@users.noreply.github.com> Date: Sun, 21 Sep 2025 06:23:41 +0000 Subject: [PATCH 1/5] feat: make audio feature extractors torch.export-able Fix #40986 Refactors the feature extraction logic for both `GraniteSpeechFeatureExtractor` and `WhisperFeatureExtractor` into separate `nn.Module` subclasses. This encapsulates the feature extraction computation in a way that is compatible with `torch.export`. A new method, `to_exportable_module()`, was added to both feature extractor classes to return an instance of these new modules. The original `__call__` and `_torch_extract_fbank_features` methods were updated to use these exportable modules. Finally, new tests were added to verify that both feature extractors can be successfully exported using `torch.export`. --- .../feature_extraction_granite_speech.py | 79 ++++++++++++++----- .../whisper/feature_extraction_whisper.py | 79 +++++++++++-------- tests/test_executorch.py | 30 ++++++- 3 files changed, 137 insertions(+), 51 deletions(-) diff --git a/src/transformers/models/granite_speech/feature_extraction_granite_speech.py b/src/transformers/models/granite_speech/feature_extraction_granite_speech.py index 7528fc7ea5bd..b01f053ce569 100644 --- a/src/transformers/models/granite_speech/feature_extraction_granite_speech.py +++ b/src/transformers/models/granite_speech/feature_extraction_granite_speech.py @@ -30,11 +30,61 @@ if is_torch_available(): import torch + from torch import nn if is_torchaudio_available(): import torchaudio +class _GraniteSpeechFeatureExtractorModule(nn.Module): + def __init__(self, feature_extractor: "GraniteSpeechFeatureExtractor"): + super().__init__() + self.melspec_kwargs = feature_extractor.melspec_kwargs + self.mel_filters = torchaudio.transforms.MelSpectrogram(**self.melspec_kwargs) + self.projector_window_size = feature_extractor.projector_window_size + self.projector_downsample_rate = feature_extractor.projector_downsample_rate + + def _get_num_audio_features(self, audio_lengths: "torch.Tensor") -> "torch.Tensor": + """ + Gets the (variable length) number of features (i.e., projector output) for the sequences + being considered. + + Args: + audio_lengths (`torch.Tensor`): + Sequence of one or more raw audio lengths. + """ + hop_length = self.melspec_kwargs["hop_length"] + effective_window_size = self.projector_window_size // self.projector_downsample_rate + + # mel sequence length computation + mel_length = audio_lengths // hop_length + 1 + # encoder frame takes two mel features + encoder_length = mel_length // 2 + nblocks = (encoder_length + self.projector_window_size - 1) // self.projector_window_size + # projector output length + projector_length = nblocks * effective_window_size + return projector_length + + def forward(self, audio: "torch.Tensor"): + """ + Compute the Mel features to be passed to the conformer encoder. + """ + bsz = audio.shape[0] + # Compute mel features + mel = self.mel_filters(audio.float()) + logmel = mel.transpose(-1, -2).clip_(min=1e-10).log10_() + mx = logmel.amax(dim=(-2, -1), keepdim=True) + logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1) + # remove last frame if odd + if logmel.shape[1] % 2 == 1: + logmel = logmel[:, :-1] + + # stacking and skipping by 2 + audio = logmel.reshape(bsz, -1, 2 * logmel.shape[-1]) + + return audio + + class GraniteSpeechFeatureExtractor(FeatureExtractionMixin): model_input_names = ["input_features"] @@ -59,10 +109,16 @@ def __init__( "n_mels": n_mels, } requires_backends(self, ["torchaudio"]) - self.mel_filters = torchaudio.transforms.MelSpectrogram(**self.melspec_kwargs) self.projector_window_size = projector_window_size self.projector_downsample_rate = projector_downsample_rate + def to_exportable_module(self) -> "nn.Module": + """ + Returns an exportable version of the feature extractor, which can be used with `torch.export`. + """ + requires_backends(self, "torch") + return _GraniteSpeechFeatureExtractorModule(self) + def __call__( self, audios: AudioInput, @@ -96,27 +152,13 @@ def _extract_mel_spectrograms(self, audio: "torch.Tensor", device="cpu"): """ Compute the Mel features to be passed to the conformer encoder. """ - requires_backends(self, ["torchaudio"]) + module = self.to_exportable_module() if device is not None: - melspec = self.mel_filters.to(device) + module = module.to(device) audio = audio.to(device) - else: - melspec = self.mel_filters - bsz = audio.shape[0] with torch.no_grad(): - # Compute mel features - mel = melspec(audio.float()) - logmel = mel.transpose(-1, -2).clip_(min=1e-10).log10_() - mx = logmel.amax(dim=(-2, -1), keepdim=True) - logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1) - # remove last frame if odd - if logmel.shape[1] % 2 == 1: - logmel = logmel[:, :-1] - - # stacking and skipping by 2 - audio = logmel.reshape(bsz, -1, 2 * logmel.shape[-1]) - + audio = module(audio) return audio def _get_num_audio_features(self, audio_lengths: Sequence[int]) -> Sequence[int]: @@ -128,6 +170,7 @@ def _get_num_audio_features(self, audio_lengths: Sequence[int]) -> Sequence[int] audio_lengths (`Sequence[int]`): Sequence of one or more raw audio lengths. """ + # TODO: make this torch based and exportable hop_length = self.melspec_kwargs["hop_length"] effective_window_size = self.projector_window_size // self.projector_downsample_rate diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index e11895191f95..4cee17eea16e 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -25,14 +25,47 @@ from ...feature_extraction_sequence_utils import SequenceFeatureExtractor from ...feature_extraction_utils import BatchFeature from ...utils import TensorType, logging +from ...utils.import_utils import requires_backends if is_torch_available(): import torch + from torch import nn logger = logging.get_logger(__name__) +class _WhisperFeatureExtractorModule(nn.Module): + def __init__(self, feature_extractor: "WhisperFeatureExtractor"): + super().__init__() + self.n_fft = feature_extractor.n_fft + self.hop_length = feature_extractor.hop_length + self.dither = feature_extractor.dither + self.register_buffer("window", torch.hann_window(self.n_fft)) + self.register_buffer("mel_filters", torch.from_numpy(feature_extractor.mel_filters).float()) + + def forward(self, waveform): + # 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) + + stft = torch.stft(waveform, self.n_fft, self.hop_length, window=self.window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + + mel_spec = self.mel_filters.T @ magnitudes + + 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] + log_spec = torch.maximum(log_spec, max_val - 8.0) + else: + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + return log_spec + + class WhisperFeatureExtractor(SequenceFeatureExtractor): r""" Constructs a Whisper feature extractor. @@ -105,17 +138,18 @@ def __init__( mel_scale="slaney", ) - def _np_extract_fbank_features(self, waveform_batch: np.ndarray, device: str) -> np.ndarray: + def to_exportable_module(self) -> "nn.Module": + """ + Returns an exportable version of the feature extractor, which can be used with `torch.export`. + """ + requires_backends(self, "torch") + return _WhisperFeatureExtractorModule(self) + + def _np_extract_fbank_features(self, waveform_batch: np.ndarray) -> np.ndarray: """ Compute the log-mel spectrogram of the provided audio, gives similar results to Whisper's original torch implementation with 1e-5 tolerance. """ - if device != "cpu": - raise ValueError( - f"Got device `{device}` for feature extraction, but feature extraction on CUDA accelerator " - "devices requires torch, which is not installed. Either set `device='cpu'`, or " - "install torch according to the official instructions: https://pytorch.org/get-started/locally/" - ) log_spec_batch = [] for waveform in waveform_batch: log_spec = spectrogram( @@ -141,27 +175,8 @@ def _torch_extract_fbank_features(self, waveform: np.ndarray, device: str = "cpu yielding results similar to cpu computing with 1e-5 tolerance. """ waveform = torch.from_numpy(waveform).to(device, torch.float32) - window = torch.hann_window(self.n_fft, device=device) - - # 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) - - stft = torch.stft(waveform, self.n_fft, self.hop_length, window=window, return_complex=True) - magnitudes = stft[..., :-1].abs() ** 2 - - mel_filters = torch.from_numpy(self.mel_filters).to(device, torch.float32) - mel_spec = mel_filters.T @ magnitudes - - 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] - log_spec = torch.maximum(log_spec, max_val - 8.0) - else: - log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) - log_spec = (log_spec + 4.0) / 4.0 + module = self.to_exportable_module().to(device) + log_spec = module(waveform) if device != "cpu": log_spec = log_spec.detach().cpu() return log_spec.numpy() @@ -312,10 +327,10 @@ def __call__( # make sure list is in array format input_features = padded_inputs.get("input_features").transpose(2, 0, 1) - extract_fbank_features = ( - self._torch_extract_fbank_features if is_torch_available() else self._np_extract_fbank_features - ) - input_features = extract_fbank_features(input_features[0], device) + if is_torch_available() and device != "cpu": + input_features = self._torch_extract_fbank_features(input_features[0], device=device) + else: + input_features = self._np_extract_fbank_features(input_features[0]) if isinstance(input_features[0], list): padded_inputs["input_features"] = [np.asarray(feature, dtype=np.float32) for feature in input_features] diff --git a/tests/test_executorch.py b/tests/test_executorch.py index 0e33253c08f1..57809e67cd25 100644 --- a/tests/test_executorch.py +++ b/tests/test_executorch.py @@ -23,8 +23,10 @@ TorchExportableModuleWithHybridCache, TorchExportableModuleWithStaticCache, ) +from transformers.models.granite_speech.feature_extraction_granite_speech import GraniteSpeechFeatureExtractor +from transformers.models.whisper.feature_extraction_whisper import WhisperFeatureExtractor from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_3 -from transformers.testing_utils import require_torch +from transformers.testing_utils import require_torch, require_torchaudio @require_torch @@ -127,3 +129,29 @@ def test_decoder_only_lm_export(self): inputs_embeds=self.inputs_embeds, cache_position=self.cache_position ) torch.testing.assert_close(eager_output_embeds, exported_output_embeds, atol=1e-4, rtol=1e-4) + + +@require_torchaudio +@require_torch +class FeatureExtractorExportTest(unittest.TestCase): + def setUp(self): + if not is_torch_greater_or_equal_than_2_3: + self.skipTest("torch >= 2.3 is required") + + def test_whisper_export(self): + feature_extractor = WhisperFeatureExtractor() + exportable_module = feature_extractor.to_exportable_module() + waveform = torch.randn(1, 16000, dtype=torch.float32) + exported_program = torch.export.export(exportable_module, args=(waveform,)) + self.assertIsNotNone(exported_program) + exported_output = exported_program.module()(waveform) + self.assertIsNotNone(exported_output) + + def test_granite_speech_export(self): + feature_extractor = GraniteSpeechFeatureExtractor() + exportable_module = feature_extractor.to_exportable_module() + waveform = torch.randn(1, 16000, dtype=torch.float32) + exported_program = torch.export.export(exportable_module, args=(waveform,)) + self.assertIsNotNone(exported_program) + exported_output = exported_program.module()(waveform) + self.assertIsNotNone(exported_output) From 8cc1bebea4f9cc47980ebc63f59626cd29a7576c Mon Sep 17 00:00:00 2001 From: Pk Patel <46714886+The5cheduler@users.noreply.github.com> Date: Tue, 23 Sep 2025 21:50:14 -0400 Subject: [PATCH 2/5] fixed the failing testcase for pytorch --- .../feature_extraction_granite_speech.py | 92 +++++++++---------- 1 file changed, 46 insertions(+), 46 deletions(-) diff --git a/src/transformers/models/granite_speech/feature_extraction_granite_speech.py b/src/transformers/models/granite_speech/feature_extraction_granite_speech.py index b01f053ce569..92f4c20cafb9 100644 --- a/src/transformers/models/granite_speech/feature_extraction_granite_speech.py +++ b/src/transformers/models/granite_speech/feature_extraction_granite_speech.py @@ -36,53 +36,53 @@ import torchaudio -class _GraniteSpeechFeatureExtractorModule(nn.Module): - def __init__(self, feature_extractor: "GraniteSpeechFeatureExtractor"): - super().__init__() - self.melspec_kwargs = feature_extractor.melspec_kwargs - self.mel_filters = torchaudio.transforms.MelSpectrogram(**self.melspec_kwargs) - self.projector_window_size = feature_extractor.projector_window_size - self.projector_downsample_rate = feature_extractor.projector_downsample_rate - - def _get_num_audio_features(self, audio_lengths: "torch.Tensor") -> "torch.Tensor": - """ - Gets the (variable length) number of features (i.e., projector output) for the sequences - being considered. - - Args: - audio_lengths (`torch.Tensor`): - Sequence of one or more raw audio lengths. - """ - hop_length = self.melspec_kwargs["hop_length"] - effective_window_size = self.projector_window_size // self.projector_downsample_rate - - # mel sequence length computation - mel_length = audio_lengths // hop_length + 1 - # encoder frame takes two mel features - encoder_length = mel_length // 2 - nblocks = (encoder_length + self.projector_window_size - 1) // self.projector_window_size - # projector output length - projector_length = nblocks * effective_window_size - return projector_length + class _GraniteSpeechFeatureExtractorModule(nn.Module): + def __init__(self, feature_extractor: "GraniteSpeechFeatureExtractor"): + super().__init__() + self.melspec_kwargs = feature_extractor.melspec_kwargs + self.mel_filters = torchaudio.transforms.MelSpectrogram(**self.melspec_kwargs) + self.projector_window_size = feature_extractor.projector_window_size + self.projector_downsample_rate = feature_extractor.projector_downsample_rate + + def _get_num_audio_features(self, audio_lengths: "torch.Tensor") -> "torch.Tensor": + """ + Gets the (variable length) number of features (i.e., projector output) for the sequences + being considered. + + Args: + audio_lengths (`torch.Tensor`): + Sequence of one or more raw audio lengths. + """ + hop_length = self.melspec_kwargs["hop_length"] + effective_window_size = self.projector_window_size // self.projector_downsample_rate - def forward(self, audio: "torch.Tensor"): - """ - Compute the Mel features to be passed to the conformer encoder. - """ - bsz = audio.shape[0] - # Compute mel features - mel = self.mel_filters(audio.float()) - logmel = mel.transpose(-1, -2).clip_(min=1e-10).log10_() - mx = logmel.amax(dim=(-2, -1), keepdim=True) - logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1) - # remove last frame if odd - if logmel.shape[1] % 2 == 1: - logmel = logmel[:, :-1] - - # stacking and skipping by 2 - audio = logmel.reshape(bsz, -1, 2 * logmel.shape[-1]) - - return audio + # mel sequence length computation + mel_length = audio_lengths // hop_length + 1 + # encoder frame takes two mel features + encoder_length = mel_length // 2 + nblocks = (encoder_length + self.projector_window_size - 1) // self.projector_window_size + # projector output length + projector_length = nblocks * effective_window_size + return projector_length + + def forward(self, audio: "torch.Tensor"): + """ + Compute the Mel features to be passed to the conformer encoder. + """ + bsz = audio.shape[0] + # Compute mel features + mel = self.mel_filters(audio.float()) + logmel = mel.transpose(-1, -2).clip_(min=1e-10).log10_() + mx = logmel.amax(dim=(-2, -1), keepdim=True) + logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1) + # remove last frame if odd + if logmel.shape[1] % 2 == 1: + logmel = logmel[:, :-1] + + # stacking and skipping by 2 + audio = logmel.reshape(bsz, -1, 2 * logmel.shape[-1]) + + return audio class GraniteSpeechFeatureExtractor(FeatureExtractionMixin): From f8bdbde356bb63b729bf1619ae49106796ae1a30 Mon Sep 17 00:00:00 2001 From: Pk Patel <46714886+The5cheduler@users.noreply.github.com> Date: Tue, 23 Sep 2025 21:53:05 -0400 Subject: [PATCH 3/5] Fixing the main method feature_extraction_whisper.py --- .../whisper/feature_extraction_whisper.py | 60 +++++++++---------- 1 file changed, 30 insertions(+), 30 deletions(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 4cee17eea16e..7ee7507f2f5e 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -32,38 +32,38 @@ import torch from torch import nn -logger = logging.get_logger(__name__) - - -class _WhisperFeatureExtractorModule(nn.Module): - def __init__(self, feature_extractor: "WhisperFeatureExtractor"): - super().__init__() - self.n_fft = feature_extractor.n_fft - self.hop_length = feature_extractor.hop_length - self.dither = feature_extractor.dither - self.register_buffer("window", torch.hann_window(self.n_fft)) - self.register_buffer("mel_filters", torch.from_numpy(feature_extractor.mel_filters).float()) - - def forward(self, waveform): - # 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) - stft = torch.stft(waveform, self.n_fft, self.hop_length, window=self.window, return_complex=True) - magnitudes = stft[..., :-1].abs() ** 2 - - mel_spec = self.mel_filters.T @ magnitudes + class _WhisperFeatureExtractorModule(nn.Module): + def __init__(self, feature_extractor: "WhisperFeatureExtractor"): + super().__init__() + self.n_fft = feature_extractor.n_fft + self.hop_length = feature_extractor.hop_length + self.dither = feature_extractor.dither + self.register_buffer("window", torch.hann_window(self.n_fft)) + self.register_buffer("mel_filters", torch.from_numpy(feature_extractor.mel_filters).float()) + + def forward(self, waveform): + # 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) + + stft = torch.stft(waveform, self.n_fft, self.hop_length, window=self.window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + + mel_spec = self.mel_filters.T @ magnitudes + + 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] + log_spec = torch.maximum(log_spec, max_val - 8.0) + else: + log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) + log_spec = (log_spec + 4.0) / 4.0 + return log_spec - 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] - log_spec = torch.maximum(log_spec, max_val - 8.0) - else: - log_spec = torch.maximum(log_spec, log_spec.max() - 8.0) - log_spec = (log_spec + 4.0) / 4.0 - return log_spec +logger = logging.get_logger(__name__) class WhisperFeatureExtractor(SequenceFeatureExtractor): From 270c47dd63dcfe457affe5a05a330e9827d5c6fb Mon Sep 17 00:00:00 2001 From: Pk Patel <46714886+The5cheduler@users.noreply.github.com> Date: Tue, 23 Sep 2025 22:04:44 -0400 Subject: [PATCH 4/5] removed-additional space --- .../models/granite_speech/feature_extraction_granite_speech.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/transformers/models/granite_speech/feature_extraction_granite_speech.py b/src/transformers/models/granite_speech/feature_extraction_granite_speech.py index 92f4c20cafb9..c0d6ac775946 100644 --- a/src/transformers/models/granite_speech/feature_extraction_granite_speech.py +++ b/src/transformers/models/granite_speech/feature_extraction_granite_speech.py @@ -35,7 +35,6 @@ if is_torchaudio_available(): import torchaudio - class _GraniteSpeechFeatureExtractorModule(nn.Module): def __init__(self, feature_extractor: "GraniteSpeechFeatureExtractor"): super().__init__() From cef648e94900efa10e2a152ea32c6014450a85df Mon Sep 17 00:00:00 2001 From: Pk Patel <46714886+The5cheduler@users.noreply.github.com> Date: Tue, 23 Sep 2025 22:07:24 -0400 Subject: [PATCH 5/5] applied proper formatting --- src/transformers/models/whisper/feature_extraction_whisper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/transformers/models/whisper/feature_extraction_whisper.py b/src/transformers/models/whisper/feature_extraction_whisper.py index 7ee7507f2f5e..79b90de66a87 100644 --- a/src/transformers/models/whisper/feature_extraction_whisper.py +++ b/src/transformers/models/whisper/feature_extraction_whisper.py @@ -32,7 +32,6 @@ import torch from torch import nn - class _WhisperFeatureExtractorModule(nn.Module): def __init__(self, feature_extractor: "WhisperFeatureExtractor"): super().__init__() @@ -63,6 +62,7 @@ def forward(self, waveform): log_spec = (log_spec + 4.0) / 4.0 return log_spec + logger = logging.get_logger(__name__)