Music_Genres_classification.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fb9d011d-afcc-424c-8a5a-8956358cf48b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['CUDA_VISIBLE_DEVICES'] = '3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "22e4df26-4485-4e9a-be50-834e1a6cb95e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/anaconda3/envs/pytorch/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from src.efficient_kan import KAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e11b5943-fc0b-469f-b9d8-0190642c77a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a8d8fe9f-dd92-4726-87c3-2ccc7bc368d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3fdb2336-dcb3-425c-b0a3-d5652ebd878a",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = {}\n",
    "config['encoder'] = 'Wav2Vec'\n",
    "config['input_dim'] = 1024*16\n",
    "config['feature_dim'] = 36864\n",
    "config['batch_size'] = 1024\n",
    "config['hidden_layers'] = [1024, 64]\n",
    "config['lr'] = 1e-3\n",
    "config['epoch'] = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a91ba2d5-c941-4cdd-b853-6ec06bb884d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import librosa\n",
    "import librosa.display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ca9d6de5-88a5-46cc-81b0-234f8743566a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class2label = {'blues': 0,\n",
    "               'classical': 1,\n",
    "               'country': 2,\n",
    "               'metal': 3,\n",
    "               'disco': 4,\n",
    "               'hiphop': 5,\n",
    "               'jazz': 6,\n",
    "               'pop': 7,\n",
    "               'reggae': 8,\n",
    "               'rock': 9,}\n",
    "               \n",
    "class genres_classification_dataset(Dataset):\n",
    "    def __init__(self, root='data/Data/genres_original/', input_dim=1024, train=True):\n",
    "        super(genres_classification_dataset, self).__init__()\n",
    "        self.classes = os.listdir(root)\n",
    "        self.data = []\n",
    "        self.train = train\n",
    "        for _class in self.classes:\n",
    "            wavs = os.listdir(os.path.join(root, _class))\n",
    "            if train:\n",
    "                sample_wav = wavs[:len(wavs)//2]\n",
    "            else:\n",
    "                sample_wav = wavs[len(wavs)//2:]\n",
    "            for wav in sample_wav:\n",
    "                sound, sample_rate = librosa.load(os.path.join(root, _class, wav))\n",
    "                for i in range(len(sound)//input_dim):\n",
    "                    d={}\n",
    "                    clip = sound[i*input_dim:(i+1)*input_dim]\n",
    "                    d['clip'] = clip\n",
    "                    d['class'] = _class\n",
    "                    self.data.append(d)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        data = self.data[index]\n",
    "        _class = data['class']\n",
    "        clip = data['clip']\n",
    "        # sound, sample_rate = librosa.load(wav_path)\n",
    "        target = torch.from_numpy(clip).squeeze().unsqueeze(0).to(torch.float32)\n",
    "        onehot = torch.zeros(10)\n",
    "        onehot[class2label[_class]] = 1\n",
    "        return target, onehot\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "\n",
    "trainset = genres_classification_dataset(input_dim=config['input_dim'], train=True)\n",
    "valset = genres_classification_dataset(input_dim=config['input_dim'], train=False)\n",
    "trainloader = DataLoader(trainset, batch_size=config['batch_size'], shuffle=True)\n",
    "valloader = DataLoader(trainset, batch_size=config['batch_size'], shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e8d0e273-8ae0-4e21-a70b-075ba38092af",
   "metadata": {},
   "outputs": [],
   "source": [
    "class KAE(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        layers_hidden,\n",
    "        input_dim,\n",
    "    ):\n",
    "        super(KAE, self).__init__()\n",
    "        self.encoder = KAN([input_dim, layers_hidden[0], layers_hidden[1]])\n",
    "        self.decoder = KAN([layers_hidden[1], layers_hidden[0], input_dim])\n",
    "\n",
    "    def forward(self, x: torch.Tensor):\n",
    "        x = self.decoder(self.encoder(x))\n",
    "        return x\n",
    "class simple_classification(torch.nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_dim,\n",
    "        layers_hidden,\n",
    "        output_dim=10,\n",
    "    ):\n",
    "        super(simple_classification, self).__init__()\n",
    "        self.model = nn.Sequential(\n",
    "            nn.Linear(input_dim, layers_hidden[0]),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(layers_hidden[0], layers_hidden[1]),\n",
    "            nn.ReLU(inplace=True),\n",
    "            # nn.LayerNorm(layers_hidden[1]),\n",
    "            nn.Linear(layers_hidden[1], output_dim)\n",
    "            )\n",
    "\n",
    "    def forward(self, x: torch.Tensor):\n",
    "        x = self.model(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "700115ee-8201-42be-a3c9-17e87bd0b4d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "import sys\n",
    "if config['encoder'] == 'beats':\n",
    "    new_path = 'third_part/beats'\n",
    "    sys.path.append(new_path)\n",
    "    from BEATs import BEATs, BEATsConfig\n",
    "    # load the pre-trained checkpoints\n",
    "    checkpoint = torch.load('BEATs_iter3_plus_AS2M.pt')\n",
    "    \n",
    "    cfg = BEATsConfig(checkpoint['cfg'])\n",
    "    BEATs_model = BEATs(cfg)\n",
    "    BEATs_model.load_state_dict(checkpoint['model'])\n",
    "    BEATs_model.eval()\n",
    "    BEATs_model.to(device)\n",
    "    model = simple_classification(config['feature_dim'], config['hidden_layers'])\n",
    "elif config['encoder'] == 'KAE':\n",
    "    # Define model\n",
    "    kae_pretrain = torch.load('checkpoints/KAN_music_genres_pretrain.pth')\n",
    "    kae_pretrain.to(device)\n",
    "    model = simple_classification(config['feature_dim'], config['hidden_layers'])\n",
    "elif config['encoder'] == 'KAN':\n",
    "    model = KAN([config['input_dim'], 4096, 1024, 64, 10], grid_size=5, spline_order=4,)\n",
    "elif config['encoder'] == 'SSAST':\n",
    "    new_path = 'third_part/ssast'\n",
    "    sys.path.append(new_path)\n",
    "    from src.models import ASTModel\n",
    "    pretrain_model_path = 'SSAST-Base-Patch-400.pth'\n",
    "    import torchaudio.compliance.kaldi as ta_kaldi\n",
    "    def ssast_preprocess(\n",
    "                source: torch.Tensor,\n",
    "                fbank_mean: float = 15.41663,\n",
    "                fbank_std: float = 6.55582,\n",
    "        ) -> torch.Tensor:\n",
    "            fbanks = []\n",
    "            for waveform in source:\n",
    "                waveform = waveform.unsqueeze(0)\n",
    "                waveform = waveform - waveform.mean()\n",
    "                fbank = ta_kaldi.fbank(waveform, htk_compat=True, sample_frequency=33000, use_energy=False,\n",
    "                                                      window_type='hanning', num_mel_bins=128, dither=0.0, frame_shift=10)\n",
    "                fbanks.append(fbank)\n",
    "            fbank = torch.stack(fbanks, dim=0)\n",
    "            fbank = (fbank - fbank_mean) / (2 * fbank_std)\n",
    "            return fbank\n",
    "    input_tdim = 48  # fine-tuning data length can be different with pretraining data length\n",
    "    ast_mdl = ASTModel(label_dim=35,\n",
    "                 fshape=16, tshape=16, fstride=10, tstride=10,\n",
    "                 input_fdim=128, input_tdim=input_tdim, model_size='base',\n",
    "                 pretrain_stage=False, load_pretrained_mdl_path=pretrain_model_path)\n",
    "    model = simple_classification(config['feature_dim'], config['hidden_layers'])\n",
    "elif config['encoder'] == 'Wav2Vec':\n",
    "    new_path = 'third_part/soxan'\n",
    "    sys.path.append(new_path)\n",
    "    from models import Wav2Vec2ForSpeechClassification\n",
    "    pretrain_model_path = 'soxan_checkpoints'\n",
    "    import torchaudio\n",
    "    from transformers import AutoConfig, Wav2Vec2FeatureExtractor\n",
    "    feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(pretrain_model_path)\n",
    "    sampling_rate = feature_extractor.sampling_rate\n",
    "    \n",
    "    # for wav2vec\n",
    "    w2v = Wav2Vec2ForSpeechClassification.from_pretrained(pretrain_model_path).to(device)\n",
    "    resampler = torchaudio.transforms.Resample(16800, sampling_rate).to(device)\n",
    "    model = simple_classification(config['feature_dim'], config['hidden_layers'])\n",
    "else:\n",
    "    raise NotImplementError\n",
    "model.to(device)\n",
    "# Define optimizer\n",
    "optimizer = optim.AdamW(model.parameters(), lr=config['lr'], weight_decay=1e-4)\n",
    "# Define learning rate scheduler\n",
    "scheduler = optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.8)\n",
    "\n",
    "# Define loss\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "11654d6b-fd9c-4a1c-ab22-d427f14a3915",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.08it/s, loss=1.65, lr=0.001]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Val Loss: 1.6312610328197479\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.63, lr=0.0008]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2, Val Loss: 1.5856821477413177\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:21<00:00,  1.06s/it, loss=1.52, lr=0.00064]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 3, Val Loss: 1.5677696883678436\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:21<00:00,  1.06s/it, loss=1.48, lr=0.000512]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 4, Val Loss: 1.550961172580719\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.54, lr=0.00041]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 5, Val Loss: 1.5517475485801697\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.52, lr=0.000328]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 6, Val Loss: 1.5335833489894868\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.53, lr=0.000262]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 7, Val Loss: 1.525952821969986\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.53, lr=0.00021]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 8, Val Loss: 1.5200729072093964\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.62, lr=0.000168]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 9, Val Loss: 1.5126644670963287\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.48, lr=0.000134]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10, Val Loss: 1.5128972023725509\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.5, lr=0.000107] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 11, Val Loss: 1.5062468856573106\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.04it/s, loss=1.42, lr=8.59e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 12, Val Loss: 1.4989856511354447\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.04it/s, loss=1.48, lr=6.87e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 13, Val Loss: 1.500696748495102\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.37, lr=5.5e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 14, Val Loss: 1.4973071992397309\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.51, lr=4.4e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 15, Val Loss: 1.4959874749183655\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.04it/s, loss=1.49, lr=3.52e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 16, Val Loss: 1.4929079294204712\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.04it/s, loss=1.44, lr=2.81e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 17, Val Loss: 1.4931572258472443\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.04it/s, loss=1.51, lr=2.25e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 18, Val Loss: 1.4912997961044312\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.04it/s, loss=1.54, lr=1.8e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 19, Val Loss: 1.490935242176056\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.46, lr=1.44e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 20, Val Loss: 1.4903143256902696\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.56, lr=1.15e-5]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 21, Val Loss: 1.490217262506485\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.5, lr=9.22e-6] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 22, Val Loss: 1.4896993815898896\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.47, lr=7.38e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 23, Val Loss: 1.489520800113678\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.51, lr=5.9e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 24, Val Loss: 1.489421832561493\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.51, lr=4.72e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 25, Val Loss: 1.4891971409320832\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.48, lr=3.78e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 26, Val Loss: 1.4891792505979538\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.4, lr=3.02e-6] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 27, Val Loss: 1.4891421258449555\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.48, lr=2.42e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 28, Val Loss: 1.4889690548181533\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.48, lr=1.93e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 29, Val Loss: 1.4889730513095856\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.54, lr=1.55e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 30, Val Loss: 1.4888917446136474\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.47, lr=1.24e-6]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 31, Val Loss: 1.4888496547937393\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.47, lr=9.9e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 32, Val Loss: 1.4888130277395248\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.37, lr=7.92e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 33, Val Loss: 1.488776046037674\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.05it/s, loss=1.54, lr=6.34e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 34, Val Loss: 1.4887786149978637\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.54, lr=5.07e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 35, Val Loss: 1.4887514531612396\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.54, lr=4.06e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 36, Val Loss: 1.488738626241684\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.48, lr=3.25e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 37, Val Loss: 1.4887233793735504\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.5, lr=2.6e-7] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 38, Val Loss: 1.4887136667966843\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.45, lr=2.08e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 39, Val Loss: 1.488707810640335\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.47, lr=1.66e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 40, Val Loss: 1.4887009739875794\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.48, lr=1.33e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 41, Val Loss: 1.488697847723961\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.05it/s, loss=1.48, lr=1.06e-7]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 42, Val Loss: 1.4886965245008468\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.43, lr=8.51e-8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 43, Val Loss: 1.4886909246444702\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.5, lr=6.81e-8] \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 44, Val Loss: 1.4886894971132278\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:21<00:00,  1.06s/it, loss=1.52, lr=5.44e-8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 45, Val Loss: 1.4886873662471771\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.41, lr=4.36e-8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 46, Val Loss: 1.4886861741542816\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.43, lr=3.48e-8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 47, Val Loss: 1.4886856436729432\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.05it/s, loss=1.54, lr=2.79e-8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 48, Val Loss: 1.4886843681335449\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:19<00:00,  1.05it/s, loss=1.51, lr=2.23e-8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 49, Val Loss: 1.4886827677488328\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 20/20 [00:18<00:00,  1.06it/s, loss=1.43, lr=1.78e-8]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 50, Val Loss: 1.4886819303035737\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(config['epoch']):\n",
    "    # Train\n",
    "    model.train()\n",
    "    with tqdm(trainloader) as pbar:\n",
    "        for i, (_input, label) in enumerate(pbar):\n",
    "            _input = _input.view(-1, config['input_dim']).to(device)\n",
    "            label = label.to(device)\n",
    "            optimizer.zero_grad()\n",
    "            with torch.no_grad():\n",
    "                if config['encoder'] == 'beats':\n",
    "                    feature = BEATs_model.extract_features(_input, padding_mask=torch.zeros(_input.shape[0], config['input_dim']).bool().to(device))[0]\n",
    "                    feature = feature.reshape(_input.shape[0], -1)\n",
    "                elif config['encoder'] == 'KAE':\n",
    "                    feature = kae_pretrain.encoder(_input)\n",
    "                elif config['encoder'] == 'SSAST':\n",
    "                    filter_bank = ssast_preprocess(_input)\n",
    "                    feature = ast_mdl(filter_bank, task='extract_feature')\n",
    "                    feature = feature.reshape(_input.shape[0], -1)\n",
    "                elif config['encoder'] == 'Wav2Vec':\n",
    "                    speech = resampler(_input).squeeze()\n",
    "                    inputs = feature_extractor(speech, sampling_rate=sampling_rate, return_tensors=\"pt\", padding=True)\n",
    "                    inputs = {key: inputs[key].to(device).squeeze() for key in inputs}\n",
    "                    feature = w2v.extract_feature(**inputs)\n",
    "                    feature = feature.reshape(_input.shape[0], -1)\n",
    "                else:\n",
    "                    feature = _input\n",
    "                    \n",
    "            # if config['encoder'] == 'KAN':\n",
    "            #     output = model(feature, update_grid=(i % 20 == 0))\n",
    "            # else:\n",
    "            output = model(feature)\n",
    "            loss = criterion(output, label)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            pbar.set_postfix(loss=loss.detach().cpu().item(), lr=optimizer.param_groups[0]['lr'])\n",
    "\n",
    "    # # Validation\n",
    "    model.eval()\n",
    "    val_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for i, (_input, label) in enumerate(valloader):\n",
    "            _input = _input.view(-1, config['input_dim']).to(device)\n",
    "            label = label.to(device)\n",
    "            if config['encoder'] == 'beats':\n",
    "                feature = BEATs_model.extract_features(_input, padding_mask=torch.zeros(_input.shape[0], config['input_dim']).bool().to(device))[0]\n",
    "                feature = feature.reshape(_input.shape[0], -1)\n",
    "            elif config['encoder'] == 'KAE':\n",
    "                feature = kae_pretrain.encoder(_input)\n",
    "            elif config['encoder'] == 'SSAST':\n",
    "                filter_bank = ssast_preprocess(_input)\n",
    "                feature = ast_mdl(filter_bank, task='extract_feature')\n",
    "                feature = feature.reshape(_input.shape[0], -1)\n",
    "            elif config['encoder'] == 'Wav2Vec':\n",
    "                speech = resampler(_input).squeeze()\n",
    "                inputs = feature_extractor(speech, sampling_rate=sampling_rate, return_tensors=\"pt\", padding=True)\n",
    "                inputs = {key: inputs[key].to(device).squeeze() for key in inputs}\n",
    "                feature = w2v.extract_feature(**inputs)\n",
    "                feature = feature.reshape(_input.shape[0], -1)\n",
    "            else:\n",
    "                feature = _input\n",
    "            output = model(feature)\n",
    "            val_loss += criterion(output, label).item()\n",
    "            \n",
    "    val_loss /= len(valloader)\n",
    "\n",
    "    # Update learning rate\n",
    "    scheduler.step()\n",
    "\n",
    "    print(\n",
    "        f\"Epoch {epoch + 1}, Val Loss: {val_loss}\"\n",
    "    )\n",
    "torch.save(model, os.path.join('checkpoints', config['encoder']+'_'+str(config['epoch'])+'_'+str(val_loss)+'.pth'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "899950eb-6d62-489c-82b7-271786e77ccd",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = valset[0]\n",
    "if config['encoder'] == 'beats':\n",
    "    feature = BEATs_model.extract_features(test_data[0].to(device), padding_mask=torch.zeros(1, config['input_dim']).bool().to(device))[0]\n",
    "    feature = feature.reshape(1, -1)\n",
    "elif config['encoder'] == 'KAE':\n",
    "    feature = kae_pretrain.encoder(test_data[0].to(device))\n",
    "else:\n",
    "    feature = test_data[0].to(device)\n",
    "pred = model(feature)\n",
    "print(pred, test_data[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5498c206-ad77-47e4-97f3-212dd78abefd",
   "metadata": {},
   "outputs": [],
   "source": [
    "if config['encoder'] == 'beats':\n",
    "    padding_mask = torch.zeros(1, 16384).bool().to(device)\n",
    "    \n",
    "    representation = BEATs_model.extract_features(test_data[0].to(device), padding_mask=padding_mask)[0]\n",
    "    plt.imshow(representation[0].detach().cpu().numpy())\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6081114a-c14d-462d-a25c-db210fec334a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}