{"nbformat":4,"nbformat_minor":0,"metadata":{"colab":{"provenance":[]},"kernelspec":{"name":"python3","display_name":"Python 3"},"language_info":{"name":"python"}},"cells":[{"cell_type":"markdown","source":["# Recitation 0: Paper2Code\n","\n","Prepared by: Massa Baali (mbaali@andrew.cmu.edu)\n","\n"],"metadata":{"id":"GAR2NBbvyhSN"}},{"cell_type":"markdown","source":["### Dataset"],"metadata":{"id":"r1x1xUC4PBhb"}},{"cell_type":"code","execution_count":null,"metadata":{"id":"E9phVg224JH5"},"outputs":[],"source":["# dataloader.py\n","import os\n","import torchaudio\n","from torch.utils.data import Dataset, DataLoader, random_split\n","import pandas as pd\n","\n","class VoxCelebDataset(Dataset):\n","    def __init__(self, train_csv_path, transform=None):\n","        \"\"\"\n","        Args:\n","            train_csv_path (str): Path to the dataset csv file.\n","            transform (callable, optional): Optional transforms to apply to the audio data.\n","        \"\"\"\n","        df = pd.read_csv(train_csv_path)\n","        self.labels = df[\"utt_spk_int_labels\"].values\n","        self.paths = df[\"utt_paths\"].values\n","\n","    def __len__(self):\n","        return len(self.paths)\n","\n","    def __getitem__(self, index):\n","        waveform = torchaudio.load(self.paths[index])\n","\n","        waveform_length = waveform.shape[-1]\n","\n","        sample = {\n","        'waveform':  waveform,\n","        'path': self.paths[index],\n","        'mapped_id': self.labels[index],\n","        'lens': waveform_length\n","        }\n","        return sample\n","\n","# Path to your VoxCeleb dataset directory\n","data_dir = \"/path/to/voxceleb.csv\"\n","\n","# Create the dataset\n","voxceleb_dataset = VoxCelebDataset(data_dir)\n","\n","# Split the dataset into train, validation, and test sets\n","train_size = int(0.8 * len(voxceleb_dataset))\n","val_size = int(0.1 * len(voxceleb_dataset))\n","test_size = len(voxceleb_dataset) - train_size - val_size\n","\n","train_dataset, val_dataset, test_dataset = random_split(\n","    voxceleb_dataset, [train_size, val_size, test_size]\n",")\n","\n","# Create DataLoaders for each split\n","batch_size = 16\n","voxceleb_dataloader_train = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n","voxceleb_dataloader_val = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)\n","voxceleb_dataloader_test = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n"]},{"cell_type":"markdown","source":["### Model\n"],"metadata":{"id":"aToAR43kQQ3M"}},{"cell_type":"code","source":["# model.py\n","import torch\n","import torch.nn as nn\n","import torch.nn.functional as F\n","\n","class MultiFrameAttention(nn.Module):\n","    def __init__(self, embed_dim, num_heads):\n","        super(MultiFrameAttention, self).__init__()\n","        self.attention = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)\n","        self.linear = nn.Linear(embed_dim, embed_dim)\n","\n","    def forward(self, x):\n","        # Input shape: (batch_size, sequence_length, embed_dim)\n","        attn_output, _ = self.attention(x, x, x)\n","        return self.linear(attn_output)\n","\n","class ConformerBlock(nn.Module):\n","    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):\n","        super(ConformerBlock, self).__init__()\n","        self.mfa = MultiFrameAttention(embed_dim, num_heads)\n","        self.ff1 = nn.Linear(embed_dim, ff_dim)\n","        self.ff2 = nn.Linear(ff_dim, embed_dim)\n","        self.norm1 = nn.LayerNorm(embed_dim)\n","        self.norm2 = nn.LayerNorm(embed_dim)\n","        self.dropout = nn.Dropout(dropout)\n","        self.conv = nn.Conv1d(embed_dim, embed_dim, kernel_size=31, padding=15, groups=embed_dim)  # Depthwise conv\n","\n","    def forward(self, x):\n","        # Multi-Frame Attention\n","        residual = x\n","        x = self.mfa(x)\n","        x = self.norm1(residual + self.dropout(x))\n","\n","        # Convolutional Module\n","        residual = x\n","        x = x.transpose(1, 2)  # (batch_size, embed_dim, sequence_length)\n","        x = F.gelu(self.conv(x))\n","        x = x.transpose(1, 2)  # Back to (batch_size, sequence_length, embed_dim)\n","        x = self.norm2(residual + self.dropout(x))\n","\n","        # Feedforward Network\n","        residual = x\n","        x = F.gelu(self.ff1(x))\n","        x = self.ff2(x)\n","        return residual + self.dropout(x)\n","\n","class ConformerMFA(nn.Module):\n","    def __init__(self, num_blocks, embed_dim, num_heads, ff_dim, num_classes, dropout=0.1):\n","        super(ConformerMFA, self).__init__()\n","        self.blocks = nn.ModuleList([\n","            ConformerBlock(embed_dim, num_heads, ff_dim, dropout) for _ in range(num_blocks)\n","        ])\n","        self.fc_out = nn.Linear(embed_dim, num_classes)\n","\n","    def forward(self, x):\n","        # Input shape: (batch_size, sequence_length, embed_dim)\n","        for block in self.blocks:\n","            x = block(x)\n","        # Average pooling for sequence aggregation\n","        x = x.mean(dim=1)  # (batch_size, embed_dim)\n","        return self.fc_out(x)\n","\n","# Example Usage\n","if __name__ == \"__main__\":\n","    # Define model\n","    model = ConformerMFA(\n","        num_blocks=4,\n","        embed_dim=256,\n","        num_heads=8,\n","        ff_dim=1024,\n","        num_classes=10,\n","        dropout=0.1\n","    )\n","\n","    # Dummy input (batch_size=16, sequence_length=100, embed_dim=256)\n","    dummy_input = torch.randn(16, 100, 256)\n","    output = model(dummy_input)\n","\n","    print(\"Output shape:\", output.shape)  # Should be (16, num_classes)\n"],"metadata":{"colab":{"base_uri":"https://localhost:8080/"},"id":"hc8MSr8CGLhF","outputId":"b75814c2-36cd-4809-c4ea-fc295ce5b2c6"},"execution_count":null,"outputs":[{"output_type":"stream","name":"stdout","text":["Output shape: torch.Size([16, 10])\n"]}]},{"cell_type":"markdown","source":["### Train\n"],"metadata":{"id":"EcYo7iT8PF-P"}},{"cell_type":"code","source":["# train.py\n","import torch\n","import torch.nn as nn\n","import torch.optim as optim\n","from torch.utils.data import DataLoader\n","from tqdm import tqdm\n","\n","def train_one_epoch(model, dataloader, criterion, optimizer, device):\n","    model.train()\n","    epoch_loss, correct, total = 0, 0, 0\n","\n","    for inputs, labels in tqdm(dataloader, desc=\"Training\", leave=False):\n","        inputs, labels = inputs.to(device), labels.to(device)\n","\n","        # Forward pass\n","        outputs = model(inputs)\n","        loss = criterion(outputs, labels)\n","\n","        # Backward pass and optimization\n","        optimizer.zero_grad()\n","        loss.backward()\n","        optimizer.step()\n","\n","        # Metrics\n","        epoch_loss += loss.item() * inputs.size(0)\n","        _, preds = torch.max(outputs, 1)\n","        correct += (preds == labels).sum().item()\n","        total += labels.size(0)\n","\n","    avg_loss = epoch_loss / total\n","    accuracy = correct / total\n","    return avg_loss, accuracy\n","\n","def validate_one_epoch(model, dataloader, criterion, device):\n","    model.eval()\n","    epoch_loss, correct, total = 0, 0, 0\n","\n","    with torch.no_grad():\n","        for inputs, labels in tqdm(dataloader, desc=\"Validation\", leave=False):\n","            inputs, labels = inputs.to(device), labels.to(device)\n","\n","            # Forward pass\n","            outputs = model(inputs)\n","            loss = criterion(outputs, labels)\n","\n","            # Metrics\n","            epoch_loss += loss.item() * inputs.size(0)\n","            _, preds = torch.max(outputs, 1)\n","            correct += (preds == labels).sum().item()\n","            total += labels.size(0)\n","\n","    avg_loss = epoch_loss / total\n","    accuracy = correct / total\n","    return avg_loss, accuracy\n","\n","def save_checkpoint(model, optimizer, epoch, file_path):\n","    torch.save({\n","        'model_state_dict': model.state_dict(),\n","        'optimizer_state_dict': optimizer.state_dict(),\n","        'epoch': epoch,\n","    }, file_path)\n","\n","def train_model(model, train_loader, val_loader, criterion, optimizer, num_epochs, device, checkpoint_path):\n","    for epoch in range(1, num_epochs + 1):\n","        print(f\"\\nEpoch {epoch}/{num_epochs}\")\n","\n","        # Training phase\n","        train_loss, train_accuracy = train_one_epoch(model, train_loader, criterion, optimizer, device)\n","        print(f\"Training Loss: {train_loss:.4f}, Accuracy: {train_accuracy:.4f}\")\n","\n","        # Validation phase\n","        val_loss, val_accuracy = validate_one_epoch(model, val_loader, criterion, device)\n","        print(f\"Validation Loss: {val_loss:.4f}, Accuracy: {val_accuracy:.4f}\")\n","\n","        # U can add here a criteria to check the accuracy and the loss of the validation set..\n","\n","        # Save checkpoint\n","        save_checkpoint(model, optimizer, epoch, f\"{checkpoint_path}_epoch_{epoch}.pth\")\n","\n","# Example Usage\n","if __name__ == \"__main__\":\n","    criterion = nn.CrossEntropyLoss()\n","    optimizer = optim.Adam(model.parameters(), lr=0.001)\n","\n","    train_model(model, voxceleb_dataloader_train, voxceleb_dataloader_val, criterion, optimizer, num_epochs=5, device=device, checkpoint_path=\"model_checkpoint\")\n"],"metadata":{"id":"rJhnDAwiGSj4"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["### Eval\n"],"metadata":{"id":"T6pDEnYaPJH2"}},{"cell_type":"code","source":["#eval.py\n","import torch\n","from sklearn.metrics import accuracy_score\n","\n","def evaluate_model(model, test_loader, device):\n","    \"\"\"\n","    Evaluate the model on the test dataset.\n","\n","    Args:\n","        model: Trained MFA Conformer model.\n","        test_loader: Voxceleb DataLoader for the test dataset.\n","        device: Device to run the evaluation (CPU or GPU).\n","\n","    Returns:\n","        float: Accuracy score on the test set.\n","    \"\"\"\n","    model.eval()  # Set the model to evaluation mode\n","    all_preds = []\n","    all_labels = []\n","\n","    with torch.no_grad():\n","        for inputs, labels in test_loader:\n","            inputs, labels = inputs.to(device), labels.to(device)\n","\n","            # Forward pass\n","            outputs = model(inputs)\n","            _, preds = torch.max(outputs, 1)  # Get class predictions\n","\n","            all_preds.extend(preds.cpu().numpy())\n","            all_labels.extend(labels.cpu().numpy())\n","\n","    # Compute the accuracy\n","    accuracy = accuracy_score(all_labels, all_preds)\n","    print(f\"Test Accuracy: {accuracy:.4f}\")\n","    return accuracy\n","\n","# Example Usage\n","if __name__ == \"__main__\":\n","\n","    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","    model = torch.load(\"model_checkpoint_best.pth\")  # Load the best model checkpoint\n","    model = model['model_state_dict']  # Adjust if checkpoint includes state_dict\n","    model.to(device)\n","\n","    # Evaluate the model\n","    evaluate_model(model, voxceleb_dataloader_test, device)\n"],"metadata":{"id":"FhhZyKyHJ9l2"},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":["# References and further reading\n","\n","1.   https://arxiv.org/abs/2203.15249\n","2.   https://youtu.be/fdJxIqVBImU?list=PLp-0K3kfddPzbe1JqsQ7nmzZ38joYelj6\n"],"metadata":{"id":"QlvVfdOnRWnn"}},{"cell_type":"markdown","source":[],"metadata":{"id":"KXhP24gMKSEU"}}]}