trainer.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torchaudio\n",
    "import numpy as np\n",
    "import random\n",
    "import os\n",
    "import librosa\n",
    "from torchvision import models\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def omitHiddenFiles(inpArray):\n",
    "    return [x for x in inpArray if '.' != x[0]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, chunkedNoisePath, datasetPath,\n",
    "                 training_dataset=True,\n",
    "                 max_noise_factor=0.2,\n",
    "                 min_noise_factor=0.05,\n",
    "                 sampling_rate=16000,\n",
    "                 spectrogram=True):\n",
    "        self.spectrogram = spectrogram\n",
    "        self.sampling_rate = sampling_rate\n",
    "        self.chunkedNoisePath = chunkedNoisePath\n",
    "        self.typesOfNoise = omitHiddenFiles(os.listdir(self.chunkedNoisePath))\n",
    "        self.datasetPath = datasetPath\n",
    "        self.max_noise_factor = max_noise_factor\n",
    "        self.min_noise_factor = min_noise_factor\n",
    "        self.wordsInDataset = omitHiddenFiles(os.listdir(datasetPath))\n",
    "        countOfWords = len(self.wordsInDataset)\n",
    "\n",
    "        if training_dataset:\n",
    "            self.wordsInDataset = self.wordsInDataset[:int(0.9 * countOfWords)]\n",
    "            print(\"Train size: \", len(self.wordsInDataset))\n",
    "        else:\n",
    "            self.wordsInDataset = self.wordsInDataset[int(0.9 * countOfWords):]\n",
    "            print(\"Test size: \", len(self.wordsInDataset))\n",
    "\n",
    "        # Shuffle the dataset\n",
    "        random.shuffle(self.wordsInDataset)\n",
    "\n",
    "        self.n = 2 * len(self.wordsInDataset)  # Positive and negative pairs\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.n\n",
    "\n",
    "    def fixPaddingIssues(self, x, length=16000):\n",
    "        if len(x) > length:\n",
    "            start = random.randint(0, len(x) - length)\n",
    "            return x[start:start+length]\n",
    "        elif len(x) < length:\n",
    "            return np.pad(x, (0, length - len(x)), 'constant')\n",
    "        else:\n",
    "            return x\n",
    "\n",
    "    def addNoise(self, x, noise, noise_factor=0.4):\n",
    "        out = (1-noise_factor)*x/np.max(np.abs(x)) + noise_factor*(noise/np.max(np.abs(noise)))\n",
    "        return out/np.max(np.abs(out))\n",
    "\n",
    "    def giveJoinedAudio(self, word1: str, word2: str):\n",
    "        if word1 == word2:\n",
    "            sample1, sample2 = random.sample(omitHiddenFiles(os.listdir(self.datasetPath + \"/\" + word1)), 2)\n",
    "        else:\n",
    "            sample1 = random.choice(omitHiddenFiles(os.listdir(self.datasetPath + \"/\" + word1)))\n",
    "            sample2 = random.choice(omitHiddenFiles(os.listdir(self.datasetPath + \"/\" + word2)))\n",
    "\n",
    "        voiceVector1, _ = librosa.load(self.datasetPath + \"/\" + word1 + \"/\" + sample1, sr=self.sampling_rate)\n",
    "        voiceVector2, _ = librosa.load(self.datasetPath + \"/\" + word2 + \"/\" + sample2, sr=self.sampling_rate)\n",
    "\n",
    "        # Padding\n",
    "        voiceVector1 = self.fixPaddingIssues(voiceVector1)\n",
    "        voiceVector2 = self.fixPaddingIssues(voiceVector2)\n",
    "\n",
    "        # Noise generation\n",
    "        randomNoiseType1, randomNoiseType2 = random.sample(self.typesOfNoise, 2)\n",
    "        randomNoise1 = random.choice(omitHiddenFiles(os.listdir(self.chunkedNoisePath + \"/\" + randomNoiseType1 + \"/\")))\n",
    "        randomNoise2 = random.choice(omitHiddenFiles(os.listdir(self.chunkedNoisePath + \"/\" + randomNoiseType2 + \"/\")))\n",
    "\n",
    "        noiseVector1, _ = librosa.load(self.chunkedNoisePath + \"/\" + randomNoiseType1 + \"/\" + randomNoise1, sr=self.sampling_rate)\n",
    "        noiseVector2, _ = librosa.load(self.chunkedNoisePath + \"/\" + randomNoiseType2 + \"/\" + randomNoise2, sr=self.sampling_rate)\n",
    "\n",
    "        # Noise factors\n",
    "        randomNoiseFactor1 = random.uniform(self.min_noise_factor, self.max_noise_factor)\n",
    "        randomNoiseFactor2 = random.uniform(self.min_noise_factor, self.max_noise_factor)\n",
    "\n",
    "        voice_with_noise1 = self.addNoise(voiceVector1, noiseVector1, randomNoiseFactor1)\n",
    "        voice_with_noise2 = self.addNoise(voiceVector2, noiseVector2, randomNoiseFactor2)\n",
    "\n",
    "        if self.spectrogram:\n",
    "            melspec = torchaudio.transforms.MelSpectrogram(\n",
    "                sample_rate=self.sampling_rate,\n",
    "                n_fft=400,\n",
    "                hop_length=160,\n",
    "                n_mels=64\n",
    "            )\n",
    "            voice_with_noise_spectrogram1 = melspec(torch.from_numpy(voice_with_noise1).float())\n",
    "            voice_with_noise_spectrogram2 = melspec(torch.from_numpy(voice_with_noise2).float())\n",
    "            \n",
    "            voice_with_noise_spectrogram1 = (torch.log(voice_with_noise_spectrogram1 + 1e-9) - 1.4) / 1.184\n",
    "            voice_with_noise_spectrogram2 = (torch.log(voice_with_noise_spectrogram2 + 1e-9) - 1.4) / 1.184\n",
    "\n",
    "            return torch.stack([voice_with_noise_spectrogram1, voice_with_noise_spectrogram2])\n",
    "\n",
    "        return torch.from_numpy(np.array([voice_with_noise1, voice_with_noise2])).float()\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        x_data = self.giveJoinedAudio(\n",
    "            self.wordsInDataset[idx // 2],\n",
    "            self.wordsInDataset[(idx // 2 + idx % 2) % (self.n // 2)]\n",
    "        )\n",
    "        y = torch.tensor(1.0 if idx % 2 == 0 else 0.0, dtype=torch.float32)\n",
    "\n",
    "        return x_data[0], x_data[1], y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SiameseNetwork(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SiameseNetwork, self).__init__()\n",
    "        \n",
    "        # Load pre-trained ResNet50\n",
    "        resnet = models.resnet50(pretrained=True)\n",
    "        \n",
    "        # Remove the last fully connected layer\n",
    "        self.features = nn.Sequential(*list(resnet.children())[:-1])\n",
    "        \n",
    "        # Add custom layers\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(2048, 128),\n",
    "            nn.BatchNorm1d(128),\n",
    "        )\n",
    "\n",
    "    def forward_one(self, x):\n",
    "        # Replicate grayscale input to 3 channels\n",
    "        x = x.repeat(1, 3, 1, 1)\n",
    "        \n",
    "        x = self.features(x)\n",
    "        x = x.view(x.size()[0], -1)\n",
    "        x = self.fc(x)\n",
    "        x = nn.functional.normalize(x, p=2, dim=1)\n",
    "        return x\n",
    "\n",
    "    def forward(self, input1, input2):\n",
    "        output1 = self.forward_one(input1)\n",
    "        output2 = self.forward_one(input2)\n",
    "        return output1, output2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def triplet_loss(y_true, y_pred):\n",
    "    distance = torch.norm(y_pred[0] - y_pred[1], dim=1)\n",
    "    match_loss = y_true * -2.0 * torch.log(1 - distance/2)\n",
    "    mismatch_loss = torch.clamp((1 - y_true) * (-torch.log(distance/0.2)), min=0)\n",
    "    return torch.mean(match_loss + mismatch_loss)\n",
    "\n",
    "def accuracy(y_true, y_pred):\n",
    "    distance = torch.norm(y_pred[0] - y_pred[1], dim=1)\n",
    "    threshold_check = (distance <= 0.2).float()\n",
    "    return torch.mean((threshold_check == y_true).float())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_loader, val_loader, num_epochs, device):\n",
    "    optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', patience=2, factor=0.1, min_lr=1e-6)\n",
    "    \n",
    "    for epoch in range(num_epochs):\n",
    "        model.train()\n",
    "        train_loss = 0.0\n",
    "        train_acc = 0.0\n",
    "        \n",
    "        # Use tqdm for the progress bar\n",
    "        pbar = tqdm(train_loader, desc=f'Epoch {epoch+1}/{num_epochs}')\n",
    "        for batch_idx, (data1, data2, target) in enumerate(pbar):\n",
    "            data1, data2, target = data1.to(device), data2.to(device), target.to(device)\n",
    "            \n",
    "            optimizer.zero_grad()\n",
    "            output1, output2 = model(data1, data2)\n",
    "            loss = triplet_loss(target, (output1, output2))\n",
    "            acc = accuracy(target, (output1, output2))\n",
    "            \n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "            train_loss += loss.item()\n",
    "            train_acc += acc.item()\n",
    "            \n",
    "            # Update progress bar\n",
    "            pbar.set_postfix({'loss': f'{loss.item():.4f}', 'acc': f'{acc.item():.4f}'})\n",
    "        \n",
    "        train_loss /= len(train_loader)\n",
    "        train_acc /= len(train_loader)\n",
    "        \n",
    "        # Validation\n",
    "        model.eval()\n",
    "        val_loss = 0.0\n",
    "        val_acc = 0.0\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            for data1, data2, target in val_loader:\n",
    "                data1, data2, target = data1.to(device), data2.to(device), target.to(device)\n",
    "                output1, output2 = model(data1, data2)\n",
    "                loss = triplet_loss(target, (output1, output2))\n",
    "                acc = accuracy(target, (output1, output2))\n",
    "                \n",
    "                val_loss += loss.item()\n",
    "                val_acc += acc.item()\n",
    "        \n",
    "        val_loss /= len(val_loader)\n",
    "        val_acc /= len(val_loader)\n",
    "        \n",
    "        print(f'Epoch {epoch+1}/{num_epochs}:')\n",
    "        print(f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}')\n",
    "        print(f'Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}')\n",
    "        \n",
    "        scheduler.step(val_loss)\n",
    "        \n",
    "        # Save the model\n",
    "        torch.save(model.state_dict(), f'model_epoch_{epoch+1}_val_acc_{val_acc:.3f}.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train size:  21157\n",
      "Test size:  2351\n",
      "Using device: cuda\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\salos\\OneDrive\\Desktop\\EfficientWord-Net\\venv\\lib\\site-packages\\torchvision\\models\\_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "c:\\Users\\salos\\OneDrive\\Desktop\\EfficientWord-Net\\venv\\lib\\site-packages\\torchvision\\models\\_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet50_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet50_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n",
      "Epoch 1/10:   0%|          | 0/21 [00:00<?, ?it/s]"
     ]
    }
   ],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    chunkedNoisePath = r\"C:\\Users\\salos\\OneDrive\\Desktop\\EfficientWord-Net\\Efficient_word_net\\NoiseChunked\"\n",
    "    datasetPath = r\"C:\\Users\\salos\\OneDrive\\Desktop\\EfficientWord-Net\\Efficient_word_net\\test\"\n",
    "    \n",
    "    train_dataset = CustomDataset(chunkedNoisePath, datasetPath, training_dataset=True)\n",
    "    test_dataset = CustomDataset(chunkedNoisePath, datasetPath, training_dataset=False)\n",
    "    \n",
    "    batch_size = 2048  # You can try even larger values like 1024 or 2048\n",
    "    \n",
    "    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)\n",
    "    \n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    print(f\"Using device: {device}\")\n",
    "\n",
    "    model = SiameseNetwork().to(device)\n",
    "    \n",
    "    # Enable cuDNN benchmark for faster training\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "    \n",
    "    train(model, train_loader, test_loader, num_epochs=10, device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}