LibriSpeechDataLoader.py

"""
Custom collate function for batching variable-length speech data.

This function:
- Pads mel-spectrograms along the time dimension
- Pads transcript sequences for CTC training
- Returns original sequence lengths for proper loss computation
"""
import torch
from torch.nn.utils.rnn import pad_sequence
import torch.nn.functional as F

def collate_fn(batch):
    """
        Collate function for DataLoader handling variable-length audio
        and transcript sequences.

        Args:
            batch (list[tuple]):
                List of (mel_spectrogram, transcript_indices) tuples.

                - mel_spectrogram: Tensor of shape (1, n_mels, time)
                - transcript_indices: 1D LongTensor of variable length

        Returns:
            tuple:
                padded_specs (Tensor):
                    Padded mel-spectrograms of shape (B, 1, n_mels, max_time)

                padded_transcripts (Tensor):
                    Padded transcript indices of shape (B, max_target_length)

                spec_lengths (LongTensor):
                    Original time lengths of each spectrogram

                transcript_lengths (LongTensor):
                    Original transcript lengths
    """
    specs, transcript = zip(*batch)

    spec_lengths = [spec.shape[-1] for spec in specs]
    max_length = max(spec_lengths)
    padded_specs = [F.pad(spec, (0, max_length - spec.shape[-1])) for spec in specs]
    padded_specs = torch.stack(padded_specs)

    #specs = [spec.squeeze(0).transpose(0, 1) for spec in specs]

    #padded_specs = pad_sequence(specs, batch_first=True)

    #padded_specs = padded_specs.transpose(1, 2).unsqueeze(1)

    transcripts = [t.detach().clone() for t in transcript]

    padded_transcripts = pad_sequence(transcripts, batch_first=True, padding_value=0)


    spec_lengths = torch.tensor(spec_lengths, dtype=torch.long)
    transcript_lengths = torch.tensor([len(t) for t in transcripts], dtype=torch.long)

    return padded_specs, padded_transcripts, spec_lengths, transcript_lengths