VoiceRecognition.py

from torch import nn
import torch
from torch.utils.data import DataLoader
from LibriSpeechMelDataset import LibriSpeechMelDataset
from LibriSpeechDataLoader import collate_fn
import time
import os

class SpeechRecognitionModel(nn.Module):
    """
       CNN + BiLSTM based speech recognition model for mel-spectrogram inputs.

       Architecture:
       - 2D CNN feature extractor over (frequency, time)
       - Bidirectional LSTM over time dimension
       - Linear classifier for per-timestep character prediction (CTC-compatible)
    """
    def __init__(self, input_channels=1, n_mels=64, hidden_size=256, output_size=None, dropout=0.3):
        """
            Initialize the speech recognition model.

            Args:
                input_channels (int): Number of input channels (default 1 for mel spectrograms).
                n_mels (int): Number of mel frequency bins.
                hidden_size (int): Hidden size of the LSTM.
                output_size (int): Vocabulary size (number of output classes).
                dropout (float): Dropout probability applied in CNN layers.

            Raises:
                ValueError: If output_size is not provided.
        """
        super().__init__()
        if output_size is None:
            raise ValueError("Must provide output size \nThis should be the length of your vocab")

        self.cnn = nn.Sequential(
            nn.Conv2d(input_channels, 32,kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),

            nn.Conv2d(32, 64,kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=(1,2)),
            nn.Dropout2d(dropout)
        )

        conv_output_freq = n_mels
        rnn_input_size = conv_output_freq * 64

        self.rnn = nn.LSTM(input_size=rnn_input_size, hidden_size=hidden_size, num_layers=2, batch_first=True, bidirectional=True)
        self.classifier = nn.Linear(hidden_size * 2, output_size)
    def forward(self, x, input_length=None):
        """
               Forward pass through the model.

               Args:
                   x (Tensor): Input spectrograms of shape (B, C, F, T).
                   input_length (Tensor, optional): Original time lengths for each sample.

               Returns:
                   Tensor or (Tensor, Tensor):
                       - Output logits of shape (B, T', vocab_size)
                       - Reduced input lengths after CNN downsampling (if provided)
        """
        if input_length is not None:
            input_length = input_length.clone()
        x = self.cnn(x)
        if input_length is not None:
            for layer in self.cnn:
                if isinstance(layer, torch.nn.Conv2d):
                    kernel_size = layer.kernel_size
                    stride = layer.stride
                    padding = layer.padding
                    dilation = layer.dilation

                    time_kernel = kernel_size[1] if isinstance(kernel_size, tuple) else kernel_size
                    time_stride = stride[1] if isinstance(stride, tuple) else stride
                    time_padding = padding[1] if isinstance(padding, tuple) else padding
                    time_dilation = dilation[1] if isinstance(dilation, tuple) else dilation

                    input_length = ((input_length + 2 * time_padding - time_dilation * (time_kernel - 1)-1) // time_stride) + 1

                elif isinstance(layer, torch.nn.MaxPool2d):

                    if isinstance(layer.stride, tuple) or isinstance(layer.stride, list):
                        time_stride = layer.stride[1]
                    elif layer.stride is None:
                        time_stride = layer.kernel_size if isinstance(layer.kernel_size, int) else layer.kernel_size[1]
                    else:
                        time_stride = layer.stride

                    input_length = input_length // time_stride

        batch_size, channels, freq, time = x.size()
        x = x.permute(0,3,1,2)
        x = x.contiguous().view(batch_size, time, channels * freq)

        x, _ = self.rnn(x)
        x = self.classifier(x)

        if input_length is not None:
            return x, input_length
        else:
            return x

#Test Utilities
def test_model_output_shape():
    """
       Verify that model output dimensions match expectations.
    """
    model = SpeechRecognitionModel(output_size=30).eval()
    dummy_input = torch.randn(4, 1, 64, 100)
    input_lengths = torch.tensor([100, 95, 90, 85])
    output, reduced_lengths = model(dummy_input, input_lengths)
    assert output.shape[0] == 4 and output.shape[2] == 30, "Output shape mismatch"
    print("test_model_output_shape: PASSED")

def test_overfit_single_batch():
    """
        Test whether the model can overfit a single batch.
        Useful for debugging training and loss correctness.
    """
    dataset = LibriSpeechMelDataset("./data", url="test-clean", download=True)
    loader = DataLoader(dataset, batch_size=2, collate_fn=collate_fn)
    batch = next(iter(loader))
    model = SpeechRecognitionModel(output_size=len(dataset.vocab)).to("cuda")
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    ctc_loss = nn.CTCLoss(blank=0, zero_infinity=True)

    model.train()
    specs, transcripts, input_lengths, target_lengths = [x.to("cuda") for x in batch]
    for i in range(50):
        optimizer.zero_grad()
        out, reduced_lens = model(specs, input_lengths)
        log_probs = out.log_softmax(dim=-1).permute(1, 0, 2)
        targets = torch.cat([transcripts[i, :target_lengths[i]] for i in range(len(transcripts))])
        loss = ctc_loss(log_probs, targets, reduced_lens, target_lengths)
        loss.backward()
        optimizer.step()
        print(f"[Step {i}] Loss: {loss.item():.4f}")
        if loss.item() < 0.1:
            print("test_overfit_single_batch: PASSED")
            break

def test_gradients():
    """
        Ensure gradients propagate through all trainable parameters.
    """
    model = SpeechRecognitionModel(output_size=30).to("cuda")
    dummy_input = torch.randn(2, 1, 64, 100).to("cuda")
    input_lengths = torch.tensor([100, 95], dtype=torch.long).to("cuda")
    targets = torch.randint(1, 30, (20,), dtype=torch.long).to("cuda")
    target_lengths = torch.tensor([10, 10], dtype=torch.long).to("cuda")
    out, red_len = model(dummy_input, input_lengths)
    log_probs = out.log_softmax(dim=-1).permute(1, 0, 2)
    loss = nn.CTCLoss(blank=0, zero_infinity=True)(log_probs, targets, red_len, target_lengths)
    loss.backward()
    for name, param in model.named_parameters():
        if param.grad is not None:
            print(f"{name}: gradient exists")
    print("test_gradients: PASSED")

def test_inference_speed():
    """
        Measure average inference time per batch.
    """
    model = SpeechRecognitionModel(output_size=30).to("cuda")
    model.eval()
    dummy_input = torch.randn(16, 1, 64, 100).to("cuda")
    input_lengths = torch.tensor([100]*16).to("cuda")
    start = time.time()
    with torch.no_grad():
        for _ in range(50):
            model(dummy_input, input_lengths)
    print(f"Inference speed: {(time.time() - start)/50:.4f} s per batch")
    print("test_inference_speed: PASSED")

####TRAINING AND TESTING
def train(model, dataloader, optimizer, device, epochs, start_epoch=0):
    """
        Train the speech recognition model using CTC loss.

        Args:
            model (nn.Module): Model to train.
            dataloader (DataLoader): Training data loader.
            optimizer (Optimizer): Optimizer instance.
            device (torch.device): Training device.
            epochs (int): Total number of epochs.
            start_epoch (int): Epoch to resume from.
    """
    avg_losses = []
    ctc_loss = nn.CTCLoss(blank=1, zero_infinity=True)
    model.to(device)
    model.train()

    for epoch in range(start_epoch, epochs):
        total_loss = 0

        for batch_index, (specs, transcripts, input_lengths, target_lengths) in enumerate(dataloader):
            specs = specs.to(device)
            transcripts = transcripts.to(device)
            target_lengths = target_lengths.to(device)

            optimizer.zero_grad()

            output, reduced_input_lengths = model(specs, input_lengths.to(device))
            log_probs = output.log_softmax(dim=-1)
            log_probs = log_probs.permute(1, 0, 2)

            targets = []

            for i, length in enumerate(target_lengths):
                targets.append(transcripts[i, :length])
            targets = torch.cat(targets)


            print("Input lengths:", input_lengths.cpu().tolist())
            print("Reduced lengths:", reduced_input_lengths.cpu().tolist())
            print("Target lengths:", target_lengths.cpu().tolist())
            loss = ctc_loss(log_probs, targets, reduced_input_lengths, target_lengths)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

            if batch_index % 10 == 0:
                print(f"Epoch [{epoch + 1}/{epochs}], Batch [{batch_index}], Loss: {loss.item():.4f}")
        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch + 1} average loss: {avg_loss:.4f}")
        avg_losses.append(avg_loss)
    print("Average Losses:", avg_losses)
def test(model, dataloader, device):
    """
        Evaluate the model on a test dataset.

        Args:
            model (nn.Module): Trained model.
            dataloader (DataLoader): Test data loader.
            device (torch.device): Evaluation device.
    """
    AverageLosses = []
    model.eval()
    ctc_loss = nn.CTCLoss(blank=1, zero_infinity=True)
    total_loss = 0

    with torch.no_grad():
        for batch_index, (specs, transcripts, input_lengths, target_lengths) in enumerate(dataloader):
            specs = specs.to(device)
            transcripts = transcripts.to(device)
            target_lengths = target_lengths.to(device)


            output, reduced_input_lengths = model(specs, input_lengths.to(device))
            log_probs = output.log_softmax(dim=-1)
            log_probs = log_probs.permute(1, 0, 2)

            targets = []

            for i, length in enumerate(target_lengths):
                targets.append(transcripts[i, :length])
            targets = torch.cat(targets)


            print("Input lengths:", input_lengths.cpu().tolist())
            print("Reduced lengths:", reduced_input_lengths.cpu().tolist())
            print("Target lengths:", target_lengths.cpu().tolist())
            loss = ctc_loss(log_probs, targets, reduced_input_lengths, target_lengths)

            total_loss += loss.item()

    avg_loss = total_loss / len(dataloader)
    print(f"Test average loss: {avg_loss:.4f}")
    AverageLosses.append(avg_loss)
    print("Average Losses by Epoch:", AverageLosses)


if __name__ == "__main__":
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    # Optional tests
    test_model_output_shape()
    test_gradients()
    test_inference_speed()
    test_overfit_single_batch()

    # Dataset and Dataloader
    Train_dataset = LibriSpeechMelDataset(
        root="./data",
        url="train-other-500",
        download=False
    )
    print("Dataset Loaded:", len(Train_dataset))

    dataloader = DataLoader(
        Train_dataset,
        batch_size=64,
        shuffle=True,
        collate_fn=collate_fn,
        num_workers=8,
        pin_memory=True
    )
    print()
    print("DataLoader initialized")
    print("---------------------------------------------------")

    # Model and optimizer
    model = SpeechRecognitionModel(
        output_size=len(Train_dataset.vocab)
    ).to(device)

    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    start_epoch = 0

    # ---- CHECKPOINT SELECTION ----
    checkpoints = sorted(f for f in os.listdir(".") if f.endswith(".pth"))

    if checkpoints:
        print("\nAvailable checkpoints:")
        for i, ckpt in enumerate(checkpoints):
            print(f"[{i}] {ckpt}")

        print("[Enter] Start training from scratch")
        choice = input("Select a checkpoint to load: ").strip()

        if choice.isdigit() and int(choice) < len(checkpoints):
            ckpt_path = checkpoints[int(choice)]
            print(f"\nLoading checkpoint: {ckpt_path}")

            checkpoint = torch.load(ckpt_path, map_location=device)
            model.load_state_dict(checkpoint["model_state_dict"])
            optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
            start_epoch = checkpoint.get("epoch", 0)
        else:
            print("\nStarting training from scratch.")
    else:
        print("\nNo checkpoints found. Starting training from scratch.")

    # ---- EPOCH SELECTION ----
    default_epochs = 230
    epoch_input = input(
        f"\nEnter total number of epochs to train "
        f"(default: {default_epochs}): "
    ).strip()

    if epoch_input.isdigit():
        epochs = int(epoch_input)
        if epochs <= start_epoch:
            raise ValueError(
                f"Total epochs ({epochs}) must be greater than "
                f"start_epoch ({start_epoch})."
            )
    else:
        epochs = default_epochs

    print(f"\nTraining from epoch {start_epoch} to {epochs}")

    # ---- TRAIN ----
    train(
        model,
        dataloader,
        optimizer=optimizer,
        device=device,
        epochs=epochs,
        start_epoch=start_epoch
    )

    # ---- SAVE FINAL CHECKPOINT ----
    save_name = f"checkpoint_epoch_{epochs}.pth"
    torch.save(
        {
            "model_state_dict": model.state_dict(),
            "optimizer_state_dict": optimizer.state_dict(),
            "epoch": epochs,
        },
        save_name
    )

    print(f"\nCheckpoint saved as: {save_name}")