{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "C_jdZ5vHJ4A9" }, "source": [ "# Task description\n", "- Classify the speakers of given features.\n", "- Main goal: Learn how to use transformer.\n", "- Baselines:\n", " - Easy: Run sample code and know how to use transformer.\n", " - Medium: Know how to adjust parameters of transformer.\n", " - Strong: Construct [conformer](https://arxiv.org/abs/2005.08100) which is a variety of transformer. \n", " - Boss: Implement [Self-Attention Pooling](https://arxiv.org/pdf/2008.01077v1.pdf) & [Additive Margin Softmax](https://arxiv.org/pdf/1801.05599.pdf) to further boost the performance.\n", "\n", "- Other links\n", " - Competiton: [link](https://www.kaggle.com/t/49ea0c385a974db5919ec67299ba2e6b)\n", " - Slide: [link](https://docs.google.com/presentation/d/1LDAW0GGrC9B6D7dlNdYzQL6D60-iKgFr/edit?usp=sharing&ouid=104280564485377739218&rtpof=true&sd=true)\n", " - Data: [link](https://github.com/googly-mingto/ML2023HW4/releases)\n", "\n", "# Download dataset\n", "- Data is [here](https://drive.google.com/drive/folders/1vI1kuLB-q1VilIftiwnPOCAeOOFfBZge?usp=sharing)" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "gtKxUzSgXKj3", "outputId": "3f59402c-95a7-4fbd-a39c-57606590a89c" }, "outputs": [], "source": [ "#!wget https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partaa\n", "#!wget https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partab\n", "#!wget https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partac\n", "#!wget https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partad\n", "\n", "#!cat Dataset.tar.gz.part* > Dataset.tar.gz\n", "#!rm Dataset.tar.gz.partaa\n", "#!rm Dataset.tar.gz.partab\n", "#!rm Dataset.tar.gz.partac\n", "#!rm Dataset.tar.gz.partad\n", "# unzip the file\n", "#!tar zxf Dataset.tar.gz\n", "#!rm Dataset.tar.gz" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "U6Y1cfpDfpON", "outputId": "6ba26637-5c7b-48a9-be0b-1f10ba76590a" }, "outputs": [], "source": [ "#!tar zxf Dataset.tar.gz" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "id": "E6burzCXIyuA" }, "outputs": [], "source": [ "import numpy as np\n", "import torch\n", "import random\n", "from common_function import train_start_log, train_log\n", "\n", "def set_seed(seed):\n", " np.random.seed(seed)\n", " random.seed(seed)\n", " torch.manual_seed(seed)\n", " if torch.cuda.is_available():\n", " torch.cuda.manual_seed(seed)\n", " torch.cuda.manual_seed_all(seed)\n", " torch.mps.manual_seed(seed)\n", " torch.mps.seed(seed)\n", " torch.backends.cudnn.benchmark = False\n", " torch.backends.cudnn.deterministic = True\n", "\n", "set_seed(8787)" ] }, { "cell_type": "markdown", "metadata": { "id": "k7dVbxW2LASN" }, "source": [ "# Data\n", "\n", "## Dataset\n", "- Original dataset is [Voxceleb2](https://www.robots.ox.ac.uk/~vgg/data/voxceleb/vox2.html).\n", "- The [license](https://creativecommons.org/licenses/by/4.0/) and [complete version](https://www.robots.ox.ac.uk/~vgg/data/voxceleb/files/license.txt) of Voxceleb2.\n", "- We randomly select 600 speakers from Voxceleb2.\n", "- Then preprocess the raw waveforms into mel-spectrograms.\n", "\n", "- Args:\n", " - data_dir: The path to the data directory.\n", " - metadata_path: The path to the metadata.\n", " - segment_len: The length of audio segment for training. \n", "- The architecture of data directory \\\\\n", " - data directory \\\\\n", " |---- metadata.json \\\\\n", " |---- testdata.json \\\\\n", " |---- mapping.json \\\\\n", " |---- uttr-{random string}.pt \\\\\n", "\n", "- The information in metadata\n", " - \"n_mels\": The dimention of mel-spectrogram.\n", " - \"speakers\": A dictionary. \n", " - Key: speaker ids.\n", " - value: \"feature_path\" and \"mel_len\"\n", "\n", "\n", "For efficiency, we segment the mel-spectrograms into segments in the traing step." ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "id": "KpuGxl4CI2pr" }, "outputs": [], "source": [ "import os\n", "import json\n", "import torch\n", "import random\n", "from pathlib import Path\n", "from torch.utils.data import Dataset\n", "from torch.nn.utils.rnn import pad_sequence\n", "from common_function import myDataset" ] }, { "cell_type": "markdown", "metadata": { "id": "668hverTMlGN" }, "source": [ "## Dataloader\n", "- Split dataset into training dataset(90%) and validation dataset(10%).\n", "- Create dataloader to iterate the data." ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "id": "B7c2gZYoJDRS" }, "outputs": [], "source": [ "import torch\n", "from torch.utils.data import DataLoader, random_split\n", "from torch.nn.utils.rnn import pad_sequence\n", "\n", "\n", "from common_function import collate_batch\n", "\n", "\n", "def get_dataloader(data_dir, batch_size, n_workers):\n", " \"\"\"Generate dataloader\"\"\"\n", " dataset = myDataset(data_dir)\n", " speaker_num = dataset.get_speaker_number()\n", " # Split dataset into training dataset and validation dataset\n", " trainlen = int(0.9 * len(dataset))\n", " lengths = [trainlen, len(dataset) - trainlen]\n", " trainset, validset = random_split(dataset, lengths)\n", "\n", " train_loader = DataLoader(\n", " trainset,\n", " batch_size=batch_size,\n", " shuffle=True,\n", " drop_last=True,\n", " num_workers=n_workers,\n", " pin_memory=True,\n", " collate_fn=collate_batch,\n", " )\n", " valid_loader = DataLoader(\n", " validset,\n", " batch_size=batch_size,\n", " num_workers=n_workers,\n", " drop_last=True,\n", " pin_memory=True,\n", " collate_fn=collate_batch,\n", " )\n", "\n", " return train_loader, valid_loader, speaker_num" ] }, { "cell_type": "markdown", "metadata": { "id": "5FOSZYxrMqhc" }, "source": [ "# Model\n", "- TransformerEncoderLayer:\n", " - Base transformer encoder layer in [Attention Is All You Need](https://arxiv.org/abs/1706.03762)\n", " - Parameters:\n", " - d_model: the number of expected features of the input (required).\n", "\n", " - nhead: the number of heads of the multiheadattention models (required).\n", "\n", " - dim_feedforward: the dimension of the feedforward network model (default=2048).\n", "\n", " - dropout: the dropout value (default=0.1).\n", "\n", " - activation: the activation function of intermediate layer, relu or gelu (default=relu).\n", "\n", "- TransformerEncoder:\n", " - TransformerEncoder is a stack of N transformer encoder layers\n", " - Parameters:\n", " - encoder_layer: an instance of the TransformerEncoderLayer() class (required).\n", "\n", " - num_layers: the number of sub-encoder-layers in the encoder (required).\n", "\n", " - norm: the layer normalization component (optional)." ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "id": "iXZ5B0EKJGs8" }, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "import torchaudio\n", "\n", "class Classifier(nn.Module):\n", " def __init__(self, d_model=160, n_spks=600, dropout=0.6):\n", " super().__init__()\n", " # Project the dimension of features from that of input into d_model.\n", " self.prenet = nn.Linear(40, d_model)\n", " # TODO:\n", " # Change Transformer to Conformer.\n", " # https://arxiv.org/abs/2005.08100\n", " #self.encoder_layer = nn.TransformerEncoderLayer(\n", " # d_model=d_model, dim_feedforward=256, nhead=8\n", " #)\n", " #self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=2)\n", " self.encoder = torchaudio.models.Conformer(d_model, 8, 256, 2, 31, dropout=dropout) # From torchaudio\n", " \n", " # Project the the dimension of features from d_model into speaker nums.\n", " self.pred_layer = nn.Linear(d_model, n_spks)\n", " \n", " # softmax\n", " self.softmax = nn.functional.softmax\n", " \n", " self.W = nn.Linear(d_model, 1)\n", "\n", " def forward(self, mels):\n", " \"\"\"\n", " args:\n", " mels: (batch size, length, 40)\n", " return:\n", " out: (batch size, n_spks)\n", " \"\"\"\n", " # out: (batch size, length, d_model)\n", " out = self.prenet(mels)\n", " # out: (length, batch size, d_model)\n", " #out = out.permute(1, 0, 2)\n", " # The encoder layer expect features in the shape of (length, batch size, d_model).\n", " device = \"cuda\" if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'\n", " #device = \"cpu\"\n", " #lengths = torch.randint(out.size(1), (int(out.size(0)),)).to(device)\n", " #lengths[torch.argmax(lengths)] = out.size(1)\n", " lengths = torch.full((int(out.size(0)),), out.size(1)).to(device)\n", " # out: (batch size, length, d_model)\n", " out, _ = self.encoder(out, lengths)\n", " # mean pooling\n", " stats = out.mean(dim=1)\n", " # self attention pooling \n", " # reference from https://gist.github.com/pohanchi/c77f6dbfbcbc21c5215acde4f62e4362\n", " # input batch_rep : size (N, T, H), N: batch size, T: sequence length, H: Hidden dimension\n", " # output utter_rep: size (N, H)\n", " att_w = self.softmax(self.W(out).squeeze(-1), dim=1).unsqueeze(-1)\n", " utter_rep = torch.sum(out * att_w, dim=1)\n", " \n", " out = self.pred_layer(utter_rep)\n", " return out, _" ] }, { "cell_type": "markdown", "metadata": { "id": "W7yX8JinM5Ly" }, "source": [ "# Learning rate schedule\n", "- For transformer architecture, the design of learning rate schedule is different from that of CNN.\n", "- Previous works show that the warmup of learning rate is useful for training models with transformer architectures.\n", "- The warmup schedule\n", " - Set learning rate to 0 in the beginning.\n", " - The learning rate increases linearly from 0 to initial learning rate during warmup period." ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "id": "ykt0N1nVJJi2" }, "outputs": [], "source": [ "import math\n", "\n", "import torch\n", "from torch.optim import Optimizer\n", "from torch.optim.lr_scheduler import LambdaLR\n", "\n", "\n", "def get_cosine_schedule_with_warmup(\n", " optimizer: Optimizer,\n", " num_warmup_steps: int,\n", " num_training_steps: int,\n", " num_cycles: float = 0.5,\n", " last_epoch: int = -1,\n", "):\n", " \"\"\"\n", " Create a schedule with a learning rate that decreases following the values of the cosine function between the\n", " initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the\n", " initial lr set in the optimizer.\n", "\n", " Args:\n", " optimizer (:class:`~torch.optim.Optimizer`):\n", " The optimizer for which to schedule the learning rate.\n", " num_warmup_steps (:obj:`int`):\n", " The number of steps for the warmup phase.\n", " num_training_steps (:obj:`int`):\n", " The total number of training steps.\n", " num_cycles (:obj:`float`, `optional`, defaults to 0.5):\n", " The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0\n", " following a half-cosine).\n", " last_epoch (:obj:`int`, `optional`, defaults to -1):\n", " The index of the last epoch when resuming training.\n", "\n", " Return:\n", " :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.\n", " \"\"\"\n", " def lr_lambda(current_step):\n", " # Warmup\n", " if current_step < num_warmup_steps:\n", " return float(current_step) / float(max(1, num_warmup_steps))\n", " # decadence\n", " progress = float(current_step - num_warmup_steps) / float(\n", " max(1, num_training_steps - num_warmup_steps)\n", " )\n", " return max(\n", " 0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))\n", " )\n", "\n", " return LambdaLR(optimizer, lr_lambda, last_epoch)" ] }, { "cell_type": "markdown", "metadata": { "id": "-LN2XkteM_uH" }, "source": [ "# Model Function\n", "- Model forward function." ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "id": "N-rr8529JMz0" }, "outputs": [], "source": [ "import torch\n", "\n", "\n", "def model_fn(batch, model, criterion, device):\n", " \"\"\"Forward a batch through the model.\"\"\"\n", "\n", " mels, labels = batch\n", " mels = mels.to(device)\n", " labels = labels.to(device)\n", "\n", " outs, outs_length = model(mels)\n", " \n", " loss = criterion(outs, labels)\n", "\n", " # Get the speaker id with highest probability.\n", " preds = outs.argmax(1)\n", " # Compute accuracy.\n", " accuracy = torch.mean((preds == labels).float())\n", "\n", " return loss, accuracy" ] }, { "cell_type": "markdown", "metadata": { "id": "cwM_xyOtNCI2" }, "source": [ "# Validate\n", "- Calculate accuracy of the validation set." ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "id": "YAiv6kpdJRTJ" }, "outputs": [], "source": [ "from tqdm import tqdm\n", "import torch\n", "\n", "\n", "def valid(dataloader, model, criterion, device): \n", " \"\"\"Validate on validation set.\"\"\"\n", "\n", " model.eval()\n", " running_loss = 0.0\n", " running_accuracy = 0.0\n", " pbar = tqdm(total=len(dataloader.dataset), ncols=0, desc=\"Valid\", unit=\" uttr\")\n", "\n", " for i, batch in enumerate(dataloader):\n", " with torch.no_grad():\n", " loss, accuracy = model_fn(batch, model, criterion, device)\n", " running_loss += loss.item()\n", " running_accuracy += accuracy.item()\n", "\n", " pbar.update(dataloader.batch_size)\n", " pbar.set_postfix(\n", " loss=f\"{running_loss / (i+1):.2f}\",\n", " accuracy=f\"{running_accuracy / (i+1):.2f}\",\n", " )\n", "\n", " pbar.close()\n", " model.train()\n", "\n", " return running_accuracy / len(dataloader)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Additive-Margin-Softmax\n", "Reference: https://github.com/Leethony/Additive-Margin-Softmax-Loss-Pytorch" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.nn.functional as F\n", "\n", "class AdMSoftmaxLoss(nn.Module):\n", "\n", " def __init__(self, in_features, out_features, s=30.0, m=0.4):\n", " '''\n", " AM Softmax Loss\n", " '''\n", " super(AdMSoftmaxLoss, self).__init__()\n", " self.s = s\n", " self.m = m\n", " self.in_features = in_features\n", " self.out_features = out_features\n", " #self.fc = nn.Linear(in_features, out_features, bias=False)\n", "\n", " def forward(self, x, labels):\n", " '''\n", " input shape (N, in_features)\n", " '''\n", " assert len(x) == len(labels)\n", " assert torch.min(labels) >= 0\n", " assert torch.max(labels) < self.out_features\n", " \n", " #for W in self.fc.parameters():\n", " # W = F.normalize(W, dim=1)\n", "\n", " x = F.normalize(x, dim=1)\n", " #x = x.view(-1, x.size(0))\n", " wf = x\n", " numerator = self.s * (torch.diagonal(wf.transpose(0, 1)[labels]) - self.m)\n", " excl = torch.cat([torch.cat((wf[i, :y], wf[i, y+1:])).unsqueeze(0) for i, y in enumerate(labels)], dim=0)\n", " denominator = torch.exp(numerator) + torch.sum(torch.exp(self.s * excl), dim=1)\n", " L = numerator - torch.log(denominator)\n", " return -torch.mean(L)" ] }, { "cell_type": "markdown", "metadata": { "id": "g6ne9G-eNEdG" }, "source": [ "# Main function" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "Usv9s-CuJSG7", "outputId": "f4f6a983-3559-4f36-efae-402bbf790473" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[Info]: Use mps now!\n", "[Info]: Finish loading data!\n", "[Info]: Finish creating model!\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "Train: 100% 2000/2000 [05:20<00:00, 6.25 step/s, accuracy=0.16, loss=15.88, step=2000] \n", "Valid: 100% 5664/5667 [00:55<00:00, 102.81 uttr/s, accuracy=0.24, loss=15.80]\n", "Train: 100% 2000/2000 [05:17<00:00, 6.29 step/s, accuracy=0.31, loss=14.76, step=4000] \n", "Valid: 100% 5664/5667 [00:55<00:00, 101.96 uttr/s, accuracy=0.39, loss=14.58]\n", "Train: 100% 2000/2000 [05:17<00:00, 6.29 step/s, accuracy=0.56, loss=13.00, step=6000] \n", "Valid: 100% 5664/5667 [00:56<00:00, 100.81 uttr/s, accuracy=0.46, loss=13.85]\n", "Train: 100% 2000/2000 [06:03<00:00, 5.51 step/s, accuracy=0.53, loss=13.60, step=8000] \n", "Valid: 100% 5664/5667 [00:55<00:00, 102.79 uttr/s, accuracy=0.48, loss=13.37]\n", "Train: 100% 2000/2000 [05:16<00:00, 6.31 step/s, accuracy=0.53, loss=13.56, step=1e+4] \n", "Valid: 100% 5664/5667 [00:55<00:00, 101.78 uttr/s, accuracy=0.51, loss=12.95]\n", "Train: 0% 0/2000 [00:00 best_accuracy:\n", " best_accuracy = valid_accuracy\n", " best_state_dict = model.state_dict()\n", "\n", " pbar = tqdm(total=valid_steps, ncols=0, desc=\"Train\", unit=\" step\")\n", "\n", " # Save the best model so far.\n", " if (step + 1) % save_steps == 0 and best_state_dict is not None:\n", " torch.save(best_state_dict, save_path)\n", " pbar.write(f\"Step {step + 1}, best model saved. (accuracy={best_accuracy:.4f})\")\n", " train_log({\n", " 'status': 'best_model_save',\n", " 'step': step+1,\n", " 'accuracy': f\"{best_accuracy:.6f}\",\n", " })\n", "\n", " pbar.close()\n", " train_log({\n", " 'status': 'completed'\n", " })\n", "\n", " del train_loader, valid_loader\n", " gc.collect()\n", " \n", "if __name__ == \"__main__\":\n", " main(**parse_args())" ] }, { "cell_type": "markdown", "metadata": { "id": "NLatBYAhNNMx" }, "source": [ "# Inference\n", "\n", "## Dataset of inference" ] }, { "cell_type": "code", "execution_count": 12, "metadata": { "id": "efS4pCmAJXJH" }, "outputs": [], "source": [ "import os\n", "import json\n", "import torch\n", "from pathlib import Path\n", "from torch.utils.data import Dataset\n", "\n", "\n", "from common_function import InferenceDataset, inference_collate_batch\n" ] }, { "cell_type": "markdown", "metadata": { "id": "tl0WnYwxNK_S" }, "source": [ "## Main funcrion of Inference" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 103, "referenced_widgets": [ "6786c2b0e2614ad389620246cb2178f2", "0d592098920140dab61aac5410568c36", "da681e3cc353420cb142d56df0fce231", "401ae722b95c4ff59b836422dbe71edc", "4efbfb7c7cb54276862e5321209d57fa", "f6dcb3ec9c624171966bb889808dbcb3", "100abf072991474abfcee871d2b83f37", "decc12da6f5742ec8b7ec7789ee434ab", "0117ac88c98440b29ab1f452107cbe1a", "3d64942a0eaa409a93df049bb594c062", "8c8721e504cf434eb9e263fb9983969a" ] }, "id": "i8SAbuXEJb2A", "outputId": "3808f409-19c9-426c-dc15-1b88b0c21645" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[Info]: Use mps now!\n", "[Info]: Finish loading data!\n", "[Info]: Finish creating model!\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "a6e9118fb3ff4e96a988a7df19ca7b98", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/8000 [00:00