Feature extraction on GPU

[alexgorodetski1802]

Hello everyone,

I am using the TEDLIUM recipe as a base for my work with in-house data. My training dataset has about 15,000 hours of audio.

In the TEDLIUM recipe, feature extraction is done on the CPU, not the GPU. Because of this, feature extraction takes more than 12 hours.

I would like to ask if you have any implementation that supports GPU-based feature extraction.

Thank you,
AlexG

[lxp3]

Lhotse supports computing FBANK features directly on the GPU. Below is a minimal example:

from lhotse import Fbank, FbankConfig

fbank_config = FbankConfig(
    num_mel_bins=80,
    dither=0.0,
    device="cuda",
)
fbank = Fbank(fbank_config)

features = fbank.extract_batch(waveforms, sampling_rate=16000)

Here, waveforms is a torch.Tensor with shape [batch_size, max_sequence_length].
To construct this tensor, you should pad a list of 1-D waveform tensors using torch.nn.utils.rnn.pad_sequence:

# Input: waveforms is List[torch.Tensor], each with shape [sequence_length]
waveforms: torch.Tensor = torch.nn.utils.rnn.pad_sequence(
    waveforms,
    batch_first=True,
    padding_value=0.0,
)

Important: After padding, the original waveform lengths are lost. You should explicitly record the original number of samples for each utterance and compute the corresponding feature lengths based on those values.

You can compute feature lengths as follows:

from lhotse.utils import compute_num_frames_from_samples

def compute_feature_lens(
    waveform_lens: List[int],
    sampling_rate: int,
) -> List[int]:
    feature_lens = []
    for num_samples in waveform_lens:
        feature_len = compute_num_frames_from_samples(
            num_samples=num_samples,
            frame_shift=0.01,
            sampling_rate=sampling_rate,
        )
        feature_lens.append(feature_len)
    return feature_lens

Finally, if feature extraction on the GPU is still slow, the bottleneck is often resampling. Make sure that torchaudio.functional.resample is also executed on the GPU.