soap-binary-classifier.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "DATA_PATH = \"data/soap/soap-80-20-20/\""
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "396a511760b83dd8"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "\n",
    "from aivolved.dataset import get_img_dataset_normalisation\n",
    "\n",
    "train_dataset = torchvision.datasets.ImageFolder(\n",
    "    DATA_PATH + \"train\",\n",
    "    transform=torchvision.transforms.Compose([\n",
    "        torchvision.transforms.Resize((224, 224)),\n",
    "    ])\n",
    ")\n",
    "\n",
    "mean, stdev = get_img_dataset_normalisation(train_dataset)"
   ],
   "metadata": {
    "collapsed": false
   },
   "id": "c68f586ca97fec90"
  },
  {
   "cell_type": "code",
   "execution_count": 237,
   "outputs": [
    {
     "data": {
      "text/plain": "Dataset ImageFolder\n    Number of datapoints: 757\n    Root location: data/soap/soap-80-20-20/train\n    StandardTransform\nTransform: Compose(\n               Resize(size=(224, 224), interpolation=bilinear, max_size=None, antialias=True)\n               AutoAugment(policy=AutoAugmentPolicy.CIFAR10, fill=None)\n               ToTensor()\n               Normalize(mean=[0.24232754 0.29568848 0.3936678 ], std=0.18121950328350067)\n           )"
     },
     "execution_count": 237,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_transforms = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize((224, 224)),\n",
    "    torchvision.transforms.AutoAugment(policy=torchvision.transforms.AutoAugmentPolicy.CIFAR10),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize(mean, stdev)\n",
    "])\n",
    "\n",
    "test_transforms = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.Resize((224, 224)),\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize(mean, stdev)\n",
    "])\n",
    "\n",
    "train_dataset = torchvision.datasets.ImageFolder(DATA_PATH + \"train\", transform=train_transforms)\n",
    "valid_dataset = torchvision.datasets.ImageFolder(DATA_PATH + \"val\", transform=train_transforms)\n",
    "test_dataset = torchvision.datasets.ImageFolder(DATA_PATH + \"test\", transform=test_transforms)\n",
    "train_dataset"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-30T18:54:24.671083Z",
     "start_time": "2024-06-30T18:54:24.660322Z"
    }
   },
   "id": "94a036377f031b30"
  },
  {
   "cell_type": "code",
   "execution_count": 261,
   "outputs": [],
   "source": [
    "from torch.utils.data import DataLoader\n",
    "\n",
    "WORKERS = 8\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, batch_size=128, num_workers=WORKERS, pin_memory=True, shuffle=True, persistent_workers=True)\n",
    "valid_dataloader = DataLoader(valid_dataset, batch_size=256, num_workers=WORKERS, pin_memory=True, persistent_workers=True)\n",
    "test_dataloader = DataLoader(test_dataset, batch_size=len(test_dataset), num_workers=WORKERS, persistent_workers=True)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-30T19:04:31.756125Z",
     "start_time": "2024-06-30T19:04:31.748012Z"
    }
   },
   "id": "f0a4cdc8c341d560"
  },
  {
   "cell_type": "code",
   "execution_count": 254,
   "outputs": [],
   "source": [
    "import lightning as pl\n",
    "\n",
    "\n",
    "class SoapBinaryClassifier(pl.LightningModule):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.resnet = torchvision.models.resnet34(weights=torchvision.models.ResNet34_Weights.DEFAULT)\n",
    "        self.criterion = torch.nn.CrossEntropyLoss()\n",
    "        \n",
    "        self.resnet.fc = torch.nn.Linear(self.resnet.fc.in_features, 2)\n",
    "        torch.nn.init.xavier_uniform_(self.resnet.fc.weight)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return self.resnet(x)\n",
    "\n",
    "    def training_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self.forward(x)\n",
    "        \n",
    "        loss = self.criterion(y_hat, y)\n",
    "        acc = 100 * (torch.argmax(y_hat, dim=1) == y).sum() / len(x)\n",
    "        self.log(\"train_loss\", loss)\n",
    "        self.log(\"train_acc\", acc)\n",
    "        \n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self.forward(x)\n",
    "\n",
    "        loss = self.criterion(y_hat, y)\n",
    "        acc = 100 * (torch.argmax(y_hat, dim=1) == y).sum() / len(x)\n",
    "        \n",
    "        self.log(\"val_loss\", loss)\n",
    "        self.log(\"val_acc\", acc)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        y_hat = self.forward(x)\n",
    "        \n",
    "        loss = self.criterion(y_hat, y)\n",
    "        acc = 100 * (torch.argmax(y_hat, dim=1) == y).sum() / len(x)\n",
    "        \n",
    "        print(f\"Test loss: {loss:.4f}\")\n",
    "        print(f\"Test accuracy: {acc:.2f}%\")\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        lr = 1e-5\n",
    "        weight_decay = 5e-4\n",
    "        \n",
    "        params_1x = [param for name, param in self.resnet.named_parameters() if \"fc\" not in str(name)]\n",
    "        optimizer = torch.optim.Adam([{\"params\": params_1x}, {\"params\": self.resnet.fc.parameters(), \"lr\": lr*10}], lr=lr, weight_decay=weight_decay)\n",
    "        return optimizer\n"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-30T19:02:41.998361Z",
     "start_time": "2024-06-30T19:02:41.994366Z"
    }
   },
   "id": "e2d04576af093730"
  },
  {
   "cell_type": "code",
   "execution_count": 255,
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (mps), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "HPU available: False, using: 0 HPUs\n",
      "\n",
      "  | Name      | Type             | Params | Mode \n",
      "-------------------------------------------------------\n",
      "0 | resnet    | ResNet           | 21.3 M | train\n",
      "1 | criterion | CrossEntropyLoss | 0      | train\n",
      "-------------------------------------------------------\n",
      "21.3 M    Trainable params\n",
      "0         Non-trainable params\n",
      "21.3 M    Total params\n",
      "85.143    Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "text/plain": "Sanity Checking: |          | 0/? [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "8a228f9d37c44c758f5f1f5f8c6547eb"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Training: |          | 0/? [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "47679ca80eff4c73879cd4dad3381c97"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Validation: |          | 0/? [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "177454cf56b1441abb4389f982218d62"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Validation: |          | 0/? [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "649f849f2763450c9bae0a38630aabf9"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": "Validation: |          | 0/? [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "85cecf59549a40598b40e4efa233b463"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=3` reached.\n"
     ]
    }
   ],
   "source": [
    "model = SoapBinaryClassifier()\n",
    "trainer = pl.Trainer(max_epochs=3, log_every_n_steps=1)\n",
    "trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=valid_dataloader)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-30T19:03:06.003254Z",
     "start_time": "2024-06-30T19:02:42.902147Z"
    }
   },
   "id": "a7aaba094ad98e77"
  },
  {
   "cell_type": "code",
   "execution_count": 262,
   "outputs": [
    {
     "data": {
      "text/plain": "Testing: |          | 0/? [00:00<?, ?it/s]",
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "7dd7dd0e7731459ab19d7f9589d3ce54"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test loss: 0.0202\n",
      "Test accuracy: 100.00%\n"
     ]
    },
    {
     "data": {
      "text/plain": "[{}]"
     },
     "execution_count": 262,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(model, dataloaders=test_dataloader)"
   ],
   "metadata": {
    "collapsed": false,
    "ExecuteTime": {
     "end_time": "2024-06-30T19:04:39.295131Z",
     "start_time": "2024-06-30T19:04:35.429613Z"
    }
   },
   "id": "cdefbb6e8ea24bcf"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}