Speech_Recognition_(PyTorch).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q1fXgsDQmK09"
      },
      "source": [
        "## installing the requirements"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gwfN8o17Bdp2"
      },
      "source": [
        "!pip install torchaudio==0.4.0 torch==1.4.0 comet-ml==3.0.2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tSKHvy8DmOCQ"
      },
      "source": [
        "## Setting up your data pipeline"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RVJs4Bk8FjjO"
      },
      "source": [
        "import os\n",
        "from comet_ml import Experiment\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.utils.data as data\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torchaudio\n",
        "import numpy as np\n",
        "\n",
        "def avg_wer(wer_scores, combined_ref_len):\n",
        "    return float(sum(wer_scores)) / float(combined_ref_len)\n",
        "\n",
        "\n",
        "def _levenshtein_distance(ref, hyp):\n",
        "    \"\"\"Levenshtein distance is a string metric for measuring the difference\n",
        "    between two sequences. Informally, the levenshtein disctance is defined as\n",
        "    the minimum number of single-character edits (substitutions, insertions or\n",
        "    deletions) required to change one word into the other. We can naturally\n",
        "    extend the edits to word level when calculate levenshtein disctance for\n",
        "    two sentences.\n",
        "    \"\"\"\n",
        "    m = len(ref)\n",
        "    n = len(hyp)\n",
        "\n",
        "    # special case\n",
        "    if ref == hyp:\n",
        "        return 0\n",
        "    if m == 0:\n",
        "        return n\n",
        "    if n == 0:\n",
        "        return m\n",
        "\n",
        "    if m < n:\n",
        "        ref, hyp = hyp, ref\n",
        "        m, n = n, m\n",
        "\n",
        "    # use O(min(m, n)) space\n",
        "    distance = np.zeros((2, n + 1), dtype=np.int32)\n",
        "\n",
        "    # initialize distance matrix\n",
        "    for j in range(0,n + 1):\n",
        "        distance[0][j] = j\n",
        "\n",
        "    # calculate levenshtein distance\n",
        "    for i in range(1, m + 1):\n",
        "        prev_row_idx = (i - 1) % 2\n",
        "        cur_row_idx = i % 2\n",
        "        distance[cur_row_idx][0] = i\n",
        "        for j in range(1, n + 1):\n",
        "            if ref[i - 1] == hyp[j - 1]:\n",
        "                distance[cur_row_idx][j] = distance[prev_row_idx][j - 1]\n",
        "            else:\n",
        "                s_num = distance[prev_row_idx][j - 1] + 1\n",
        "                i_num = distance[cur_row_idx][j - 1] + 1\n",
        "                d_num = distance[prev_row_idx][j] + 1\n",
        "                distance[cur_row_idx][j] = min(s_num, i_num, d_num)\n",
        "\n",
        "    return distance[m % 2][n]\n",
        "\n",
        "\n",
        "def word_errors(reference, hypothesis, ignore_case=False, delimiter=' '):\n",
        "    \"\"\"Compute the levenshtein distance between reference sequence and\n",
        "    hypothesis sequence in word-level.\n",
        "    :param reference: The reference sentence.\n",
        "    :type reference: basestring\n",
        "    :param hypothesis: The hypothesis sentence.\n",
        "    :type hypothesis: basestring\n",
        "    :param ignore_case: Whether case-sensitive or not.\n",
        "    :type ignore_case: bool\n",
        "    :param delimiter: Delimiter of input sentences.\n",
        "    :type delimiter: char\n",
        "    :return: Levenshtein distance and word number of reference sentence.\n",
        "    :rtype: list\n",
        "    \"\"\"\n",
        "    if ignore_case == True:\n",
        "        reference = reference.lower()\n",
        "        hypothesis = hypothesis.lower()\n",
        "\n",
        "    ref_words = reference.split(delimiter)\n",
        "    hyp_words = hypothesis.split(delimiter)\n",
        "\n",
        "    edit_distance = _levenshtein_distance(ref_words, hyp_words)\n",
        "    return float(edit_distance), len(ref_words)\n",
        "\n",
        "\n",
        "def char_errors(reference, hypothesis, ignore_case=False, remove_space=False):\n",
        "    \"\"\"Compute the levenshtein distance between reference sequence and\n",
        "    hypothesis sequence in char-level.\n",
        "    :param reference: The reference sentence.\n",
        "    :type reference: basestring\n",
        "    :param hypothesis: The hypothesis sentence.\n",
        "    :type hypothesis: basestring\n",
        "    :param ignore_case: Whether case-sensitive or not.\n",
        "    :type ignore_case: bool\n",
        "    :param remove_space: Whether remove internal space characters\n",
        "    :type remove_space: bool\n",
        "    :return: Levenshtein distance and length of reference sentence.\n",
        "    :rtype: list\n",
        "    \"\"\"\n",
        "    if ignore_case == True:\n",
        "        reference = reference.lower()\n",
        "        hypothesis = hypothesis.lower()\n",
        "\n",
        "    join_char = ' '\n",
        "    if remove_space == True:\n",
        "        join_char = ''\n",
        "\n",
        "    reference = join_char.join(filter(None, reference.split(' ')))\n",
        "    hypothesis = join_char.join(filter(None, hypothesis.split(' ')))\n",
        "\n",
        "    edit_distance = _levenshtein_distance(reference, hypothesis)\n",
        "    return float(edit_distance), len(reference)\n",
        "\n",
        "\n",
        "def wer(reference, hypothesis, ignore_case=False, delimiter=' '):\n",
        "    \"\"\"Calculate word error rate (WER). WER compares reference text and\n",
        "    hypothesis text in word-level. WER is defined as:\n",
        "    .. math::\n",
        "        WER = (Sw + Dw + Iw) / Nw\n",
        "    where\n",
        "    .. code-block:: text\n",
        "        Sw is the number of words subsituted,\n",
        "        Dw is the number of words deleted,\n",
        "        Iw is the number of words inserted,\n",
        "        Nw is the number of words in the reference\n",
        "    We can use levenshtein distance to calculate WER. Please draw an attention\n",
        "    that empty items will be removed when splitting sentences by delimiter.\n",
        "    :param reference: The reference sentence.\n",
        "    :type reference: basestring\n",
        "    :param hypothesis: The hypothesis sentence.\n",
        "    :type hypothesis: basestring\n",
        "    :param ignore_case: Whether case-sensitive or not.\n",
        "    :type ignore_case: bool\n",
        "    :param delimiter: Delimiter of input sentences.\n",
        "    :type delimiter: char\n",
        "    :return: Word error rate.\n",
        "    :rtype: float\n",
        "    :raises ValueError: If word number of reference is zero.\n",
        "    \"\"\"\n",
        "    edit_distance, ref_len = word_errors(reference, hypothesis, ignore_case,\n",
        "                                         delimiter)\n",
        "\n",
        "    if ref_len == 0:\n",
        "        raise ValueError(\"Reference's word number should be greater than 0.\")\n",
        "\n",
        "    wer = float(edit_distance) / ref_len\n",
        "    return wer\n",
        "\n",
        "\n",
        "def cer(reference, hypothesis, ignore_case=False, remove_space=False):\n",
        "    \"\"\"Calculate charactor error rate (CER). CER compares reference text and\n",
        "    hypothesis text in char-level. CER is defined as:\n",
        "    .. math::\n",
        "        CER = (Sc + Dc + Ic) / Nc\n",
        "    where\n",
        "    .. code-block:: text\n",
        "        Sc is the number of characters substituted,\n",
        "        Dc is the number of characters deleted,\n",
        "        Ic is the number of characters inserted\n",
        "        Nc is the number of characters in the reference\n",
        "    We can use levenshtein distance to calculate CER. Chinese input should be\n",
        "    encoded to unicode. Please draw an attention that the leading and tailing\n",
        "    space characters will be truncated and multiple consecutive space\n",
        "    characters in a sentence will be replaced by one space character.\n",
        "    :param reference: The reference sentence.\n",
        "    :type reference: basestring\n",
        "    :param hypothesis: The hypothesis sentence.\n",
        "    :type hypothesis: basestring\n",
        "    :param ignore_case: Whether case-sensitive or not.\n",
        "    :type ignore_case: bool\n",
        "    :param remove_space: Whether remove internal space characters\n",
        "    :type remove_space: bool\n",
        "    :return: Character error rate.\n",
        "    :rtype: float\n",
        "    :raises ValueError: If the reference length is zero.\n",
        "    \"\"\"\n",
        "    edit_distance, ref_len = char_errors(reference, hypothesis, ignore_case,\n",
        "                                         remove_space)\n",
        "\n",
        "    if ref_len == 0:\n",
        "        raise ValueError(\"Length of reference should be greater than 0.\")\n",
        "\n",
        "    cer = float(edit_distance) / ref_len\n",
        "    return cer\n",
        "\n",
        "class TextTransform:\n",
        "    \"\"\"Maps characters to integers and vice versa\"\"\"\n",
        "    def __init__(self):\n",
        "        char_map_str = \"\"\"\n",
        "        ' 0\n",
        "        <SPACE> 1\n",
        "        a 2\n",
        "        b 3\n",
        "        c 4\n",
        "        d 5\n",
        "        e 6\n",
        "        f 7\n",
        "        g 8\n",
        "        h 9\n",
        "        i 10\n",
        "        j 11\n",
        "        k 12\n",
        "        l 13\n",
        "        m 14\n",
        "        n 15\n",
        "        o 16\n",
        "        p 17\n",
        "        q 18\n",
        "        r 19\n",
        "        s 20\n",
        "        t 21\n",
        "        u 22\n",
        "        v 23\n",
        "        w 24\n",
        "        x 25\n",
        "        y 26\n",
        "        z 27\n",
        "        \"\"\"\n",
        "        self.char_map = {}\n",
        "        self.index_map = {}\n",
        "        for line in char_map_str.strip().split('\\n'):\n",
        "            ch, index = line.split()\n",
        "            self.char_map[ch] = int(index)\n",
        "            self.index_map[int(index)] = ch\n",
        "        self.index_map[1] = ' '\n",
        "\n",
        "    def text_to_int(self, text):\n",
        "        \"\"\" Use a character map and convert text to an integer sequence \"\"\"\n",
        "        int_sequence = []\n",
        "        for c in text:\n",
        "            if c == ' ':\n",
        "                ch = self.char_map['<SPACE>']\n",
        "            else:\n",
        "                ch = self.char_map[c]\n",
        "            int_sequence.append(ch)\n",
        "        return int_sequence\n",
        "\n",
        "    def int_to_text(self, labels):\n",
        "        \"\"\" Use a character map and convert integer labels to an text sequence \"\"\"\n",
        "        string = []\n",
        "        for i in labels:\n",
        "            string.append(self.index_map[i])\n",
        "        return ''.join(string).replace('<SPACE>', ' ')\n",
        "\n",
        "train_audio_transforms = nn.Sequential(\n",
        "    torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_mels=128),\n",
        "    torchaudio.transforms.FrequencyMasking(freq_mask_param=30),\n",
        "    torchaudio.transforms.TimeMasking(time_mask_param=100)\n",
        ")\n",
        "\n",
        "valid_audio_transforms = torchaudio.transforms.MelSpectrogram()\n",
        "\n",
        "text_transform = TextTransform()\n",
        "\n",
        "def data_processing(data, data_type=\"train\"):\n",
        "    spectrograms = []\n",
        "    labels = []\n",
        "    input_lengths = []\n",
        "    label_lengths = []\n",
        "    for (waveform, _, utterance, _, _, _) in data:\n",
        "        if data_type == 'train':\n",
        "            spec = train_audio_transforms(waveform).squeeze(0).transpose(0, 1)\n",
        "        elif data_type == 'valid':\n",
        "            spec = valid_audio_transforms(waveform).squeeze(0).transpose(0, 1)\n",
        "        else:\n",
        "            raise Exception('data_type should be train or valid')\n",
        "        spectrograms.append(spec)\n",
        "        label = torch.Tensor(text_transform.text_to_int(utterance.lower()))\n",
        "        labels.append(label)\n",
        "        input_lengths.append(spec.shape[0]//2)\n",
        "        label_lengths.append(len(label))\n",
        "\n",
        "    spectrograms = nn.utils.rnn.pad_sequence(spectrograms, batch_first=True).unsqueeze(1).transpose(2, 3)\n",
        "    labels = nn.utils.rnn.pad_sequence(labels, batch_first=True)\n",
        "\n",
        "    return spectrograms, labels, input_lengths, label_lengths\n",
        "\n",
        "\n",
        "def GreedyDecoder(output, labels, label_lengths, blank_label=28, collapse_repeated=True):\n",
        "\targ_maxes = torch.argmax(output, dim=2)\n",
        "\tdecodes = []\n",
        "\ttargets = []\n",
        "\tfor i, args in enumerate(arg_maxes):\n",
        "\t\tdecode = []\n",
        "\t\ttargets.append(text_transform.int_to_text(labels[i][:label_lengths[i]].tolist()))\n",
        "\t\tfor j, index in enumerate(args):\n",
        "\t\t\tif index != blank_label:\n",
        "\t\t\t\tif collapse_repeated and j != 0 and index == args[j -1]:\n",
        "\t\t\t\t\tcontinue\n",
        "\t\t\t\tdecode.append(index.item())\n",
        "\t\tdecodes.append(text_transform.int_to_text(decode))\n",
        "\treturn decodes, targets\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4XdSlhAQnDEA"
      },
      "source": [
        "## The Model\n",
        "Base of of Deep Speech 2 with some personal improvements"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "65H1-PCjm-FB"
      },
      "source": [
        "class CNNLayerNorm(nn.Module):\n",
        "    \"\"\"Layer normalization built for cnns input\"\"\"\n",
        "    def __init__(self, n_feats):\n",
        "        super(CNNLayerNorm, self).__init__()\n",
        "        self.layer_norm = nn.LayerNorm(n_feats)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # x (batch, channel, feature, time)\n",
        "        x = x.transpose(2, 3).contiguous() # (batch, channel, time, feature)\n",
        "        x = self.layer_norm(x)\n",
        "        return x.transpose(2, 3).contiguous() # (batch, channel, feature, time) \n",
        "\n",
        "\n",
        "class ResidualCNN(nn.Module):\n",
        "    \"\"\"Residual CNN inspired by https://arxiv.org/pdf/1603.05027.pdf\n",
        "        except with layer norm instead of batch norm\n",
        "    \"\"\"\n",
        "    def __init__(self, in_channels, out_channels, kernel, stride, dropout, n_feats):\n",
        "        super(ResidualCNN, self).__init__()\n",
        "\n",
        "        self.cnn1 = nn.Conv2d(in_channels, out_channels, kernel, stride, padding=kernel//2)\n",
        "        self.cnn2 = nn.Conv2d(out_channels, out_channels, kernel, stride, padding=kernel//2)\n",
        "        self.dropout1 = nn.Dropout(dropout)\n",
        "        self.dropout2 = nn.Dropout(dropout)\n",
        "        self.layer_norm1 = CNNLayerNorm(n_feats)\n",
        "        self.layer_norm2 = CNNLayerNorm(n_feats)\n",
        "\n",
        "    def forward(self, x):\n",
        "        residual = x  # (batch, channel, feature, time)\n",
        "        x = self.layer_norm1(x)\n",
        "        x = F.gelu(x)\n",
        "        x = self.dropout1(x)\n",
        "        x = self.cnn1(x)\n",
        "        x = self.layer_norm2(x)\n",
        "        x = F.gelu(x)\n",
        "        x = self.dropout2(x)\n",
        "        x = self.cnn2(x)\n",
        "        x += residual\n",
        "        return x # (batch, channel, feature, time)\n",
        "\n",
        "\n",
        "class BidirectionalGRU(nn.Module):\n",
        "\n",
        "    def __init__(self, rnn_dim, hidden_size, dropout, batch_first):\n",
        "        super(BidirectionalGRU, self).__init__()\n",
        "\n",
        "        self.BiGRU = nn.GRU(\n",
        "            input_size=rnn_dim, hidden_size=hidden_size,\n",
        "            num_layers=1, batch_first=batch_first, bidirectional=True)\n",
        "        self.layer_norm = nn.LayerNorm(rnn_dim)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.layer_norm(x)\n",
        "        x = F.gelu(x)\n",
        "        x, _ = self.BiGRU(x)\n",
        "        x = self.dropout(x)\n",
        "        return x\n",
        "\n",
        "\n",
        "class SpeechRecognitionModel(nn.Module):\n",
        "    \n",
        "    def __init__(self, n_cnn_layers, n_rnn_layers, rnn_dim, n_class, n_feats, stride=2, dropout=0.1):\n",
        "        super(SpeechRecognitionModel, self).__init__()\n",
        "        n_feats = n_feats//2\n",
        "        self.cnn = nn.Conv2d(1, 32, 3, stride=stride, padding=3//2)  # cnn for extracting heirachal features\n",
        "\n",
        "        # n residual cnn layers with filter size of 32\n",
        "        self.rescnn_layers = nn.Sequential(*[\n",
        "            ResidualCNN(32, 32, kernel=3, stride=1, dropout=dropout, n_feats=n_feats) \n",
        "            for _ in range(n_cnn_layers)\n",
        "        ])\n",
        "        self.fully_connected = nn.Linear(n_feats*32, rnn_dim)\n",
        "        self.birnn_layers = nn.Sequential(*[\n",
        "            BidirectionalGRU(rnn_dim=rnn_dim if i==0 else rnn_dim*2,\n",
        "                             hidden_size=rnn_dim, dropout=dropout, batch_first=i==0)\n",
        "            for i in range(n_rnn_layers)\n",
        "        ])\n",
        "        self.classifier = nn.Sequential(\n",
        "            nn.Linear(rnn_dim*2, rnn_dim),  # birnn returns rnn_dim*2\n",
        "            nn.GELU(),\n",
        "            nn.Dropout(dropout),\n",
        "            nn.Linear(rnn_dim, n_class)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.cnn(x)\n",
        "        x = self.rescnn_layers(x)\n",
        "        sizes = x.size()\n",
        "        x = x.view(sizes[0], sizes[1] * sizes[2], sizes[3])  # (batch, feature, time)\n",
        "        x = x.transpose(1, 2) # (batch, time, feature)\n",
        "        x = self.fully_connected(x)\n",
        "        x = self.birnn_layers(x)\n",
        "        x = self.classifier(x)\n",
        "        return x\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CuguNEzKnMOn"
      },
      "source": [
        "## The Training and Evaluating Script"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ydkqGeOwnPGY"
      },
      "source": [
        "class IterMeter(object):\n",
        "    \"\"\"keeps track of total iterations\"\"\"\n",
        "    def __init__(self):\n",
        "        self.val = 0\n",
        "\n",
        "    def step(self):\n",
        "        self.val += 1\n",
        "\n",
        "    def get(self):\n",
        "        return self.val\n",
        "\n",
        "\n",
        "def train(model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, experiment):\n",
        "    model.train()\n",
        "    data_len = len(train_loader.dataset)\n",
        "    with experiment.train():\n",
        "        for batch_idx, _data in enumerate(train_loader):\n",
        "            spectrograms, labels, input_lengths, label_lengths = _data \n",
        "            spectrograms, labels = spectrograms.to(device), labels.to(device)\n",
        "\n",
        "            optimizer.zero_grad()\n",
        "\n",
        "            output = model(spectrograms)  # (batch, time, n_class)\n",
        "            output = F.log_softmax(output, dim=2)\n",
        "            output = output.transpose(0, 1) # (time, batch, n_class)\n",
        "\n",
        "            loss = criterion(output, labels, input_lengths, label_lengths)\n",
        "            loss.backward()\n",
        "\n",
        "            experiment.log_metric('loss', loss.item(), step=iter_meter.get())\n",
        "            experiment.log_metric('learning_rate', scheduler.get_lr(), step=iter_meter.get())\n",
        "\n",
        "            optimizer.step()\n",
        "            scheduler.step()\n",
        "            iter_meter.step()\n",
        "            if batch_idx % 100 == 0 or batch_idx == data_len:\n",
        "                print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
        "                    epoch, batch_idx * len(spectrograms), data_len,\n",
        "                    100. * batch_idx / len(train_loader), loss.item()))\n",
        "\n",
        "\n",
        "def test(model, device, test_loader, criterion, epoch, iter_meter, experiment):\n",
        "    print('\\nevaluating...')\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    test_cer, test_wer = [], []\n",
        "    with experiment.test():\n",
        "        with torch.no_grad():\n",
        "            for i, _data in enumerate(test_loader):\n",
        "                spectrograms, labels, input_lengths, label_lengths = _data \n",
        "                spectrograms, labels = spectrograms.to(device), labels.to(device)\n",
        "\n",
        "                output = model(spectrograms)  # (batch, time, n_class)\n",
        "                output = F.log_softmax(output, dim=2)\n",
        "                output = output.transpose(0, 1) # (time, batch, n_class)\n",
        "\n",
        "                loss = criterion(output, labels, input_lengths, label_lengths)\n",
        "                test_loss += loss.item() / len(test_loader)\n",
        "\n",
        "                decoded_preds, decoded_targets = GreedyDecoder(output.transpose(0, 1), labels, label_lengths)\n",
        "                for j in range(len(decoded_preds)):\n",
        "                    test_cer.append(cer(decoded_targets[j], decoded_preds[j]))\n",
        "                    test_wer.append(wer(decoded_targets[j], decoded_preds[j]))\n",
        "\n",
        "\n",
        "    avg_cer = sum(test_cer)/len(test_cer)\n",
        "    avg_wer = sum(test_wer)/len(test_wer)\n",
        "    experiment.log_metric('test_loss', test_loss, step=iter_meter.get())\n",
        "    experiment.log_metric('cer', avg_cer, step=iter_meter.get())\n",
        "    experiment.log_metric('wer', avg_wer, step=iter_meter.get())\n",
        "\n",
        "    print('Test set: Average loss: {:.4f}, Average CER: {:4f} Average WER: {:.4f}\\n'.format(test_loss, avg_cer, avg_wer))\n",
        "\n",
        "\n",
        "def main(learning_rate=5e-4, batch_size=20, epochs=10,\n",
        "        train_url=\"train-clean-100\", test_url=\"test-clean\",\n",
        "        experiment=Experiment(api_key='dummy_key', disabled=True)):\n",
        "\n",
        "    hparams = {\n",
        "        \"n_cnn_layers\": 3,\n",
        "        \"n_rnn_layers\": 5,\n",
        "        \"rnn_dim\": 512,\n",
        "        \"n_class\": 29,\n",
        "        \"n_feats\": 128,\n",
        "        \"stride\":2,\n",
        "        \"dropout\": 0.1,\n",
        "        \"learning_rate\": learning_rate,\n",
        "        \"batch_size\": batch_size,\n",
        "        \"epochs\": epochs\n",
        "    }\n",
        "\n",
        "    experiment.log_parameters(hparams)\n",
        "\n",
        "    use_cuda = torch.cuda.is_available()\n",
        "    torch.manual_seed(7)\n",
        "    device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
        "\n",
        "    if not os.path.isdir(\"./data\"):\n",
        "        os.makedirs(\"./data\")\n",
        "\n",
        "    train_dataset = torchaudio.datasets.LIBRISPEECH(\"./data\", url=train_url, download=True)\n",
        "    test_dataset = torchaudio.datasets.LIBRISPEECH(\"./data\", url=test_url, download=True)\n",
        "\n",
        "    kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}\n",
        "    train_loader = data.DataLoader(dataset=train_dataset,\n",
        "                                batch_size=hparams['batch_size'],\n",
        "                                shuffle=True,\n",
        "                                collate_fn=lambda x: data_processing(x, 'train'),\n",
        "                                **kwargs)\n",
        "    test_loader = data.DataLoader(dataset=test_dataset,\n",
        "                                batch_size=hparams['batch_size'],\n",
        "                                shuffle=False,\n",
        "                                collate_fn=lambda x: data_processing(x, 'valid'),\n",
        "                                **kwargs)\n",
        "\n",
        "    model = SpeechRecognitionModel(\n",
        "        hparams['n_cnn_layers'], hparams['n_rnn_layers'], hparams['rnn_dim'],\n",
        "        hparams['n_class'], hparams['n_feats'], hparams['stride'], hparams['dropout']\n",
        "        ).to(device)\n",
        "\n",
        "    print(model)\n",
        "    print('Num Model Parameters', sum([param.nelement() for param in model.parameters()]))\n",
        "\n",
        "    optimizer = optim.AdamW(model.parameters(), hparams['learning_rate'])\n",
        "    criterion = nn.CTCLoss(blank=28).to(device)\n",
        "    scheduler = optim.lr_scheduler.OneCycleLR(optimizer, max_lr=hparams['learning_rate'], \n",
        "                                            steps_per_epoch=int(len(train_loader)),\n",
        "                                            epochs=hparams['epochs'],\n",
        "                                            anneal_strategy='linear')\n",
        "    \n",
        "    iter_meter = IterMeter()\n",
        "    for epoch in range(1, epochs + 1):\n",
        "        train(model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, experiment)\n",
        "        test(model, device, test_loader, criterion, epoch, iter_meter, experiment)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4qBGdkQSmW3a"
      },
      "source": [
        "## Setting up Comet"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "edo8shRBFt4V"
      },
      "source": [
        "comet_api_key = \"\" # api key here\n",
        "project_name = \"speechrecognition\"\n",
        "experiment_name = \"speechrecognition-colab\"\n",
        "\n",
        "if comet_api_key:\n",
        "  experiment = Experiment(api_key=comet_api_key, project_name=project_name, parse_args=False)\n",
        "  experiment.set_name(experiment_name)\n",
        "  experiment.display()\n",
        "else:\n",
        "  experiment = Experiment(api_key='dummy_key', disabled=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HxRIb_WempDq"
      },
      "source": [
        "## GPU runtime\n",
        "Adjust your batch_size depending on which GPU you use"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nlUSuAJwlzo8"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hXvlWZeVpXfX"
      },
      "source": [
        "## Train\n",
        "this will download the data on first run(takes little time).\n",
        "\\n logs can be viewed in comet. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XZodve8PGKfS"
      },
      "source": [
        "\n",
        "\n",
        "learning_rate = 5e-4  \n",
        "batch_size = 10    #can change according to GPU type and memory\n",
        "epochs = 10\n",
        "libri_train_set = \"train-clean-100\"   \n",
        "libri_test_set = \"test-clean\"\n",
        "\n",
        "main(learning_rate, batch_size, epochs, libri_train_set, libri_test_set, experiment)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "92kVVEr7GR6j"
      },
      "source": [
        "experiment.end()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ucfQX3qN21az"
      },
      "source": [],
      "execution_count": null,
      "outputs": []
    }
  ]
}