feat: make audio feature extractors torch.export-able

src/transformers/models/granite_speech/feature_extraction_granite_speech.py

@@ -30,10 +30,59 @@
 
 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 +108,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 +151,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 +169,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
 

src/transformers/models/whisper/feature_extraction_whisper.py

@@ -25,10 +25,43 @@
 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
+
+    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
+
 
 logger = logging.get_logger(__name__)
 
@@ -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]

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)