trial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import librosa\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import LabelEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 200/200 [00:46<00:00,  4.26it/s]\n"
     ]
    }
   ],
   "source": [
    "df = {\n",
    "    'file':[],\n",
    "    'y':[],\n",
    "    'sr':[],\n",
    "    'class':[],\n",
    "}\n",
    "files = sorted(os.listdir('../Project 1/Speech Recordings/'))\n",
    "for file in tqdm(files):\n",
    "    class_name = file.split('.')[0].split('_')[0]\n",
    "    # print(class_name,file)\n",
    "    y, sr  = librosa.load('../Project 1/Speech Recordings/' + file)\n",
    "    df['file'].append(file)\n",
    "    df['y'].append(y)\n",
    "    df['sr'].append(sr)\n",
    "    df['class'].append(class_name)\n",
    "    # break\n",
    "df = pd.DataFrame(df)\n",
    "df['targets'] = LabelEncoder().fit_transform(df['class'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "import librosa\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Load the audio file\n",
    "audio_path = \"../Project 1/Speech Recordings/Bikram_E_7.m4a\"\n",
    "y, sr = librosa.load(audio_path, sr=None)\n",
    "\n",
    "# Define the desired chunk length\n",
    "chunk_length = 512\n",
    "\n",
    "# Calculate the number of chunks\n",
    "num_chunks = len(y) // chunk_length\n",
    "\n",
    "# Extract chunks\n",
    "chunks = [y[i*chunk_length:(i+1)*chunk_length] for i in range(num_chunks)]\n",
    "\n",
    "# If there's any remaining audio at the end, include it as a final chunk\n",
    "if len(y) % chunk_length != 0:\n",
    "    chunks.append(y[num_chunks*chunk_length:])\n",
    "\n",
    "# Compute and plot spectrogram for each chunk\n",
    "spectrograms = []\n",
    "for i, chunk in enumerate(chunks):\n",
    "    # Compute the spectrogram using the short-time Fourier transform (STFT)\n",
    "    spectrogram = np.abs(librosa.stft(chunk))\n",
    "    spectrograms.append(spectrogram)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(512,)"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunks[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class PixelCNN(nn.Module):\n",
    "    def __init__(self, input_channels=2, num_layers=12, kernel_size=7, num_filters=64 , output_channels=1):\n",
    "        super(PixelCNN, self).__init__()\n",
    "        self.input_channels = input_channels\n",
    "        self.num_layers = num_layers\n",
    "        self.kernel_size = kernel_size\n",
    "        self.num_filters = num_filters\n",
    "        \n",
    "        # Define the initial convolutional layer\n",
    "        self.initial_conv = nn.Conv1d(input_channels, num_filters, kernel_size=kernel_size, padding=kernel_size//2)\n",
    "        \n",
    "        # Define the residual blocks\n",
    "        self.residual_blocks = nn.ModuleList([\n",
    "            ResidualBlock(num_filters, kernel_size) for _ in range(num_layers)\n",
    "        ])\n",
    "        \n",
    "        # Define the final convolutional layer for output\n",
    "        self.out_conv = nn.Conv1d(num_filters, output_channels, kernel_size=1)  # Output single-channel spectrogram\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.initial_conv(x)\n",
    "        \n",
    "        # Pass through residual blocks\n",
    "        for block in self.residual_blocks:\n",
    "            x = block(x)\n",
    "        \n",
    "        # Apply output convolution\n",
    "        x = self.out_conv(x)\n",
    "        return x\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, num_filters, kernel_size):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv1d(num_filters, num_filters, kernel_size=kernel_size, padding=kernel_size//2)\n",
    "        self.conv2 = nn.Conv1d(num_filters, num_filters, kernel_size=kernel_size, padding=kernel_size//2)\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        x = F.relu(x)\n",
    "        x = self.conv1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x += residual  # Add residual connection\n",
    "        return x\n",
    "\n",
    "# # Example usage:\n",
    "# model = PixelCNN(input_channels=3, num_layers=12, kernel_size=7, num_filters=64)\n",
    "# print(model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunk1 = torch.tensor([chunks[0],]).float()\n",
    "chunk2 = torch.tensor([chunks[1],]).float()\n",
    "chunk3 = torch.tensor([chunks[2],]).float()\n",
    "chunk = torch.cat([chunk1, chunk2,chunk3], dim=0).unsqueeze(0)\n",
    "# model(chunk).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 512])"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunk.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VectorQuantizerEMA(nn.Module):\n",
    "    def __init__(self, num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5):\n",
    "        super(VectorQuantizerEMA, self).__init__()\n",
    "\n",
    "        self._embedding_dim = embedding_dim\n",
    "        self._num_embeddings = num_embeddings\n",
    "\n",
    "        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)\n",
    "        self._embedding.weight.data.normal_()\n",
    "        self._commitment_cost = commitment_cost\n",
    "\n",
    "        self.register_buffer('_ema_cluster_size', torch.zeros(num_embeddings))\n",
    "        self._ema_w = nn.Parameter(torch.Tensor(num_embeddings, self._embedding_dim))\n",
    "        self._ema_w.data.normal_()\n",
    "\n",
    "        self._decay = decay\n",
    "        self._epsilon = epsilon\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        # convert inputs from BCHW -> BHWC\n",
    "        #BCH -> BHC\n",
    "        # print(inputs.shape)\n",
    "        inputs = inputs.permute(0, 2 , 1).contiguous()\n",
    "        input_shape = inputs.shape\n",
    "\n",
    "        # Flatten input\n",
    "        flat_input = inputs.view(-1, self._embedding_dim)\n",
    "\n",
    "        # Calculate distances\n",
    "        distances = (torch.sum(flat_input**2, dim=1, keepdim=True)\n",
    "                    + torch.sum(self._embedding.weight**2, dim=1)\n",
    "                    - 2 * torch.matmul(flat_input, self._embedding.weight.t()))\n",
    "\n",
    "        # Encoding\n",
    "        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)\n",
    "        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)\n",
    "        encodings.scatter_(1, encoding_indices, 1)\n",
    "\n",
    "        # Quantize and unflatten\n",
    "        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)\n",
    "\n",
    "        # Use EMA to update the embedding vectors\n",
    "        if self.training:\n",
    "            self._ema_cluster_size = self._ema_cluster_size * self._decay + \\\n",
    "                                     (1 - self._decay) * torch.sum(encodings, 0)\n",
    "\n",
    "            # Laplace smoothing of the cluster size\n",
    "            n = torch.sum(self._ema_cluster_size.data)\n",
    "            self._ema_cluster_size = (\n",
    "                (self._ema_cluster_size + self._epsilon)\n",
    "                / (n + self._num_embeddings * self._epsilon) * n)\n",
    "\n",
    "            dw = torch.matmul(encodings.t(), flat_input)\n",
    "            self._ema_w = nn.Parameter(self._ema_w * self._decay + (1 - self._decay) * dw)\n",
    "\n",
    "            self._embedding.weight = nn.Parameter(self._ema_w / self._ema_cluster_size.unsqueeze(1))\n",
    "\n",
    "        # Loss\n",
    "        e_latent_loss = F.mse_loss(quantized.detach(), inputs)\n",
    "        loss = self._commitment_cost * e_latent_loss\n",
    "\n",
    "        # Straight Through Estimator\n",
    "        quantized = inputs + (quantized - inputs).detach()\n",
    "        avg_probs = torch.mean(encodings, dim=0)\n",
    "        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))\n",
    "\n",
    "        # convert quantized from BHWC -> BCHW\n",
    "        #BHC -> BCH\n",
    "        return loss, quantized.permute(0, 2, 1).contiguous(), perplexity, encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "class VectorQuantizer(nn.Module):\n",
    "    def __init__(self, num_embeddings, embedding_dim, commitment_cost):\n",
    "        super(VectorQuantizer, self).__init__()\n",
    "\n",
    "        self._embedding_dim = embedding_dim\n",
    "        self._num_embeddings = num_embeddings\n",
    "\n",
    "        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)\n",
    "        self._embedding.weight.data.uniform_(-1/self._num_embeddings, 1/self._num_embeddings)\n",
    "        self._commitment_cost = commitment_cost\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        # convert inputs from BCHW -> BHWC\n",
    "        inputs = inputs.permute(0, 2 , 1).contiguous()\n",
    "        input_shape = inputs.shape\n",
    "\n",
    "        # Flatten input\n",
    "        flat_input = inputs.view(-1, self._embedding_dim)\n",
    "\n",
    "        # Calculate distances\n",
    "        distances = (torch.sum(flat_input**2, dim=1, keepdim=True)\n",
    "                    + torch.sum(self._embedding.weight**2, dim=1)\n",
    "                    - 2 * torch.matmul(flat_input, self._embedding.weight.t()))\n",
    "\n",
    "        # Encoding\n",
    "        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)\n",
    "        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)\n",
    "        encodings.scatter_(1, encoding_indices, 1)\n",
    "\n",
    "        # Quantize and unflatten\n",
    "        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)\n",
    "\n",
    "        # Loss\n",
    "        e_latent_loss = F.mse_loss(quantized.detach(), inputs)\n",
    "        q_latent_loss = F.mse_loss(quantized, inputs.detach())\n",
    "        loss = q_latent_loss + self._commitment_cost * e_latent_loss\n",
    "\n",
    "        quantized = inputs + (quantized - inputs).detach()\n",
    "        avg_probs = torch.mean(encodings, dim=0)\n",
    "        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))\n",
    "\n",
    "        # convert quantized from BHWC -> BCHW\n",
    "        return loss, quantized.permute(0, 2, 1).contiguous(), perplexity, encodings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Residual(nn.Module):\n",
    "    def __init__(self, in_channels, num_hiddens, num_residual_hiddens):\n",
    "        super(Residual, self).__init__()\n",
    "        self._block = nn.Sequential(\n",
    "            nn.ReLU(True),\n",
    "            nn.Conv1d(in_channels=in_channels,\n",
    "                      out_channels=num_residual_hiddens,\n",
    "                      kernel_size=3, stride=1, padding=1, bias=False),\n",
    "            nn.ReLU(True),\n",
    "            nn.Conv1d(in_channels=num_residual_hiddens,\n",
    "                      out_channels=num_hiddens,\n",
    "                      kernel_size=1, stride=1, bias=False)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self._block(x)\n",
    "\n",
    "\n",
    "class ResidualStack(nn.Module):\n",
    "    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):\n",
    "        super(ResidualStack, self).__init__()\n",
    "        self._num_residual_layers = num_residual_layers\n",
    "        self._layers = nn.ModuleList([Residual(in_channels, num_hiddens, num_residual_hiddens)\n",
    "                             for _ in range(self._num_residual_layers)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        for i in range(self._num_residual_layers):\n",
    "            x = self._layers[i](x)\n",
    "        return F.relu(x)\n",
    "    \n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):\n",
    "        super(Decoder, self).__init__()\n",
    "\n",
    "        self._conv_1 = nn.Conv1d(in_channels=in_channels,\n",
    "                                 out_channels=num_hiddens,\n",
    "                                 kernel_size=3,\n",
    "                                 stride=1, padding=1)\n",
    "\n",
    "        self._residual_stack = ResidualStack(in_channels=num_hiddens,\n",
    "                                             num_hiddens=num_hiddens,\n",
    "                                             num_residual_layers=num_residual_layers,\n",
    "                                             num_residual_hiddens=num_residual_hiddens)\n",
    "\n",
    "        self._conv_trans_1 = nn.ConvTranspose1d(in_channels=num_hiddens,\n",
    "                                                out_channels=num_hiddens//2,\n",
    "                                                kernel_size=4,\n",
    "                                                stride=2, padding=1)\n",
    "\n",
    "        self._conv_trans_2 = nn.ConvTranspose1d(in_channels=num_hiddens//2,\n",
    "                                                out_channels=1,\n",
    "                                                kernel_size=4,\n",
    "                                                stride=2, padding=1)\n",
    "\n",
    "    def forward(self, inputs):\n",
    "        x = self._conv_1(inputs)\n",
    "        x = self._residual_stack(x)\n",
    "\n",
    "        x = self._conv_trans_1(x)\n",
    "        x = F.relu(x)\n",
    "\n",
    "        return self._conv_trans_2(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self,chunk_size, num_hiddens, num_residual_layers, num_residual_hiddens,\n",
    "                 num_embeddings, embedding_dim, output_dim, commitment_cost, decay=0):\n",
    "        super(Model, self).__init__()\n",
    "        self.num_hiddens = num_hiddens\n",
    "        self._encoder = PixelCNN(input_channels=num_hiddens+1, num_layers=12, kernel_size=7, num_filters=64,output_channels=num_hiddens)\n",
    "# print(model)\n",
    "        self._pre_vq_conv = nn.Conv1d(in_channels=num_hiddens,\n",
    "                                      out_channels=embedding_dim,\n",
    "                                      kernel_size=1,\n",
    "                                      stride=1)\n",
    "        if decay > 0.0:\n",
    "            self._vq_vae = VectorQuantizerEMA(num_embeddings, embedding_dim,\n",
    "                                              commitment_cost, decay)\n",
    "        else:\n",
    "            self._vq_vae = VectorQuantizer(num_embeddings, embedding_dim,\n",
    "                                          commitment_cost)\n",
    "\n",
    "        self._decoder = Decoder(embedding_dim,\n",
    "                                num_hiddens,\n",
    "                                num_residual_layers,\n",
    "                                num_residual_hiddens)\n",
    "        self.mlp = nn.Linear(4*chunk_size, output_dim)\n",
    "    def forward(self, x,context=None):\n",
    "        if context is None:\n",
    "            context = torch.zeros(x.shape[0], self.num_hiddens, x.shape[2]).to(x)\n",
    "        x = torch.cat([x,context], dim=1)\n",
    "        # print(x.shape,'x')\n",
    "        z1 = self._encoder(x)\n",
    "        # print(z1.shape,'enc')\n",
    "        z = self._pre_vq_conv(z1)\n",
    "        # print(z.shape,'pre_vq_conv')\n",
    "        loss, quantized, perplexity, _ = self._vq_vae(z)\n",
    "        # print(quantized.shape,'quan')\n",
    "        x = self._decoder(quantized)\n",
    "        # print(x.shape,'dec')\n",
    "        emotion = self.mlp(x.view(x.shape[0], -1))\n",
    "        # return loss, quantized, perplexity\n",
    "        emotion = F.softmax(emotion, dim=1)\n",
    "        return emotion, z1, loss, perplexity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 256\n",
    "num_training_updates = 15000\n",
    "\n",
    "num_hiddens = 64 #128\n",
    "num_residual_hiddens = 32\n",
    "num_residual_layers = 2\n",
    "\n",
    "embedding_dim = 128\n",
    "num_embeddings = 512\n",
    "\n",
    "output_dim = 20 # 5 emotions\n",
    "\n",
    "commitment_cost = 1.0#0.25\n",
    "\n",
    "decay =  0.1\n",
    "\n",
    "learning_rate = 1e-4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import  DataLoader, Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AudioDataset(Dataset):\n",
    "    def __init__(self,df,chunk_length = 1024) -> None:\n",
    "        super().__init__()\n",
    "        self.data = df['y'].values\n",
    "        self.targets = df['targets'].values\n",
    "        self.chunk_length = chunk_length\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    def __getitem__(self, idx):\n",
    "        y = self.data[idx]\n",
    "        y, _ = librosa.effects.trim(y)\n",
    "        # Calculate the number of chunks\n",
    "        num_chunks = len(y) // self.chunk_length\n",
    "        # Extract chunks\n",
    "        chunks = [y[i*self.chunk_length:(i+1)*self.chunk_length] for i in range(num_chunks)]\n",
    "        # if the chunk only contains zeros, skip it\n",
    "        # chunks = [chunk for chunk in chunks if np.sum(chunk) != 0]\n",
    "        return chunks, self.targets[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {},
   "outputs": [],
   "source": [
    "# split the dataframe into training and validation dataframes\n",
    "df_train = df.sample(frac=0.8, random_state=0)\n",
    "df_val = df.drop(df_train.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunk_size = 1024\n",
    "train_dataset = AudioDataset(df_train,chunk_size)\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True)\n",
    "\n",
    "val_dataset = AudioDataset(df_val,chunk_size)\n",
    "val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[tensor([[-2.0662e-06,  1.8591e-06, -3.1475e-05,  ..., -6.3167e-03,\n",
       "           -6.7745e-03, -7.3603e-03]]),\n",
       "  tensor([[-0.0079, -0.0081, -0.0075,  ...,  0.0038,  0.0031,  0.0040]]),\n",
       "  tensor([[0.0051, 0.0053, 0.0045,  ..., 0.0044, 0.0024, 0.0019]]),\n",
       "  tensor([[ 0.0012, -0.0006, -0.0019,  ...,  0.0096,  0.0118,  0.0120]]),\n",
       "  tensor([[ 0.0104,  0.0081,  0.0057,  ..., -0.0058, -0.0025,  0.0005]]),\n",
       "  tensor([[ 0.0036,  0.0056,  0.0071,  ..., -0.0214, -0.0209, -0.0194]]),\n",
       "  tensor([[-0.0166, -0.0150, -0.0140,  ..., -0.0198, -0.0220, -0.0228]]),\n",
       "  tensor([[-0.0210, -0.0205, -0.0204,  ..., -0.0029, -0.0019, -0.0005]]),\n",
       "  tensor([[-0.0004, -0.0007,  0.0010,  ...,  0.0046,  0.0047,  0.0059]]),\n",
       "  tensor([[ 0.0058,  0.0060,  0.0094,  ..., -0.0079, -0.0057, -0.0014]]),\n",
       "  tensor([[-0.0001, -0.0005, -0.0019,  ..., -0.0034, -0.0053, -0.0062]]),\n",
       "  tensor([[-0.0049, -0.0023, -0.0014,  ..., -0.0053, -0.0111, -0.0157]]),\n",
       "  tensor([[-0.0183, -0.0207, -0.0224,  ...,  0.0085,  0.0065,  0.0046]]),\n",
       "  tensor([[ 0.0029,  0.0011, -0.0010,  ...,  0.0047,  0.0067,  0.0070]]),\n",
       "  tensor([[ 0.0073,  0.0079,  0.0080,  ..., -0.0064, -0.0031, -0.0004]]),\n",
       "  tensor([[0.0020, 0.0062, 0.0107,  ..., 0.0458, 0.0656, 0.0599]]),\n",
       "  tensor([[ 0.0367,  0.0128, -0.0060,  ..., -0.2414, -0.1825, -0.2334]]),\n",
       "  tensor([[-0.1784, -0.1276, -0.2585,  ...,  0.2449,  0.2233,  0.1896]]),\n",
       "  tensor([[ 0.1644,  0.1296,  0.1061,  ...,  0.0397,  0.0109, -0.0139]]),\n",
       "  tensor([[-0.0287, -0.0465, -0.0497,  ..., -0.0323, -0.0414, -0.0428]]),\n",
       "  tensor([[-0.0421, -0.0345,  0.0032,  ..., -0.0902, -0.0499, -0.0045]]),\n",
       "  tensor([[-0.0425, -0.0883, -0.0964,  ...,  0.0084, -0.0202, -0.0535]]),\n",
       "  tensor([[-0.0205,  0.0039,  0.0035,  ...,  0.1134,  0.0660,  0.0263]]),\n",
       "  tensor([[-0.0120, -0.0217, -0.0103,  ...,  0.0239,  0.0373,  0.0470]]),\n",
       "  tensor([[ 0.0627,  0.0564,  0.0491,  ..., -0.0104, -0.0174, -0.0258]]),\n",
       "  tensor([[ 0.0334,  0.1140,  0.1752,  ..., -0.0599, -0.0260,  0.0043]]),\n",
       "  tensor([[ 0.0478,  0.0905,  0.1243,  ..., -0.1435, -0.0703,  0.0375]]),\n",
       "  tensor([[ 0.0884,  0.0210, -0.1180,  ...,  0.1293,  0.1750,  0.1931]]),\n",
       "  tensor([[ 0.1900,  0.2171,  0.2418,  ..., -0.2131, -0.1935, -0.1627]]),\n",
       "  tensor([[-0.1331, -0.1177, -0.1444,  ...,  0.2438,  0.2183,  0.1895]]),\n",
       "  tensor([[ 0.1790,  0.1848,  0.2071,  ..., -0.2518, -0.2924, -0.2741]]),\n",
       "  tensor([[-0.2349, -0.1953, -0.0852,  ...,  0.1287,  0.1426,  0.1538]]),\n",
       "  tensor([[ 0.1799,  0.1994,  0.2228,  ...,  0.0739,  0.0323, -0.0035]]),\n",
       "  tensor([[ 0.0107,  0.0315,  0.0531,  ..., -0.0396,  0.0020,  0.0254]]),\n",
       "  tensor([[ 0.1041,  0.0744, -0.0530,  ..., -0.0052,  0.0225,  0.0459]]),\n",
       "  tensor([[0.0317, 0.0040, 0.0024,  ..., 0.0178, 0.0482, 0.0733]]),\n",
       "  tensor([[0.0934, 0.0996, 0.1116,  ..., 0.1139, 0.0845, 0.0635]]),\n",
       "  tensor([[ 0.0404,  0.0177,  0.0010,  ..., -0.1983, -0.1728, -0.1616]]),\n",
       "  tensor([[-0.1583, -0.1572, -0.1557,  ...,  0.0469,  0.0458,  0.0439]]),\n",
       "  tensor([[ 0.0380,  0.0310,  0.0259,  ..., -0.3096, -0.3084, -0.2846]]),\n",
       "  tensor([[-0.2598, -0.2117, -0.1590,  ..., -0.1806, -0.2369, -0.2792]]),\n",
       "  tensor([[-0.3059, -0.3354, -0.3671,  ...,  0.0050, -0.0014, -0.0062]]),\n",
       "  tensor([[-0.0036, -0.0048, -0.0008,  ..., -0.4951, -0.5077, -0.4971]]),\n",
       "  tensor([[-0.4561, -0.3933, -0.3256,  ...,  0.0292,  0.0384,  0.0443]]),\n",
       "  tensor([[ 0.0470,  0.0486,  0.0514,  ..., -0.0903, -0.0883, -0.0848]]),\n",
       "  tensor([[-0.0799, -0.0747, -0.0692,  ..., -0.0092, -0.0139, -0.0194]]),\n",
       "  tensor([[-0.0236, -0.0303, -0.0390,  ..., -0.0332, -0.0594, -0.0739]]),\n",
       "  tensor([[-0.0702, -0.0941, -0.1410,  ...,  0.1733,  0.1204,  0.0527]]),\n",
       "  tensor([[-0.0146, -0.0828, -0.1516,  ..., -0.0121,  0.0015,  0.0157]]),\n",
       "  tensor([[ 0.0142,  0.0125,  0.0032,  ..., -0.0541, -0.0530, -0.0494]]),\n",
       "  tensor([[-0.0449, -0.0542, -0.0787,  ...,  0.2862,  0.3003,  0.3122]]),\n",
       "  tensor([[ 0.3141,  0.2858,  0.2347,  ...,  0.0161, -0.0045, -0.0309]]),\n",
       "  tensor([[-0.0448, -0.0563, -0.0710,  ..., -0.0097, -0.0026,  0.0056]]),\n",
       "  tensor([[0.0107, 0.0157, 0.0185,  ..., 0.0321, 0.0341, 0.0359]]),\n",
       "  tensor([[0.0373, 0.0375, 0.0361,  ..., 0.0009, 0.0019, 0.0033]]),\n",
       "  tensor([[ 0.0044,  0.0048,  0.0047,  ..., -0.0218, -0.0221, -0.0222]]),\n",
       "  tensor([[-0.0216, -0.0193, -0.0199,  ..., -0.0202, -0.0248, -0.0295]]),\n",
       "  tensor([[-0.0362, -0.0426, -0.0441,  ..., -0.0221, -0.0428, -0.1954]]),\n",
       "  tensor([[-0.2219, -0.2668, -0.4804,  ...,  0.0914,  0.0211, -0.0575]]),\n",
       "  tensor([[-0.0917, -0.1504, -0.2156,  ...,  0.2970,  0.2733,  0.2671]]),\n",
       "  tensor([[ 0.3020,  0.3228,  0.3014,  ..., -0.0247, -0.0247, -0.0215]]),\n",
       "  tensor([[-0.0162, -0.0101, -0.0044,  ...,  0.0048,  0.0013, -0.0011]]),\n",
       "  tensor([[-0.0008,  0.0008,  0.0068,  ..., -0.3076, -0.3002, -0.2982]]),\n",
       "  tensor([[-0.3071, -0.3064, -0.2968,  ..., -0.1743, -0.1370, -0.1064]]),\n",
       "  tensor([[-0.1000, -0.0967, -0.0903,  ..., -0.3393, -0.3740, -0.3852]]),\n",
       "  tensor([[-0.3387, -0.2740, -0.2332,  ...,  0.0500,  0.0466,  0.0420]]),\n",
       "  tensor([[0.0363, 0.0293, 0.0233,  ..., 0.0831, 0.1023, 0.1848]]),\n",
       "  tensor([[0.1905, 0.1825, 0.2540,  ..., 0.5401, 0.3865, 0.3319]]),\n",
       "  tensor([[ 0.2363,  0.1664,  0.0814,  ..., -0.0304, -0.0408, -0.0479]]),\n",
       "  tensor([[-0.0461, -0.0443, -0.0414,  ..., -0.0213, -0.0218, -0.0249]]),\n",
       "  tensor([[ 0.0019,  0.0117, -0.0014,  ...,  0.1891,  0.2108,  0.1851]]),\n",
       "  tensor([[ 0.1870,  0.1783,  0.1510,  ...,  0.0077, -0.0038,  0.0106]]),\n",
       "  tensor([[0.0153, 0.0482, 0.0612,  ..., 0.3442, 0.3258, 0.2966]]),\n",
       "  tensor([[ 0.2859,  0.2703,  0.2571,  ..., -0.4128, -0.4737, -0.4477]]),\n",
       "  tensor([[-0.4318, -0.3826, -0.2775,  ...,  0.0140,  0.0149,  0.0157]]),\n",
       "  tensor([[0.0157, 0.0162, 0.0168,  ..., 0.0191, 0.0203, 0.0205]]),\n",
       "  tensor([[ 0.0192,  0.0159,  0.0111,  ..., -0.0002,  0.0035,  0.0071]]),\n",
       "  tensor([[ 0.0086,  0.0074,  0.0049,  ..., -0.0230, -0.0072,  0.0102]]),\n",
       "  tensor([[-0.0190, -0.0271, -0.0177,  ...,  0.1726,  0.2204,  0.2324]]),\n",
       "  tensor([[0.2338, 0.2012, 0.1271,  ..., 0.0086, 0.0022, 0.0150]]),\n",
       "  tensor([[ 0.0022,  0.0050,  0.0249,  ..., -0.1754, -0.1755, -0.1694]]),\n",
       "  tensor([[-0.1618, -0.1442, -0.1270,  ..., -0.0091, -0.0079, -0.0056]]),\n",
       "  tensor([[-0.0040, -0.0040, -0.0045,  ..., -0.0019, -0.0024, -0.0032]]),\n",
       "  tensor([[-0.0037, -0.0039, -0.0023,  ...,  0.0354,  0.0340,  0.0318]]),\n",
       "  tensor([[ 0.0294,  0.0272,  0.0255,  ..., -0.0153, -0.0155, -0.0152]]),\n",
       "  tensor([[-0.0144, -0.0136, -0.0129,  ..., -0.0309, -0.0338, -0.0365]]),\n",
       "  tensor([[-0.0389, -0.0401, -0.0399,  ...,  0.0196,  0.0167,  0.0135]]),\n",
       "  tensor([[0.0107, 0.0077, 0.0048,  ..., 0.0142, 0.0149, 0.0149]]),\n",
       "  tensor([[ 0.0146,  0.0143,  0.0135,  ..., -0.0713, -0.0649, -0.0584]]),\n",
       "  tensor([[-0.0518, -0.0460, -0.0419,  ..., -0.0039, -0.0017,  0.0008]]),\n",
       "  tensor([[ 0.0040,  0.0074,  0.0102,  ..., -0.0101, -0.0109, -0.0122]]),\n",
       "  tensor([[-0.0130, -0.0132, -0.0142,  ...,  0.0554,  0.0563,  0.0556]]),\n",
       "  tensor([[0.0532, 0.0503, 0.0477,  ..., 0.0093, 0.0124, 0.0146]]),\n",
       "  tensor([[0.0159, 0.0159, 0.0147,  ..., 0.0455, 0.0466, 0.0469]]),\n",
       "  tensor([[0.0464, 0.0446, 0.0413,  ..., 0.0075, 0.0075, 0.0071]]),\n",
       "  tensor([[0.0071, 0.0073, 0.0083,  ..., 0.0101, 0.0084, 0.0061]]),\n",
       "  tensor([[ 0.0031,  0.0005, -0.0016,  ...,  0.0016,  0.0009, -0.0005]]),\n",
       "  tensor([[-0.0025, -0.0043, -0.0055,  ...,  0.0052,  0.0028,  0.0003]]),\n",
       "  tensor([[-0.0025, -0.0055, -0.0089,  ..., -0.0185, -0.0157, -0.0121]]),\n",
       "  tensor([[-0.0087, -0.0066, -0.0062,  ...,  0.0137,  0.0177,  0.0214]]),\n",
       "  tensor([[ 0.0246,  0.0277,  0.0297,  ..., -0.0178, -0.0215, -0.0240]])],\n",
       " tensor([17], dtype=torch.int32)]"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "next(iter(val_dataloader))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model(chunk_size,num_hiddens, num_residual_layers, num_residual_hiddens,\n",
    "              num_embeddings, embedding_dim,output_dim,\n",
    "              commitment_cost, decay)\n",
    "ngpu = torch.cuda.device_count()\n",
    "if ngpu > 1:\n",
    "    model =  nn.DataParallel(model)\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)\n",
    "loss_fn = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 1/160 [01:01<2:41:39, 61.00s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "batch:1, Loss: 0.18692758679389954 , Accuracy: 0.002187120291616039\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 1/160 [01:11<3:08:56, 71.30s/it]\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[145], line 36\u001b[0m\n\u001b[0;32m     34\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m chunk \u001b[38;5;129;01min\u001b[39;00m x:\n\u001b[0;32m     35\u001b[0m     chunk \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mtensor(chunk)\u001b[38;5;241m.\u001b[39munsqueeze(\u001b[38;5;241m0\u001b[39m)\u001b[38;5;241m.\u001b[39mfloat()\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m---> 36\u001b[0m     emotion, context, vq_loss, perplexity \u001b[38;5;241m=\u001b[39m \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchunk\u001b[49m\u001b[43m,\u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     37\u001b[0m     Loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m loss_fn(emotion, y)\n\u001b[0;32m     38\u001b[0m     predicted \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39margmax(emotion\u001b[38;5;241m.\u001b[39mdata)\n",
      "File \u001b[1;32mc:\\ProgramData\\Anaconda3\\envs\\speech_env\\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;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\ProgramData\\Anaconda3\\envs\\speech_env\\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 forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\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",
      "Cell \u001b[1;32mIn[71], line 33\u001b[0m, in \u001b[0;36mModel.forward\u001b[1;34m(self, x, context)\u001b[0m\n\u001b[0;32m     31\u001b[0m z \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pre_vq_conv(z1)\n\u001b[0;32m     32\u001b[0m \u001b[38;5;66;03m# print(z.shape,'pre_vq_conv')\u001b[39;00m\n\u001b[1;32m---> 33\u001b[0m loss, quantized, perplexity, _ \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_vq_vae\u001b[49m\u001b[43m(\u001b[49m\u001b[43mz\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     34\u001b[0m \u001b[38;5;66;03m# print(quantized.shape,'quan')\u001b[39;00m\n\u001b[0;32m     35\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_decoder(quantized)\n",
      "File \u001b[1;32mc:\\ProgramData\\Anaconda3\\envs\\speech_env\\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;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[1;32mc:\\ProgramData\\Anaconda3\\envs\\speech_env\\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 forward_call(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\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",
      "Cell \u001b[1;32mIn[56], line 40\u001b[0m, in \u001b[0;36mVectorQuantizerEMA.forward\u001b[1;34m(self, inputs)\u001b[0m\n\u001b[0;32m     37\u001b[0m encodings\u001b[38;5;241m.\u001b[39mscatter_(\u001b[38;5;241m1\u001b[39m, encoding_indices, \u001b[38;5;241m1\u001b[39m)\n\u001b[0;32m     39\u001b[0m \u001b[38;5;66;03m# Quantize and unflatten\u001b[39;00m\n\u001b[1;32m---> 40\u001b[0m quantized \u001b[38;5;241m=\u001b[39m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmatmul\u001b[49m\u001b[43m(\u001b[49m\u001b[43mencodings\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[43m_embedding\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mweight\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mview(input_shape)\n\u001b[0;32m     42\u001b[0m \u001b[38;5;66;03m# Use EMA to update the embedding vectors\u001b[39;00m\n\u001b[0;32m     43\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtraining:\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "epochs = 100\n",
    "model.train()\n",
    "gradient_accumulation_steps = 16\n",
    "for epoch in range(epochs):\n",
    "    epoch_loss = 0.0\n",
    "    for idx, (x, y) in tqdm(enumerate(train_dataloader),total=len(train_dataloader),dynamic_ncols=True):\n",
    "            y = y.to(device).long()\n",
    "            context = None\n",
    "            VQ_loss = 0.0\n",
    "            Loss = 0.0\n",
    "            for chunk in x:\n",
    "                chunk = torch.tensor(chunk).unsqueeze(0).float().to(device)\n",
    "                emotion, context, vq_loss, perplexity = model(chunk,context)\n",
    "                VQ_loss += vq_loss\n",
    "                Loss += loss_fn(emotion, y)\n",
    "            \n",
    "            loss = (Loss + VQ_loss)/len(x)\n",
    "            loss = loss / gradient_accumulation_steps\n",
    "            loss.backward()\n",
    "            epoch_loss += loss.item()\n",
    "            if (idx+1) % gradient_accumulation_steps == 0:\n",
    "                optimizer.step()\n",
    "                optimizer.zero_grad()\n",
    "            if idx % 1 == 0:\n",
    "                correct = 0\n",
    "                total = 0\n",
    "                with torch.no_grad():\n",
    "                    for x, y in val_dataloader:\n",
    "                        y = y.to(device).long()\n",
    "                        context = None\n",
    "                        for chunk in x:\n",
    "                            chunk = torch.tensor(chunk).unsqueeze(0).float().to(device)\n",
    "                            emotion, context, vq_loss, perplexity = model(chunk,context)\n",
    "                            Loss += loss_fn(emotion, y)\n",
    "                            predicted = torch.argmax(emotion.data)\n",
    "                            total += 1\n",
    "                            correct += (predicted == y).sum().item()\n",
    "                print(f'batch:{idx +1 }, Loss: {loss.item()} , Accuracy: {correct/total}')\n",
    "    \n",
    "    epoch_loss = epoch_loss / len(train_dataloader)\n",
    "    print(f'Epoch : [{epoch}/{epochs}], Loss: {epoch_loss}')\n",
    "    state_dict = {'model': model.state_dict(),'epoch': epoch,'epoch_loss': epoch_loss}\n",
    "    torch.save(state_dict, f'./model_VQ_Emo.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "iteration over a 0-d tensor",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[131], line 8\u001b[0m\n\u001b[0;32m      6\u001b[0m     emotion, context, vq_loss, perplexity \u001b[38;5;241m=\u001b[39m model(chunk,context)\n\u001b[0;32m      7\u001b[0m     Loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m loss_fn(emotion, y)\n\u001b[1;32m----> 8\u001b[0m _, predicted \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39margmax(emotion\u001b[38;5;241m.\u001b[39mdata)\n\u001b[0;32m      9\u001b[0m \u001b[38;5;28mprint\u001b[39m(predicted, y)\n",
      "File \u001b[1;32mc:\\ProgramData\\Anaconda3\\envs\\speech_env\\lib\\site-packages\\torch\\_tensor.py:1022\u001b[0m, in \u001b[0;36mTensor.__iter__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m   1012\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__iter__\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m   1013\u001b[0m     \u001b[38;5;66;03m# NB: we use 'imap' and not 'map' here, so that in Python 2 we get a\u001b[39;00m\n\u001b[0;32m   1014\u001b[0m     \u001b[38;5;66;03m# generator and don't eagerly perform all the indexes.  This could\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m   1019\u001b[0m     \u001b[38;5;66;03m# NB: We have intentionally skipped __torch_function__ dispatch here.\u001b[39;00m\n\u001b[0;32m   1020\u001b[0m     \u001b[38;5;66;03m# See gh-54457\u001b[39;00m\n\u001b[0;32m   1021\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdim() \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m-> 1022\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124miteration over a 0-d tensor\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m   1023\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_C\u001b[38;5;241m.\u001b[39m_get_tracing_state():\n\u001b[0;32m   1024\u001b[0m         warnings\u001b[38;5;241m.\u001b[39mwarn(\n\u001b[0;32m   1025\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mIterating over a tensor might cause the trace to be incorrect. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m   1026\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPassing a tensor of different shape won\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mt change the number of \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m   1030\u001b[0m             stacklevel\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m,\n\u001b[0;32m   1031\u001b[0m         )\n",
      "\u001b[1;31mTypeError\u001b[0m: iteration over a 0-d tensor"
     ]
    }
   ],
   "source": [
    "for x, y in val_dataloader:\n",
    "    y = y.to(device).long()\n",
    "    context = None\n",
    "    for chunk in x:\n",
    "        chunk = torch.tensor(chunk).unsqueeze(0).float().to(device)\n",
    "        emotion, context, vq_loss, perplexity = model(chunk,context)\n",
    "        Loss += loss_fn(emotion, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(13, 15)"
      ]
     },
     "execution_count": 144,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "torch.argmax(emotion.data).item(), y.item()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "15"
      ]
     },
     "execution_count": 143,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([15], device='cuda:0')"
      ]
     },
     "execution_count": 141,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "speech_env",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}