train_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training Our Custom Models 🏋️\n",
    "- in this notebook we train and test our custom implementations of Deep Speech 2, Jasper, and Conformer\n",
    "- the models take as input Mel Filterbanks (2D images) and output probability distribution over characters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1/5 Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "# Custom libraries\n",
    "from utils.preprocessing import Preprocessing\n",
    "from utils.word_model import WordModel\n",
    "from utils.decoders import DecoderGreedy\n",
    "from utils.metrics import avg_cer, avg_wer\n",
    "from utils.misc import pretty_params\n",
    "\n",
    "# Plots\n",
    "import wandb\n",
    "\n",
    "# Models\n",
    "from models.deep_speech.deep_speech_base import DeepSpeechBase\n",
    "from models.deep_speech.deep_speech_attention import DeepSpeechAttention\n",
    "from models.jasper.jasper_base import Jasper\n",
    "from models.jasper.jasper_dr import JasperDR\n",
    "from models.conformer.conformer import Conformer\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Notes**\n",
    "\n",
    "Deep Speech\n",
    "- to use unigrams: set `WordModel(\"unigram\")` and `stride=2`\n",
    "- to use bigrams: set `WordModel(\"bigrams\")` and `stride=4`\n",
    "\n",
    "Jasper & Conformer\n",
    "- always set `WordModel(\"unigram\")` and `stride=2` "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# General Hyper-params\n",
    "seed = 42\n",
    "batch_size = 16\n",
    "epochs = 10\n",
    "\n",
    "n_features = 128            # freq axis\n",
    "stride = 2                  # time-axis striding\n",
    "\n",
    "word_model = WordModel(\"unigram\")\n",
    "lr = 0.0005\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda\n"
     ]
    }
   ],
   "source": [
    "use_cuda = torch.cuda.is_available()\n",
    "torch.manual_seed(seed)\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2/5 Data processing\n",
    "- set `download=True` if you are downloading for the first time\n",
    "- pick the train dataset according to your hardware constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set_name = \"train-clean-100\" # [6.3G] 100 hours \n",
    "#train_set_name = \"train-clean-360\" # [23G] 360 hours\n",
    "#train_set_name = \"train-clean-500\" # [30G] 500 hours"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\marco\\.conda\\envs\\exp\\Lib\\site-packages\\torchaudio\\functional\\functional.py:584: UserWarning: At least one mel filterbank has all zero values. The value for `n_mels` (128) may be set too high. Or, the value for `n_freqs` (201) may be set too low.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "prep = Preprocessing()\n",
    "\n",
    "train_clean = prep.download(split=train_set_name, download=False)\n",
    "\n",
    "dev_clean = prep.download(split='dev-clean', download=False)\n",
    "dev_other = prep.download(split='dev-other', download=False)\n",
    "\n",
    "test_clean = prep.download(split='test-clean', download=False)\n",
    "test_other = prep.download(split='test-other', download=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = DataLoader(dataset=train_clean,\n",
    "                          batch_size=batch_size,\n",
    "                          shuffle=True,\n",
    "                          collate_fn=lambda x: prep.preprocess(x, train_set_name, stride, word_model))\n",
    "\n",
    "dev_clean_loader = DataLoader(dataset=dev_clean,\n",
    "                               batch_size=batch_size,\n",
    "                               shuffle=False,\n",
    "                               collate_fn=lambda x: prep.preprocess(x, \"dev-clean\", stride, word_model))\n",
    "\n",
    "dev_other_loader = DataLoader(dataset=dev_other,\n",
    "                               batch_size=batch_size,\n",
    "                               shuffle=False,\n",
    "                               collate_fn=lambda x: prep.preprocess(x, \"dev-other\", stride, word_model))\n",
    "\n",
    "test_clean_loader = DataLoader(dataset=test_clean,\n",
    "                               batch_size=batch_size,\n",
    "                               shuffle=False,\n",
    "                               collate_fn=lambda x: prep.preprocess(x, \"test-clean\", stride, word_model))\n",
    "\n",
    "test_other_loader = DataLoader(dataset=test_other,\n",
    "                               batch_size=batch_size,\n",
    "                               shuffle=False,\n",
    "                               collate_fn=lambda x: prep.preprocess(x, \"test-other\", stride, word_model))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "~ TRAIN \t28539 samples ~\n",
      "~ DEV CLEAN \t2703 samples ~\n",
      "~ DEV OTHER \t2864 samples ~\n",
      "~ TEST CLEAN \t2620 samples ~\n",
      "~ TEST OTHER \t2939 samples ~\n",
      "\n",
      "+------ Dataloader length: 28539 ------+\n",
      "# Batches: 1784\n",
      "Spectrogram shape: [16, 1, 128, 1308]\n",
      "Label shape: [16, 257]\n",
      "Mel length (length of each spectrogram): [577, 654, 300, 484, 542, 584] ...\n",
      "Idx length (length of each label): [197, 257, 121, 173, 196, 205] ...\n",
      "+------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "## Sanity check ##\n",
    "\n",
    "print(f\"~ TRAIN \\t{len(train_loader.dataset)} samples ~\")\n",
    "print(f\"~ DEV CLEAN \\t{len(dev_clean_loader.dataset)} samples ~\")\n",
    "print(f\"~ DEV OTHER \\t{len(dev_other_loader.dataset)} samples ~\")\n",
    "print(f\"~ TEST CLEAN \\t{len(test_clean_loader.dataset)} samples ~\")\n",
    "print(f\"~ TEST OTHER \\t{len(test_other_loader.dataset)} samples ~\")\n",
    "print()\n",
    "\n",
    "prep.print_loader_info(train_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3/5 Model\n",
    "Available models:\n",
    "\n",
    "- DeepSpeech-Base\n",
    "- DeepSpeech-Attention\n",
    "- Jasper-Base\n",
    "- Conformer-Small"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 23.71M (23705373)\n"
     ]
    }
   ],
   "source": [
    "# DeepSpeechBase\n",
    "deep_speech = DeepSpeechBase(n_cnn=3,\n",
    "                             n_rnn=5, \n",
    "                             rnn_dim=512, \n",
    "                             n_features=n_features, \n",
    "                             n_class=word_model.get_n_class(),\n",
    "                             stride=stride,\n",
    "                             drop_rate=0.2).to(device)\n",
    "\n",
    "tot_params = sum([p.numel() for p in deep_speech.parameters()])\n",
    "model_name=\"deep_speech_base\"\n",
    "print(f\"Number of parameters: {pretty_params(tot_params)} ({tot_params})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 20.32M (20319005)\n"
     ]
    }
   ],
   "source": [
    "# DeepSpeechAttention\n",
    "deep_speech = DeepSpeechAttention(n_cnn=3,\n",
    "                                  n_enc=12,\n",
    "                                  n_features=n_features,\n",
    "                                  n_class=word_model.get_n_class(),\n",
    "                                  emb_dim=512,\n",
    "                                  n_heads=4,\n",
    "                                  stride=stride,\n",
    "                                  drop_rate=0.2).to(device)\n",
    "\n",
    "tot_params = sum([p.numel() for p in deep_speech.parameters()])\n",
    "model_name=\"deep_speech_attention\"\n",
    "print(f\"Number of parameters: {pretty_params(tot_params)} ({tot_params})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 107.87M (107873693)\n"
     ]
    }
   ],
   "source": [
    "# JasperBase\n",
    "jasper = Jasper().to(device)\n",
    "\n",
    "tot_params = sum([p.numel() for p in jasper.parameters()])\n",
    "model_name=\"jasper_base\"\n",
    "print(f\"Number of parameters: {pretty_params(tot_params)} ({tot_params})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 109.91M (109908125)\n"
     ]
    }
   ],
   "source": [
    "# Jasper DR (Dense Residual)\n",
    "jasper = JasperDR().to(device)\n",
    "\n",
    "tot_params = sum([p.numel() for p in jasper.parameters()])\n",
    "model_name=\"jasper_DR\"\n",
    "print(f\"Number of parameters: {pretty_params(tot_params)} ({tot_params})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of parameters: 27.36M (27362525)\n"
     ]
    }
   ],
   "source": [
    "# Conformer Small\n",
    "conformer = Conformer(in_features=128,\n",
    "                      encoder_dim=256,\n",
    "                      num_heads=4,\n",
    "                      kernel_size=31,\n",
    "                      hidden_size=320,\n",
    "                      n_class=29,\n",
    "                      n_blocks=16,\n",
    "                      device=device).to(device)\n",
    "\n",
    "tot_params = sum([p.numel() for p in conformer.parameters()])\n",
    "model_name=\"Conformer-Small\"\n",
    "print(f\"Number of parameters: {pretty_params(tot_params)} ({tot_params})\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = DecoderGreedy(word_model.get_blank_id())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4/5 Optimizer, loss, scheduler\n",
    "- set here the model you want to train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "#model_to_train = deep_speech\n",
    "model_to_train = jasper\n",
    "#model_to_train = conformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Optionally, load model's weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#path_to_weights = \"./weights/jasper_base_lr_variable.pth\"\n",
    "path_to_weights = \"./weights/3invres_12attention.pth\"\n",
    "model_to_train.load_state_dict(torch.load(path_to_weights))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change model\n",
    "adamW = optim.AdamW(model_to_train.parameters(), lr)\n",
    "ctc_loss = nn.CTCLoss(blank=word_model.get_blank_id()).to(device)\n",
    "one_cycle_lr = optim.lr_scheduler.OneCycleLR(adamW,\n",
    "                                             max_lr=lr,\n",
    "                                             steps_per_epoch=int(len(train_loader)),\n",
    "                                             epochs=epochs,\n",
    "                                             anneal_strategy=\"linear\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5/5 Train & Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Online fancy plots\n",
    "! wandb login\n",
    "\n",
    "wandb.init(\n",
    "    project=\"asr_librispeech\",\n",
    "\n",
    "    config= {\n",
    "        \"model\": model_name,\n",
    "        \"word_model\": word_model.get_name()\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Training info\n",
    "- After the forward pass, output must be  `[batch, time, class]`\n",
    "- CTC loss expects predictions to be `[seq, batch, class]`\n",
    "- Train your model for `epochs` number of epochs\n",
    "- At the end of each epoch, get WER/CER on validation/test dataset\n",
    "\n",
    "Convention on Input/Output dimensions; all models must respect them\n",
    "- `input [batch_size, 1, seq_len, num_features]`\n",
    "- `output [batch_size, seq_len, num_class]`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch, dataset_loader, model, optimizer, scheduler, fn_loss):\n",
    "    print(f\"Traininig... (e={epoch})\")\n",
    "    \n",
    "    # Train mode ON\n",
    "    model.train()\n",
    "    n_samples = int(len(dataset_loader.dataset))\n",
    "\n",
    "    for idx, audio_data in enumerate(dataset_loader):\n",
    "        \n",
    "        # Get audio data with shape [batch, 1, n_features, seq_len]\n",
    "        spectrograms, indices, len_spectrograms, len_indices = audio_data\n",
    "        spectrograms, indices = spectrograms.to(device), indices.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Forward pass\n",
    "        out = model(spectrograms)\n",
    "        out = F.log_softmax(out, dim=2)\n",
    "        out = out.transpose(0, 1)\n",
    "        \n",
    "        # Backward pass\n",
    "        loss = fn_loss(out, indices, len_spectrograms, len_indices)\n",
    "        loss.backward()\n",
    "\n",
    "        # Log\n",
    "        # wandb.log({\n",
    "        #    \"loss\": loss.item(),\n",
    "        #    \"lr\": scheduler.get_last_lr()[0]\n",
    "        # })\n",
    "\n",
    "        # Step\n",
    "        optimizer.step()\n",
    "        scheduler.step()\n",
    "\n",
    "        # Log\n",
    "        if idx % 20 == 0 or idx == n_samples:\n",
    "            print(\"Epoch: {}, [{}/{}], Loss: {:.6f}\".format(\n",
    "                epoch, \n",
    "                idx*len(spectrograms), \n",
    "                n_samples,\n",
    "                loss.item()))\n",
    "\n",
    "    wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(epoch, dataset_name, dataset_loader, model, optimizer, fn_loss, debug=False):\n",
    "    print(f\"Testing on {dataset_name} (epoch={epoch})\")\n",
    "    model.eval()\n",
    "\n",
    "    total_loss = 0\n",
    "    wer_list = []\n",
    "    cer_list = []\n",
    "\n",
    "    n_batch = int(len(dataset_loader))\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for idx, audio_data in enumerate(dataset_loader):\n",
    "        \n",
    "            # Get audio data\n",
    "            spectrograms, indices, len_spectrograms, len_indices = audio_data\n",
    "            spectrograms, indices = spectrograms.to(device), indices.to(device)\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            \n",
    "            # Forward pass\n",
    "            out = model(spectrograms)\n",
    "            out = F.log_softmax(out, dim=2)\n",
    "            out = out.transpose(0, 1)\n",
    "\n",
    "            # Compute loss\n",
    "            loss = fn_loss(out, indices, len_spectrograms, len_indices)\n",
    "            total_loss += loss.item() / n_batch\n",
    "\n",
    "            # Metrics\n",
    "            decode_hypothesis = decoder.decode_prob(out, word_model)\n",
    "            decode_reference = decoder.decode_labels(indices, len_indices, word_model)\n",
    "\n",
    "            wer_list.append(avg_wer(decode_hypothesis, decode_reference))\n",
    "            cer_list.append(avg_cer(decode_hypothesis, decode_reference))\n",
    "            \n",
    "            if idx % 20 == 0:\n",
    "                print(f'Idx: {idx}')\n",
    "                print(f'reference: {decode_reference[0]}')\n",
    "                print(f'hypothesis: {decode_hypothesis[0]}')\n",
    "                print(f\"WER: {wer_list[-1]:.4f}, CER: {cer_list[-1]:.4f}\")\n",
    "                print()\n",
    "\n",
    "            if debug: break\n",
    "            \n",
    "    print(f\"Loss: {total_loss:.6f}\")\n",
    "    print(f\"WER: {sum(wer_list)/len(wer_list):.4f}\")\n",
    "    print(f\"CER: {sum(cer_list)/len(cer_list):.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Traininig... (e=1)\n",
      "Epoch: 1, [0/28539], Loss: 6.652934\n",
      "Epoch: 1, [320/28539], Loss: 4.094671\n",
      "Epoch: 1, [640/28539], Loss: 3.571755\n",
      "Epoch: 1, [960/28539], Loss: 3.351143\n",
      "Epoch: 1, [1280/28539], Loss: 3.229725\n",
      "Epoch: 1, [1600/28539], Loss: 3.138221\n",
      "Epoch: 1, [1920/28539], Loss: 3.070983\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[13], line 2\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m epoch \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m1\u001b[39m, epochs\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m):\n\u001b[1;32m----> 2\u001b[0m     \u001b[43mtrain\u001b[49m\u001b[43m(\u001b[49m\u001b[43mepoch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtrain_loader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmodel_to_train\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43madamW\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mone_cycle_lr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mctc_loss\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m      3\u001b[0m     test(epoch, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mdev-clean\u001b[39m\u001b[38;5;124m\"\u001b[39m, dev_clean_loader, model_to_train, adamW, ctc_loss)\n",
      "Cell \u001b[1;32mIn[12], line 22\u001b[0m, in \u001b[0;36mtrain\u001b[1;34m(epoch, dataset_loader, model, optimizer, scheduler, fn_loss)\u001b[0m\n\u001b[0;32m     19\u001b[0m out \u001b[38;5;241m=\u001b[39m out\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m     21\u001b[0m \u001b[38;5;66;03m# Backward pass\u001b[39;00m\n\u001b[1;32m---> 22\u001b[0m loss \u001b[38;5;241m=\u001b[39m \u001b[43mfn_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[43mout\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mindices\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlen_spectrograms\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mlen_indices\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     23\u001b[0m loss\u001b[38;5;241m.\u001b[39mbackward()\n\u001b[0;32m     25\u001b[0m \u001b[38;5;66;03m# Log\u001b[39;00m\n\u001b[0;32m     26\u001b[0m \u001b[38;5;66;03m# wandb.log({\u001b[39;00m\n\u001b[0;32m     27\u001b[0m \u001b[38;5;66;03m#    \"loss\": loss.item(),\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     30\u001b[0m \n\u001b[0;32m     31\u001b[0m \u001b[38;5;66;03m# Step\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\marco\\.conda\\envs\\exp\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[0;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\marco\\.conda\\envs\\exp\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[0;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[0;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[0;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[0;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[1;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
      "File \u001b[1;32mc:\\Users\\marco\\.conda\\envs\\exp\\Lib\\site-packages\\torch\\nn\\modules\\loss.py:1779\u001b[0m, in \u001b[0;36mCTCLoss.forward\u001b[1;34m(self, log_probs, targets, input_lengths, target_lengths)\u001b[0m\n\u001b[0;32m   1778\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, log_probs: Tensor, targets: Tensor, input_lengths: Tensor, target_lengths: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[1;32m-> 1779\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mF\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mctc_loss\u001b[49m\u001b[43m(\u001b[49m\u001b[43mlog_probs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtargets\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minput_lengths\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget_lengths\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mblank\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mreduction\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m   1780\u001b[0m \u001b[43m                      \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mzero_infinity\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mc:\\Users\\marco\\.conda\\envs\\exp\\Lib\\site-packages\\torch\\nn\\functional.py:2660\u001b[0m, in \u001b[0;36mctc_loss\u001b[1;34m(log_probs, targets, input_lengths, target_lengths, blank, reduction, zero_infinity)\u001b[0m\n\u001b[0;32m   2653\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_torch_function_variadic(log_probs, targets, input_lengths, target_lengths):\n\u001b[0;32m   2654\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(\n\u001b[0;32m   2655\u001b[0m         ctc_loss,\n\u001b[0;32m   2656\u001b[0m         (log_probs, targets, input_lengths, target_lengths),\n\u001b[0;32m   2657\u001b[0m         log_probs, targets, input_lengths, target_lengths,\n\u001b[0;32m   2658\u001b[0m         blank\u001b[38;5;241m=\u001b[39mblank, reduction\u001b[38;5;241m=\u001b[39mreduction, zero_infinity\u001b[38;5;241m=\u001b[39mzero_infinity\n\u001b[0;32m   2659\u001b[0m     )\n\u001b[1;32m-> 2660\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mctc_loss\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m   2661\u001b[0m \u001b[43m    \u001b[49m\u001b[43mlog_probs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtargets\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43minput_lengths\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtarget_lengths\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mblank\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m_Reduction\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget_enum\u001b[49m\u001b[43m(\u001b[49m\u001b[43mreduction\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mzero_infinity\u001b[49m\n\u001b[0;32m   2662\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for epoch in range(1, epochs+1):\n",
    "    train(epoch, train_loader, model_to_train, adamW, one_cycle_lr, ctc_loss)\n",
    "    test(epoch, \"dev-clean\", dev_clean_loader, model_to_train, adamW, ctc_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation\n",
    "- pick the split you want to get your metrics on"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(1, \"dev-clean\", dev_clean_loader, model_to_train, adamW, ctc_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(1, \"dev-other\", dev_other_loader, model_to_train, adamW, ctc_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(1, \"test-clean\", test_clean_loader, model_to_train, adamW, ctc_loss)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test(1, \"test-other\", test_other_loader, jasper, adamW, ctc_loss)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "exp",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}