yushan 3 жил өмнө
commit
3bedf5bf20
1 өөрчлөгдсөн 1440 нэмэгдсэн , 0 устгасан
  1. 1440 0
      ML2023_hw04.ipynb

+ 1440 - 0
ML2023_hw04.ipynb

@@ -0,0 +1,1440 @@
+{
+  "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": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "gtKxUzSgXKj3",
+        "outputId": "3f59402c-95a7-4fbd-a39c-57606590a89c"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "--2023-03-02 05:04:55--  https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partaa\n",
+            "Resolving github.com (github.com)... 140.82.121.4\n",
+            "Connecting to github.com (github.com)|140.82.121.4|:443... connected.\n",
+            "HTTP request sent, awaiting response... 302 Found\n",
+            "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/7646b36b-6033-4a31-bac4-380c4d21d91e?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T050455Z&X-Amz-Expires=300&X-Amz-Signature=ce8c7b9a37cfcd2b905ab2a8aeaa65c7ef1c8b20e1ea07855edb16a56f2d1b3e&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partaa&response-content-type=application%2Foctet-stream [following]\n",
+            "--2023-03-02 05:04:56--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/7646b36b-6033-4a31-bac4-380c4d21d91e?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T050455Z&X-Amz-Expires=300&X-Amz-Signature=ce8c7b9a37cfcd2b905ab2a8aeaa65c7ef1c8b20e1ea07855edb16a56f2d1b3e&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partaa&response-content-type=application%2Foctet-stream\n",
+            "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.111.133, 185.199.110.133, ...\n",
+            "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
+            "HTTP request sent, awaiting response... 200 OK\n",
+            "Length: 1560784333 (1.5G) [application/octet-stream]\n",
+            "Saving to: ‘Dataset.tar.gz.partaa’\n",
+            "\n",
+            "Dataset.tar.gz.part 100%[===================>]   1.45G  31.0MB/s    in 86s     \n",
+            "\n",
+            "2023-03-02 05:06:22 (17.3 MB/s) - ‘Dataset.tar.gz.partaa’ saved [1560784333/1560784333]\n",
+            "\n",
+            "--2023-03-02 05:06:22--  https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partab\n",
+            "Resolving github.com (github.com)... 140.82.121.3\n",
+            "Connecting to github.com (github.com)|140.82.121.3|:443... connected.\n",
+            "HTTP request sent, awaiting response... 302 Found\n",
+            "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/95b45712-6e2f-4a52-96b1-7d88578345fc?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T050622Z&X-Amz-Expires=300&X-Amz-Signature=f53b3f800fa1be29c1c4c0a058747b3f085d64a54a617bff48f198a8704d681e&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partab&response-content-type=application%2Foctet-stream [following]\n",
+            "--2023-03-02 05:06:22--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/95b45712-6e2f-4a52-96b1-7d88578345fc?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T050622Z&X-Amz-Expires=300&X-Amz-Signature=f53b3f800fa1be29c1c4c0a058747b3f085d64a54a617bff48f198a8704d681e&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partab&response-content-type=application%2Foctet-stream\n",
+            "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.111.133, 185.199.110.133, ...\n",
+            "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
+            "HTTP request sent, awaiting response... 200 OK\n",
+            "Length: 1560784333 (1.5G) [application/octet-stream]\n",
+            "Saving to: ‘Dataset.tar.gz.partab’\n",
+            "\n",
+            "Dataset.tar.gz.part 100%[===================>]   1.45G  5.37MB/s    in 2m 49s  \n",
+            "\n",
+            "2023-03-02 05:09:11 (8.82 MB/s) - ‘Dataset.tar.gz.partab’ saved [1560784333/1560784333]\n",
+            "\n",
+            "--2023-03-02 05:09:11--  https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partac\n",
+            "Resolving github.com (github.com)... 140.82.121.4\n",
+            "Connecting to github.com (github.com)|140.82.121.4|:443... connected.\n",
+            "HTTP request sent, awaiting response... 302 Found\n",
+            "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/0c9d42d3-95b7-4ca4-b57c-ab1a66a5564d?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T050911Z&X-Amz-Expires=300&X-Amz-Signature=a10ad7a33c85a994cea24c2493cc018fa1d287ee75cef1f2b36c919cc173b03d&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partac&response-content-type=application%2Foctet-stream [following]\n",
+            "--2023-03-02 05:09:11--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/0c9d42d3-95b7-4ca4-b57c-ab1a66a5564d?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T050911Z&X-Amz-Expires=300&X-Amz-Signature=a10ad7a33c85a994cea24c2493cc018fa1d287ee75cef1f2b36c919cc173b03d&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partac&response-content-type=application%2Foctet-stream\n",
+            "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
+            "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
+            "HTTP request sent, awaiting response... 200 OK\n",
+            "Length: 1560784333 (1.5G) [application/octet-stream]\n",
+            "Saving to: ‘Dataset.tar.gz.partac’\n",
+            "\n",
+            "Dataset.tar.gz.part 100%[===================>]   1.45G  31.0MB/s    in 98s     \n",
+            "\n",
+            "2023-03-02 05:10:50 (15.2 MB/s) - ‘Dataset.tar.gz.partac’ saved [1560784333/1560784333]\n",
+            "\n",
+            "--2023-03-02 05:10:50--  https://github.com/googly-mingto/ML2023HW4/releases/download/data/Dataset.tar.gz.partad\n",
+            "Resolving github.com (github.com)... 140.82.121.3\n",
+            "Connecting to github.com (github.com)|140.82.121.3|:443... connected.\n",
+            "HTTP request sent, awaiting response... 302 Found\n",
+            "Location: https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/0ee11da6-8c96-4463-b084-cea8f95d26e9?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T051050Z&X-Amz-Expires=300&X-Amz-Signature=5b78f08d84962d376572a5bfd3e3d9beb020b61a93013b752257739f86efed6c&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partad&response-content-type=application%2Foctet-stream [following]\n",
+            "--2023-03-02 05:10:50--  https://objects.githubusercontent.com/github-production-release-asset-2e65be/606989982/0ee11da6-8c96-4463-b084-cea8f95d26e9?X-Amz-Algorithm=AWS4-HMAC-SHA256&X-Amz-Credential=AKIAIWNJYAX4CSVEH53A%2F20230302%2Fus-east-1%2Fs3%2Faws4_request&X-Amz-Date=20230302T051050Z&X-Amz-Expires=300&X-Amz-Signature=5b78f08d84962d376572a5bfd3e3d9beb020b61a93013b752257739f86efed6c&X-Amz-SignedHeaders=host&actor_id=0&key_id=0&repo_id=606989982&response-content-disposition=attachment%3B%20filename%3DDataset.tar.gz.partad&response-content-type=application%2Foctet-stream\n",
+            "Resolving objects.githubusercontent.com (objects.githubusercontent.com)... 185.199.108.133, 185.199.111.133, 185.199.110.133, ...\n",
+            "Connecting to objects.githubusercontent.com (objects.githubusercontent.com)|185.199.108.133|:443... connected.\n",
+            "HTTP request sent, awaiting response... 200 OK\n",
+            "Length: 1560784336 (1.5G) [application/octet-stream]\n",
+            "Saving to: ‘Dataset.tar.gz.partad’\n",
+            "\n",
+            "Dataset.tar.gz.part 100%[===================>]   1.45G  14.3MB/s    in 3m 52s  \n",
+            "\n",
+            "2023-03-02 05:14:43 (6.42 MB/s) - ‘Dataset.tar.gz.partad’ saved [1560784336/1560784336]\n",
+            "\n",
+            "tar: Ignoring unknown extended header keyword 'LIBARCHIVE.xattr.com.apple.macl'\n"
+          ]
+        }
+      ],
+      "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": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "U6Y1cfpDfpON",
+        "outputId": "6ba26637-5c7b-48a9-be0b-1f10ba76590a"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "tar (child): Dataset.tar.gz: Cannot open: No such file or directory\n",
+            "tar (child): Error is not recoverable: exiting now\n",
+            "tar: Child returned status 2\n",
+            "tar: Error is not recoverable: exiting now\n"
+          ]
+        }
+      ],
+      "source": [
+        "!tar zxf Dataset.tar.gz"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "E6burzCXIyuA"
+      },
+      "outputs": [],
+      "source": [
+        "import numpy as np\n",
+        "import torch\n",
+        "import random\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.backends.cudnn.benchmark = False\n",
+        "    torch.backends.cudnn.deterministic = True\n",
+        "\n",
+        "set_seed(87)"
+      ]
+    },
+    {
+      "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": null,
+      "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",
+        " \n",
+        " \n",
+        "class myDataset(Dataset):\n",
+        "\tdef __init__(self, data_dir, segment_len=128):\n",
+        "\t\tself.data_dir = data_dir\n",
+        "\t\tself.segment_len = segment_len\n",
+        "\t\n",
+        "\t\t# Load the mapping from speaker neme to their corresponding id. \n",
+        "\t\tmapping_path = Path(data_dir) / \"mapping.json\"\n",
+        "\t\tmapping = json.load(mapping_path.open())\n",
+        "\t\tself.speaker2id = mapping[\"speaker2id\"]\n",
+        "\t\n",
+        "\t\t# Load metadata of training data.\n",
+        "\t\tmetadata_path = Path(data_dir) / \"metadata.json\"\n",
+        "\t\tmetadata = json.load(open(metadata_path))[\"speakers\"]\n",
+        "\t\n",
+        "\t\t# Get the total number of speaker.\n",
+        "\t\tself.speaker_num = len(metadata.keys())\n",
+        "\t\tself.data = []\n",
+        "\t\tfor speaker in metadata.keys():\n",
+        "\t\t\tfor utterances in metadata[speaker]:\n",
+        "\t\t\t\tself.data.append([utterances[\"feature_path\"], self.speaker2id[speaker]])\n",
+        " \n",
+        "\tdef __len__(self):\n",
+        "\t\t\treturn len(self.data)\n",
+        " \n",
+        "\tdef __getitem__(self, index):\n",
+        "\t\tfeat_path, speaker = self.data[index]\n",
+        "\t\t# Load preprocessed mel-spectrogram.\n",
+        "\t\tmel = torch.load(os.path.join(self.data_dir, feat_path))\n",
+        "\n",
+        "\t\t# Segmemt mel-spectrogram into \"segment_len\" frames.\n",
+        "\t\tif len(mel) > self.segment_len:\n",
+        "\t\t\t# Randomly get the starting point of the segment.\n",
+        "\t\t\tstart = random.randint(0, len(mel) - self.segment_len)\n",
+        "\t\t\t# Get a segment with \"segment_len\" frames.\n",
+        "\t\t\tmel = torch.FloatTensor(mel[start:start+self.segment_len])\n",
+        "\t\telse:\n",
+        "\t\t\tmel = torch.FloatTensor(mel)\n",
+        "\t\t# Turn the speaker id into long for computing loss later.\n",
+        "\t\tspeaker = torch.FloatTensor([speaker]).long()\n",
+        "\t\treturn mel, speaker\n",
+        " \n",
+        "\tdef get_speaker_number(self):\n",
+        "\t\treturn self.speaker_num"
+      ]
+    },
+    {
+      "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": null,
+      "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",
+        "def collate_batch(batch):\n",
+        "\t# Process features within a batch.\n",
+        "\t\"\"\"Collate a batch of data.\"\"\"\n",
+        "\tmel, speaker = zip(*batch)\n",
+        "\t# Because we train the model batch by batch, we need to pad the features in the same batch to make their lengths the same.\n",
+        "\tmel = pad_sequence(mel, batch_first=True, padding_value=-20)    # pad log 10^(-20) which is very small value.\n",
+        "\t# mel: (batch size, length, 40)\n",
+        "\treturn mel, torch.FloatTensor(speaker).long()\n",
+        "\n",
+        "\n",
+        "def get_dataloader(data_dir, batch_size, n_workers):\n",
+        "\t\"\"\"Generate dataloader\"\"\"\n",
+        "\tdataset = myDataset(data_dir)\n",
+        "\tspeaker_num = dataset.get_speaker_number()\n",
+        "\t# Split dataset into training dataset and validation dataset\n",
+        "\ttrainlen = int(0.9 * len(dataset))\n",
+        "\tlengths = [trainlen, len(dataset) - trainlen]\n",
+        "\ttrainset, validset = random_split(dataset, lengths)\n",
+        "\n",
+        "\ttrain_loader = DataLoader(\n",
+        "\t\ttrainset,\n",
+        "\t\tbatch_size=batch_size,\n",
+        "\t\tshuffle=True,\n",
+        "\t\tdrop_last=True,\n",
+        "\t\tnum_workers=n_workers,\n",
+        "\t\tpin_memory=True,\n",
+        "\t\tcollate_fn=collate_batch,\n",
+        "\t)\n",
+        "\tvalid_loader = DataLoader(\n",
+        "\t\tvalidset,\n",
+        "\t\tbatch_size=batch_size,\n",
+        "\t\tnum_workers=n_workers,\n",
+        "\t\tdrop_last=True,\n",
+        "\t\tpin_memory=True,\n",
+        "\t\tcollate_fn=collate_batch,\n",
+        "\t)\n",
+        "\n",
+        "\treturn 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": null,
+      "metadata": {
+        "id": "iXZ5B0EKJGs8"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "import torch.nn.functional as F\n",
+        "\n",
+        "\n",
+        "class Classifier(nn.Module):\n",
+        "\tdef __init__(self, d_model=80, n_spks=600, dropout=0.1):\n",
+        "\t\tsuper().__init__()\n",
+        "\t\t# Project the dimension of features from that of input into d_model.\n",
+        "\t\tself.prenet = nn.Linear(40, d_model)\n",
+        "\t\t# TODO:\n",
+        "\t\t#   Change Transformer to Conformer.\n",
+        "\t\t#   https://arxiv.org/abs/2005.08100\n",
+        "\t\tself.encoder_layer = nn.TransformerEncoderLayer(\n",
+        "\t\t\td_model=d_model, dim_feedforward=256, nhead=2\n",
+        "\t\t)\n",
+        "\t\t# self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=2)\n",
+        "\n",
+        "\t\t# Project the the dimension of features from d_model into speaker nums.\n",
+        "\t\tself.pred_layer = nn.Sequential(\n",
+        "\t\t\tnn.Linear(d_model, d_model),\n",
+        "\t\t\tnn.Sigmoid(),\n",
+        "\t\t\tnn.Linear(d_model, n_spks),\n",
+        "\t\t)\n",
+        "\n",
+        "\tdef forward(self, mels):\n",
+        "\t\t\"\"\"\n",
+        "\t\targs:\n",
+        "\t\t\tmels: (batch size, length, 40)\n",
+        "\t\treturn:\n",
+        "\t\t\tout: (batch size, n_spks)\n",
+        "\t\t\"\"\"\n",
+        "\t\t# out: (batch size, length, d_model)\n",
+        "\t\tout = self.prenet(mels)\n",
+        "\t\t# out: (length, batch size, d_model)\n",
+        "\t\tout = out.permute(1, 0, 2)\n",
+        "\t\t# The encoder layer expect features in the shape of (length, batch size, d_model).\n",
+        "\t\tout = self.encoder_layer(out)\n",
+        "\t\t# out: (batch size, length, d_model)\n",
+        "\t\tout = out.transpose(0, 1)\n",
+        "\t\t# mean pooling\n",
+        "\t\tstats = out.mean(dim=1)\n",
+        "\n",
+        "\t\t# out: (batch, n_spks)\n",
+        "\t\tout = self.pred_layer(stats)\n",
+        "\t\treturn 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": null,
+      "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",
+        "\toptimizer: Optimizer,\n",
+        "\tnum_warmup_steps: int,\n",
+        "\tnum_training_steps: int,\n",
+        "\tnum_cycles: float = 0.5,\n",
+        "\tlast_epoch: int = -1,\n",
+        "):\n",
+        "\t\"\"\"\n",
+        "\tCreate a schedule with a learning rate that decreases following the values of the cosine function between the\n",
+        "\tinitial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the\n",
+        "\tinitial lr set in the optimizer.\n",
+        "\n",
+        "\tArgs:\n",
+        "\t\toptimizer (:class:`~torch.optim.Optimizer`):\n",
+        "\t\tThe optimizer for which to schedule the learning rate.\n",
+        "\t\tnum_warmup_steps (:obj:`int`):\n",
+        "\t\tThe number of steps for the warmup phase.\n",
+        "\t\tnum_training_steps (:obj:`int`):\n",
+        "\t\tThe total number of training steps.\n",
+        "\t\tnum_cycles (:obj:`float`, `optional`, defaults to 0.5):\n",
+        "\t\tThe number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0\n",
+        "\t\tfollowing a half-cosine).\n",
+        "\t\tlast_epoch (:obj:`int`, `optional`, defaults to -1):\n",
+        "\t\tThe index of the last epoch when resuming training.\n",
+        "\n",
+        "\tReturn:\n",
+        "\t\t:obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.\n",
+        "\t\"\"\"\n",
+        "\tdef lr_lambda(current_step):\n",
+        "\t\t# Warmup\n",
+        "\t\tif current_step < num_warmup_steps:\n",
+        "\t\t\treturn float(current_step) / float(max(1, num_warmup_steps))\n",
+        "\t\t# decadence\n",
+        "\t\tprogress = float(current_step - num_warmup_steps) / float(\n",
+        "\t\t\tmax(1, num_training_steps - num_warmup_steps)\n",
+        "\t\t)\n",
+        "\t\treturn max(\n",
+        "\t\t\t0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))\n",
+        "\t\t)\n",
+        "\n",
+        "\treturn 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": null,
+      "metadata": {
+        "id": "N-rr8529JMz0"
+      },
+      "outputs": [],
+      "source": [
+        "import torch\n",
+        "\n",
+        "\n",
+        "def model_fn(batch, model, criterion, device):\n",
+        "\t\"\"\"Forward a batch through the model.\"\"\"\n",
+        "\n",
+        "\tmels, labels = batch\n",
+        "\tmels = mels.to(device)\n",
+        "\tlabels = labels.to(device)\n",
+        "\n",
+        "\touts = model(mels)\n",
+        "\n",
+        "\tloss = criterion(outs, labels)\n",
+        "\n",
+        "\t# Get the speaker id with highest probability.\n",
+        "\tpreds = outs.argmax(1)\n",
+        "\t# Compute accuracy.\n",
+        "\taccuracy = torch.mean((preds == labels).float())\n",
+        "\n",
+        "\treturn loss, accuracy"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "cwM_xyOtNCI2"
+      },
+      "source": [
+        "# Validate\n",
+        "- Calculate accuracy of the validation set."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YAiv6kpdJRTJ"
+      },
+      "outputs": [],
+      "source": [
+        "from tqdm import tqdm\n",
+        "import torch\n",
+        "\n",
+        "\n",
+        "def valid(dataloader, model, criterion, device): \n",
+        "\t\"\"\"Validate on validation set.\"\"\"\n",
+        "\n",
+        "\tmodel.eval()\n",
+        "\trunning_loss = 0.0\n",
+        "\trunning_accuracy = 0.0\n",
+        "\tpbar = tqdm(total=len(dataloader.dataset), ncols=0, desc=\"Valid\", unit=\" uttr\")\n",
+        "\n",
+        "\tfor i, batch in enumerate(dataloader):\n",
+        "\t\twith torch.no_grad():\n",
+        "\t\t\tloss, accuracy = model_fn(batch, model, criterion, device)\n",
+        "\t\t\trunning_loss += loss.item()\n",
+        "\t\t\trunning_accuracy += accuracy.item()\n",
+        "\n",
+        "\t\tpbar.update(dataloader.batch_size)\n",
+        "\t\tpbar.set_postfix(\n",
+        "\t\t\tloss=f\"{running_loss / (i+1):.2f}\",\n",
+        "\t\t\taccuracy=f\"{running_accuracy / (i+1):.2f}\",\n",
+        "\t\t)\n",
+        "\n",
+        "\tpbar.close()\n",
+        "\tmodel.train()\n",
+        "\n",
+        "\treturn running_accuracy / len(dataloader)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "g6ne9G-eNEdG"
+      },
+      "source": [
+        "# Main function"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "Usv9s-CuJSG7",
+        "outputId": "f4f6a983-3559-4f36-efae-402bbf790473"
+      },
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "[Info]: Use cuda now!\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "/usr/local/lib/python3.8/dist-packages/torch/utils/data/dataloader.py:554: UserWarning: This DataLoader will create 8 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
+            "  warnings.warn(_create_warning_msg(\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "[Info]: Finish loading data!\n",
+            "[Info]: Finish creating model!\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Train: 100% 2000/2000 [01:39<00:00, 20.15 step/s, accuracy=0.00, loss=5.34, step=2000]\n",
+            "Valid: 100% 5664/5667 [00:08<00:00, 646.60 uttr/s, accuracy=0.03, loss=5.30]\n",
+            "Train: 100% 2000/2000 [01:05<00:00, 30.48 step/s, accuracy=0.09, loss=4.58, step=4000]\n",
+            "Valid: 100% 5664/5667 [00:09<00:00, 592.93 uttr/s, accuracy=0.09, loss=4.65]\n",
+            "Train: 100% 2000/2000 [00:55<00:00, 36.21 step/s, accuracy=0.19, loss=3.84, step=6000]\n",
+            "Valid: 100% 5664/5667 [00:06<00:00, 909.71 uttr/s, accuracy=0.14, loss=4.25] \n",
+            "Train: 100% 2000/2000 [01:03<00:00, 31.46 step/s, accuracy=0.16, loss=4.17, step=8000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1575.18 uttr/s, accuracy=0.16, loss=4.06]\n",
+            "Train: 100% 2000/2000 [00:51<00:00, 38.66 step/s, accuracy=0.16, loss=4.00, step=1e+4]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1629.40 uttr/s, accuracy=0.19, loss=3.86]\n",
+            "Train:   0% 10/2000 [00:00<00:50, 39.27 step/s, accuracy=0.22, loss=3.96, step=1e+4]"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Step 10000, best model saved. (accuracy=0.1933)\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Train: 100% 2000/2000 [00:47<00:00, 41.88 step/s, accuracy=0.19, loss=3.55, step=12000]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1279.71 uttr/s, accuracy=0.22, loss=3.72]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 43.75 step/s, accuracy=0.25, loss=3.63, step=14000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1714.78 uttr/s, accuracy=0.24, loss=3.57]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 44.12 step/s, accuracy=0.22, loss=3.98, step=16000]\n",
+            "Valid: 100% 5664/5667 [00:05<00:00, 1049.05 uttr/s, accuracy=0.26, loss=3.50]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 44.31 step/s, accuracy=0.16, loss=3.48, step=18000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1765.13 uttr/s, accuracy=0.28, loss=3.37]\n",
+            "Train: 100% 2000/2000 [00:44<00:00, 44.49 step/s, accuracy=0.28, loss=3.25, step=2e+4]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1379.63 uttr/s, accuracy=0.30, loss=3.28]\n",
+            "Train:   0% 6/2000 [00:00<01:08, 29.05 step/s, accuracy=0.34, loss=2.96, step=2e+4]"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Step 20000, best model saved. (accuracy=0.2994)\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Train: 100% 2000/2000 [00:45<00:00, 44.24 step/s, accuracy=0.47, loss=2.32, step=22000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1698.76 uttr/s, accuracy=0.32, loss=3.17]\n",
+            "Train: 100% 2000/2000 [00:44<00:00, 44.56 step/s, accuracy=0.38, loss=3.30, step=24000]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1328.95 uttr/s, accuracy=0.33, loss=3.08]\n",
+            "Train: 100% 2000/2000 [00:46<00:00, 42.84 step/s, accuracy=0.44, loss=2.49, step=26000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1708.26 uttr/s, accuracy=0.35, loss=3.01]\n",
+            "Train: 100% 2000/2000 [00:44<00:00, 44.61 step/s, accuracy=0.34, loss=2.98, step=28000]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1164.28 uttr/s, accuracy=0.36, loss=2.94]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 43.79 step/s, accuracy=0.28, loss=3.35, step=3e+4]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1691.95 uttr/s, accuracy=0.38, loss=2.87]\n",
+            "Train:   0% 10/2000 [00:00<00:40, 48.99 step/s, accuracy=0.31, loss=3.11, step=3e+4]"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Step 30000, best model saved. (accuracy=0.3766)\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Train: 100% 2000/2000 [00:46<00:00, 42.87 step/s, accuracy=0.59, loss=2.10, step=32000]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1266.74 uttr/s, accuracy=0.39, loss=2.83]\n",
+            "Train: 100% 2000/2000 [00:47<00:00, 42.50 step/s, accuracy=0.47, loss=2.61, step=34000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1695.46 uttr/s, accuracy=0.40, loss=2.79]\n",
+            "Train: 100% 2000/2000 [00:46<00:00, 42.68 step/s, accuracy=0.47, loss=2.27, step=36000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1695.12 uttr/s, accuracy=0.41, loss=2.71]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 44.31 step/s, accuracy=0.31, loss=2.92, step=38000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1699.13 uttr/s, accuracy=0.42, loss=2.66]\n",
+            "Train: 100% 2000/2000 [00:47<00:00, 41.84 step/s, accuracy=0.44, loss=3.49, step=4e+4]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1407.64 uttr/s, accuracy=0.43, loss=2.65]\n",
+            "Train:   0% 10/2000 [00:00<00:41, 48.40 step/s, accuracy=0.50, loss=2.24, step=4e+4]"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Step 40000, best model saved. (accuracy=0.4287)\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Train: 100% 2000/2000 [00:45<00:00, 43.89 step/s, accuracy=0.44, loss=2.31, step=42000]\n",
+            "Valid: 100% 5664/5667 [00:05<00:00, 1062.44 uttr/s, accuracy=0.42, loss=2.63]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 43.58 step/s, accuracy=0.50, loss=1.89, step=44000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1633.34 uttr/s, accuracy=0.44, loss=2.58]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 44.32 step/s, accuracy=0.34, loss=3.01, step=46000]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1176.49 uttr/s, accuracy=0.45, loss=2.52]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 43.93 step/s, accuracy=0.53, loss=2.38, step=48000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1690.24 uttr/s, accuracy=0.45, loss=2.49]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 44.33 step/s, accuracy=0.44, loss=1.99, step=5e+4]\n",
+            "Valid: 100% 5664/5667 [00:04<00:00, 1318.06 uttr/s, accuracy=0.45, loss=2.49]\n",
+            "Train:   0% 10/2000 [00:00<00:44, 44.97 step/s, accuracy=0.53, loss=2.17, step=5e+4]"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Step 50000, best model saved. (accuracy=0.4536)\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Train: 100% 2000/2000 [00:45<00:00, 43.60 step/s, accuracy=0.56, loss=2.07, step=52000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1704.31 uttr/s, accuracy=0.47, loss=2.44]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 44.30 step/s, accuracy=0.47, loss=2.27, step=54000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1416.50 uttr/s, accuracy=0.47, loss=2.45]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 44.31 step/s, accuracy=0.38, loss=2.47, step=56000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1636.55 uttr/s, accuracy=0.47, loss=2.41]\n",
+            "Train: 100% 2000/2000 [00:46<00:00, 43.13 step/s, accuracy=0.53, loss=2.40, step=58000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1733.87 uttr/s, accuracy=0.47, loss=2.42]\n",
+            "Train: 100% 2000/2000 [00:46<00:00, 43.38 step/s, accuracy=0.56, loss=1.91, step=6e+4]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1722.16 uttr/s, accuracy=0.48, loss=2.39]\n",
+            "Train:   0% 9/2000 [00:00<00:44, 44.52 step/s, accuracy=0.56, loss=2.01, step=6e+4]"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Step 60000, best model saved. (accuracy=0.4802)\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stderr",
+          "text": [
+            "Train: 100% 2000/2000 [00:46<00:00, 43.09 step/s, accuracy=0.38, loss=2.42, step=62000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1652.15 uttr/s, accuracy=0.48, loss=2.36]\n",
+            "Train: 100% 2000/2000 [00:45<00:00, 43.60 step/s, accuracy=0.53, loss=2.07, step=64000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1482.73 uttr/s, accuracy=0.48, loss=2.39]\n",
+            "Train: 100% 2000/2000 [00:47<00:00, 42.54 step/s, accuracy=0.53, loss=2.47, step=66000]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1674.74 uttr/s, accuracy=0.48, loss=2.40]\n",
+            "Train: 100% 2000/2000 [00:46<00:00, 43.00 step/s, accuracy=0.41, loss=2.61, step=68000]\n",
+            "Valid: 100% 5664/5667 [00:05<00:00, 990.67 uttr/s, accuracy=0.48, loss=2.37] \n",
+            "Train: 100% 2000/2000 [00:48<00:00, 41.57 step/s, accuracy=0.69, loss=1.81, step=7e+4]\n",
+            "Valid: 100% 5664/5667 [00:03<00:00, 1500.45 uttr/s, accuracy=0.47, loss=2.42]\n",
+            "Train:   0% 0/2000 [00:00<?, ? step/s]\n"
+          ]
+        },
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "Step 70000, best model saved. (accuracy=0.4834)\n"
+          ]
+        }
+      ],
+      "source": [
+        "from tqdm import tqdm\n",
+        "\n",
+        "import torch\n",
+        "import torch.nn as nn\n",
+        "from torch.optim import AdamW\n",
+        "from torch.utils.data import DataLoader, random_split\n",
+        "\n",
+        "\n",
+        "def parse_args():\n",
+        "\t\"\"\"arguments\"\"\"\n",
+        "\tconfig = {\n",
+        "\t\t\"data_dir\": \"./Dataset\",\n",
+        "\t\t\"save_path\": \"model.ckpt\",\n",
+        "\t\t\"batch_size\": 32,\n",
+        "\t\t\"n_workers\": 8,\n",
+        "\t\t\"valid_steps\": 2000,\n",
+        "\t\t\"warmup_steps\": 1000,\n",
+        "\t\t\"save_steps\": 10000,\n",
+        "\t\t\"total_steps\": 70000,\n",
+        "\t}\n",
+        "\n",
+        "\treturn config\n",
+        "\n",
+        "\n",
+        "def main(\n",
+        "\tdata_dir,\n",
+        "\tsave_path,\n",
+        "\tbatch_size,\n",
+        "\tn_workers,\n",
+        "\tvalid_steps,\n",
+        "\twarmup_steps,\n",
+        "\ttotal_steps,\n",
+        "\tsave_steps,\n",
+        "):\n",
+        "\t\"\"\"Main function.\"\"\"\n",
+        "\tdevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "\tprint(f\"[Info]: Use {device} now!\")\n",
+        "\n",
+        "\ttrain_loader, valid_loader, speaker_num = get_dataloader(data_dir, batch_size, n_workers)\n",
+        "\ttrain_iterator = iter(train_loader)\n",
+        "\tprint(f\"[Info]: Finish loading data!\",flush = True)\n",
+        "\n",
+        "\tmodel = Classifier(n_spks=speaker_num).to(device)\n",
+        "\tcriterion = nn.CrossEntropyLoss()\n",
+        "\toptimizer = AdamW(model.parameters(), lr=1e-3)\n",
+        "\tscheduler = get_cosine_schedule_with_warmup(optimizer, warmup_steps, total_steps)\n",
+        "\tprint(f\"[Info]: Finish creating model!\",flush = True)\n",
+        "\n",
+        "\tbest_accuracy = -1.0\n",
+        "\tbest_state_dict = None\n",
+        "\n",
+        "\tpbar = tqdm(total=valid_steps, ncols=0, desc=\"Train\", unit=\" step\")\n",
+        "\n",
+        "\tfor step in range(total_steps):\n",
+        "\t\t# Get data\n",
+        "\t\ttry:\n",
+        "\t\t\tbatch = next(train_iterator)\n",
+        "\t\texcept StopIteration:\n",
+        "\t\t\ttrain_iterator = iter(train_loader)\n",
+        "\t\t\tbatch = next(train_iterator)\n",
+        "\n",
+        "\t\tloss, accuracy = model_fn(batch, model, criterion, device)\n",
+        "\t\tbatch_loss = loss.item()\n",
+        "\t\tbatch_accuracy = accuracy.item()\n",
+        "\n",
+        "\t\t# Updata model\n",
+        "\t\tloss.backward()\n",
+        "\t\toptimizer.step()\n",
+        "\t\tscheduler.step()\n",
+        "\t\toptimizer.zero_grad()\n",
+        "\n",
+        "\t\t# Log\n",
+        "\t\tpbar.update()\n",
+        "\t\tpbar.set_postfix(\n",
+        "\t\t\tloss=f\"{batch_loss:.2f}\",\n",
+        "\t\t\taccuracy=f\"{batch_accuracy:.2f}\",\n",
+        "\t\t\tstep=step + 1,\n",
+        "\t\t)\n",
+        "\n",
+        "\t\t# Do validation\n",
+        "\t\tif (step + 1) % valid_steps == 0:\n",
+        "\t\t\tpbar.close()\n",
+        "\n",
+        "\t\t\tvalid_accuracy = valid(valid_loader, model, criterion, device)\n",
+        "\n",
+        "\t\t\t# keep the best model\n",
+        "\t\t\tif valid_accuracy > best_accuracy:\n",
+        "\t\t\t\tbest_accuracy = valid_accuracy\n",
+        "\t\t\t\tbest_state_dict = model.state_dict()\n",
+        "\n",
+        "\t\t\tpbar = tqdm(total=valid_steps, ncols=0, desc=\"Train\", unit=\" step\")\n",
+        "\n",
+        "\t\t# Save the best model so far.\n",
+        "\t\tif (step + 1) % save_steps == 0 and best_state_dict is not None:\n",
+        "\t\t\ttorch.save(best_state_dict, save_path)\n",
+        "\t\t\tpbar.write(f\"Step {step + 1}, best model saved. (accuracy={best_accuracy:.4f})\")\n",
+        "\n",
+        "\tpbar.close()\n",
+        "\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "\tmain(**parse_args())"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "NLatBYAhNNMx"
+      },
+      "source": [
+        "# Inference\n",
+        "\n",
+        "## Dataset of inference"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "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",
+        "class InferenceDataset(Dataset):\n",
+        "\tdef __init__(self, data_dir):\n",
+        "\t\ttestdata_path = Path(data_dir) / \"testdata.json\"\n",
+        "\t\tmetadata = json.load(testdata_path.open())\n",
+        "\t\tself.data_dir = data_dir\n",
+        "\t\tself.data = metadata[\"utterances\"]\n",
+        "\n",
+        "\tdef __len__(self):\n",
+        "\t\treturn len(self.data)\n",
+        "\n",
+        "\tdef __getitem__(self, index):\n",
+        "\t\tutterance = self.data[index]\n",
+        "\t\tfeat_path = utterance[\"feature_path\"]\n",
+        "\t\tmel = torch.load(os.path.join(self.data_dir, feat_path))\n",
+        "\n",
+        "\t\treturn feat_path, mel\n",
+        "\n",
+        "\n",
+        "def inference_collate_batch(batch):\n",
+        "\t\"\"\"Collate a batch of data.\"\"\"\n",
+        "\tfeat_paths, mels = zip(*batch)\n",
+        "\n",
+        "\treturn feat_paths, torch.stack(mels)"
+      ]
+    },
+    {
+      "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": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "[Info]: Use cuda now!\n",
+            "[Info]: Finish loading data!\n",
+            "[Info]: Finish creating model!\n"
+          ]
+        },
+        {
+          "output_type": "display_data",
+          "data": {
+            "text/plain": [
+              "  0%|          | 0/8000 [00:00<?, ?it/s]"
+            ],
+            "application/vnd.jupyter.widget-view+json": {
+              "version_major": 2,
+              "version_minor": 0,
+              "model_id": "6786c2b0e2614ad389620246cb2178f2"
+            }
+          },
+          "metadata": {}
+        }
+      ],
+      "source": [
+        "import json\n",
+        "import csv\n",
+        "from pathlib import Path\n",
+        "from tqdm.notebook import tqdm\n",
+        "\n",
+        "import torch\n",
+        "from torch.utils.data import DataLoader\n",
+        "\n",
+        "def parse_args():\n",
+        "\t\"\"\"arguments\"\"\"\n",
+        "\tconfig = {\n",
+        "\t\t\"data_dir\": \"./Dataset\",\n",
+        "\t\t\"model_path\": \"./model.ckpt\",\n",
+        "\t\t\"output_path\": \"./output.csv\",\n",
+        "\t}\n",
+        "\n",
+        "\treturn config\n",
+        "\n",
+        "\n",
+        "def main(\n",
+        "\tdata_dir,\n",
+        "\tmodel_path,\n",
+        "\toutput_path,\n",
+        "):\n",
+        "\t\"\"\"Main function.\"\"\"\n",
+        "\tdevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+        "\tprint(f\"[Info]: Use {device} now!\")\n",
+        "\n",
+        "\tmapping_path = Path(data_dir) / \"mapping.json\"\n",
+        "\tmapping = json.load(mapping_path.open())\n",
+        "\n",
+        "\tdataset = InferenceDataset(data_dir)\n",
+        "\tdataloader = DataLoader(\n",
+        "\t\tdataset,\n",
+        "\t\tbatch_size=1,\n",
+        "\t\tshuffle=False,\n",
+        "\t\tdrop_last=False,\n",
+        "\t\tnum_workers=8,\n",
+        "\t\tcollate_fn=inference_collate_batch,\n",
+        "\t)\n",
+        "\tprint(f\"[Info]: Finish loading data!\",flush = True)\n",
+        "\n",
+        "\tspeaker_num = len(mapping[\"id2speaker\"])\n",
+        "\tmodel = Classifier(n_spks=speaker_num).to(device)\n",
+        "\tmodel.load_state_dict(torch.load(model_path))\n",
+        "\tmodel.eval()\n",
+        "\tprint(f\"[Info]: Finish creating model!\",flush = True)\n",
+        "\n",
+        "\tresults = [[\"Id\", \"Category\"]]\n",
+        "\tfor feat_paths, mels in tqdm(dataloader):\n",
+        "\t\twith torch.no_grad():\n",
+        "\t\t\tmels = mels.to(device)\n",
+        "\t\t\touts = model(mels)\n",
+        "\t\t\tpreds = outs.argmax(1).cpu().numpy()\n",
+        "\t\t\tfor feat_path, pred in zip(feat_paths, preds):\n",
+        "\t\t\t\tresults.append([feat_path, mapping[\"id2speaker\"][str(pred)]])\n",
+        "\n",
+        "\twith open(output_path, 'w', newline='') as csvfile:\n",
+        "\t\twriter = csv.writer(csvfile)\n",
+        "\t\twriter.writerows(results)\n",
+        "\n",
+        "\n",
+        "if __name__ == \"__main__\":\n",
+        "\tmain(**parse_args())"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": []
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "widgets": {
+      "application/vnd.jupyter.widget-state+json": {
+        "6786c2b0e2614ad389620246cb2178f2": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HBoxModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HBoxModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HBoxView",
+            "box_style": "",
+            "children": [
+              "IPY_MODEL_0d592098920140dab61aac5410568c36",
+              "IPY_MODEL_da681e3cc353420cb142d56df0fce231",
+              "IPY_MODEL_401ae722b95c4ff59b836422dbe71edc"
+            ],
+            "layout": "IPY_MODEL_4efbfb7c7cb54276862e5321209d57fa"
+          }
+        },
+        "0d592098920140dab61aac5410568c36": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_f6dcb3ec9c624171966bb889808dbcb3",
+            "placeholder": "​",
+            "style": "IPY_MODEL_100abf072991474abfcee871d2b83f37",
+            "value": "100%"
+          }
+        },
+        "da681e3cc353420cb142d56df0fce231": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "FloatProgressModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "FloatProgressModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "ProgressView",
+            "bar_style": "success",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_decc12da6f5742ec8b7ec7789ee434ab",
+            "max": 8000,
+            "min": 0,
+            "orientation": "horizontal",
+            "style": "IPY_MODEL_0117ac88c98440b29ab1f452107cbe1a",
+            "value": 8000
+          }
+        },
+        "401ae722b95c4ff59b836422dbe71edc": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "HTMLModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_dom_classes": [],
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "HTMLModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/controls",
+            "_view_module_version": "1.5.0",
+            "_view_name": "HTMLView",
+            "description": "",
+            "description_tooltip": null,
+            "layout": "IPY_MODEL_3d64942a0eaa409a93df049bb594c062",
+            "placeholder": "​",
+            "style": "IPY_MODEL_8c8721e504cf434eb9e263fb9983969a",
+            "value": " 8000/8000 [00:33&lt;00:00, 256.07it/s]"
+          }
+        },
+        "4efbfb7c7cb54276862e5321209d57fa": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "f6dcb3ec9c624171966bb889808dbcb3": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "100abf072991474abfcee871d2b83f37": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        },
+        "decc12da6f5742ec8b7ec7789ee434ab": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "0117ac88c98440b29ab1f452107cbe1a": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "ProgressStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "ProgressStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "bar_color": null,
+            "description_width": ""
+          }
+        },
+        "3d64942a0eaa409a93df049bb594c062": {
+          "model_module": "@jupyter-widgets/base",
+          "model_name": "LayoutModel",
+          "model_module_version": "1.2.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/base",
+            "_model_module_version": "1.2.0",
+            "_model_name": "LayoutModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "LayoutView",
+            "align_content": null,
+            "align_items": null,
+            "align_self": null,
+            "border": null,
+            "bottom": null,
+            "display": null,
+            "flex": null,
+            "flex_flow": null,
+            "grid_area": null,
+            "grid_auto_columns": null,
+            "grid_auto_flow": null,
+            "grid_auto_rows": null,
+            "grid_column": null,
+            "grid_gap": null,
+            "grid_row": null,
+            "grid_template_areas": null,
+            "grid_template_columns": null,
+            "grid_template_rows": null,
+            "height": null,
+            "justify_content": null,
+            "justify_items": null,
+            "left": null,
+            "margin": null,
+            "max_height": null,
+            "max_width": null,
+            "min_height": null,
+            "min_width": null,
+            "object_fit": null,
+            "object_position": null,
+            "order": null,
+            "overflow": null,
+            "overflow_x": null,
+            "overflow_y": null,
+            "padding": null,
+            "right": null,
+            "top": null,
+            "visibility": null,
+            "width": null
+          }
+        },
+        "8c8721e504cf434eb9e263fb9983969a": {
+          "model_module": "@jupyter-widgets/controls",
+          "model_name": "DescriptionStyleModel",
+          "model_module_version": "1.5.0",
+          "state": {
+            "_model_module": "@jupyter-widgets/controls",
+            "_model_module_version": "1.5.0",
+            "_model_name": "DescriptionStyleModel",
+            "_view_count": null,
+            "_view_module": "@jupyter-widgets/base",
+            "_view_module_version": "1.2.0",
+            "_view_name": "StyleView",
+            "description_width": ""
+          }
+        }
+      }
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}