examples/PyTorch/MNIST_pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import pandas as pd\n",
    "import seaborn as sn\n",
    "\n",
    "import torch\n",
    "from pytorch_lightning import LightningModule, Trainer\n",
    "from pytorch_lightning.callbacks.progress import TQDMProgressBar\n",
    "from pytorch_lightning.loggers import CSVLogger\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader, random_split\n",
    "from torchmetrics import Accuracy\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import MNIST\n",
    "from torch.nn import functional as F\n",
    "\n",
    "PATH_DATASETS = os.environ.get(\"PATH_DATASETS\", \".\")\n",
    "BATCH_SIZE = 256 if torch.cuda.is_available() else 64\n",
    "\n",
    "import sys\n",
    "sys.path.append(\"../../src\")\n",
    "from concretedropout.pytorch import ConcreteDropout, ConcreteDropout2D, get_weight_regularizer, get_dropout_regularizer\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the model and utility functions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MNISTModel(LightningModule):\n",
    "    def __init__(self, weight_regularizer, dropout_regularizer):\n",
    "        super().__init__()\n",
    "\n",
    "        self.weight_regularizer = weight_regularizer\n",
    "        self.dropout_regularizer = dropout_regularizer\n",
    "\n",
    "        self.data_dir = PATH_DATASETS\n",
    "\n",
    "        self.transform = transforms.Compose(\n",
    "            [\n",
    "                transforms.ToTensor(),\n",
    "                transforms.Normalize((0.1307,), (0.3081,)),\n",
    "            ]\n",
    "        )\n",
    "\n",
    "\n",
    "        self.conv1 = nn.Conv2d(1, 32, 3, padding=\"same\")\n",
    "        self.cd1 = ConcreteDropout2D(weight_regularizer=self.weight_regularizer, dropout_regularizer=self.dropout_regularizer)\n",
    "        self.maxpool1 = nn.MaxPool2d(2)\n",
    "        self.activation1 = nn.ReLU()\n",
    "\n",
    "        self.conv2 = nn.Conv2d(32, 64, 3, padding=\"same\")\n",
    "        self.cd2 = ConcreteDropout2D(weight_regularizer=self.weight_regularizer, dropout_regularizer=self.dropout_regularizer)\n",
    "        self.maxpool2 = nn.MaxPool2d(2)\n",
    "        self.activation2 = nn.ReLU()\n",
    "\n",
    "        self.flatten = nn.Flatten()\n",
    "\n",
    "        self.l1 = torch.nn.Linear(7 * 7*64, 10)\n",
    "\n",
    "        self.val_accuracy = Accuracy()\n",
    "        self.test_accuracy = Accuracy()\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.cd1(x, self.conv1)\n",
    "        x = self.activation1(x)\n",
    "        x = self.maxpool1(x)\n",
    "\n",
    "        x = self.cd2(x, self.conv2)\n",
    "        x = self.activation2(x)\n",
    "        x = self.maxpool2(x)\n",
    "\n",
    "        x = self.flatten(x)\n",
    "\n",
    "        x = self.l1(x)\n",
    "        return F.log_softmax(x, dim=1)\n",
    "        \n",
    "\n",
    "    def training_step(self, batch, batch_nb):\n",
    "        x, y = batch\n",
    "        logits = self(x)\n",
    "        \n",
    "        reg = torch.zeros(1)\n",
    "        for module in filter(lambda x: isinstance(x, ConcreteDropout), self.modules()):\n",
    "            reg += module.regularization\n",
    "\n",
    "        loss = F.nll_loss(logits, y) + reg\n",
    "        return loss\n",
    "\n",
    "    def validation_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        logits = self(x)\n",
    "        loss = F.nll_loss(logits, y)\n",
    "        preds = torch.argmax(logits, dim=1)\n",
    "        self.val_accuracy.update(preds, y)\n",
    "\n",
    "        # Calling self.log will surface up scalars for you in TensorBoard\n",
    "        self.log(\"val_loss\", loss, prog_bar=True)\n",
    "        self.log(\"val_acc\", self.val_accuracy, prog_bar=True)\n",
    "\n",
    "    def test_step(self, batch, batch_idx):\n",
    "        x, y = batch\n",
    "        logits = self(x)\n",
    "        loss = F.nll_loss(logits, y)\n",
    "        preds = torch.argmax(logits, dim=1)\n",
    "        self.test_accuracy.update(preds, y)\n",
    "\n",
    "        # Calling self.log will surface up scalars for you in TensorBoard\n",
    "        self.log(\"test_loss\", loss, prog_bar=True)\n",
    "        self.log(\"test_acc\", self.test_accuracy, prog_bar=True)\n",
    "\n",
    "    def configure_optimizers(self):\n",
    "        return torch.optim.Adam(self.parameters(), lr=0.02)\n",
    "\n",
    "\n",
    "    def prepare_data(self):\n",
    "        # download\n",
    "        MNIST(self.data_dir, train=True, download=True)\n",
    "        MNIST(self.data_dir, train=False, download=True)\n",
    "\n",
    "    def setup(self, stage=None):\n",
    "\n",
    "        # Assign train/val datasets for use in dataloaders\n",
    "        if stage == \"fit\" or stage is None:\n",
    "            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)\n",
    "            self.mnist_train, self.mnist_val = random_split(mnist_full, [55000, 5000])\n",
    "\n",
    "        # Assign test dataset for use in dataloader(s)\n",
    "        if stage == \"test\" or stage is None:\n",
    "            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)\n",
    "\n",
    "    def train_dataloader(self):\n",
    "        return DataLoader(self.mnist_train, batch_size=BATCH_SIZE, num_workers=4)\n",
    "\n",
    "    def val_dataloader(self):\n",
    "        return DataLoader(self.mnist_val, batch_size=BATCH_SIZE, num_workers=4)\n",
    "\n",
    "    def test_dataloader(self):\n",
    "        return DataLoader(self.mnist_test, batch_size=BATCH_SIZE, num_workers=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "wr = get_weight_regularizer(55000, l=1e-2, tau=0.1) # tau and l may need tuning, depending on the dataset\n",
    "dr = get_dropout_regularizer(55000, tau=0.1, cross_entropy_loss=True)\n",
    "\n",
    "mnist_model = MNISTModel(wr, dr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True, used: False\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n",
      "/home/aure/miniconda3/envs/pytorch/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1814: PossibleUserWarning: GPU available but not used. Set `accelerator` and `devices` using `Trainer(accelerator='gpu', devices=1)`.\n",
      "  rank_zero_warn(\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    accelerator=\"cpu\",\n",
    "    devices=1 if torch.cuda.is_available() else None,  # limiting got iPython runs\n",
    "    max_epochs=3,\n",
    "    callbacks=[TQDMProgressBar(refresh_rate=20)],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz to ./MNIST/raw/train-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53e2da8f36a2493998ec51bac004258d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/9912422 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/train-images-idx3-ubyte.gz to ./MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ./MNIST/raw/train-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ee492ac5f37443a68cceeb6625393216",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/28881 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/train-labels-idx1-ubyte.gz to ./MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz to ./MNIST/raw/t10k-images-idx3-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e6f0455a6b7c4807908f6f30a7e4ab10",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1648877 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/t10k-images-idx3-ubyte.gz to ./MNIST/raw\n",
      "\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz\n",
      "Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz to ./MNIST/raw/t10k-labels-idx1-ubyte.gz\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0e790e1af3f64a2096ec46933cd02cf9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/4542 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extracting ./MNIST/raw/t10k-labels-idx1-ubyte.gz to ./MNIST/raw\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Missing logger folder: /mnt/d/Aure/Documenti/Github/ConcreteDropout-TF2/examples/PyTorch/lightning_logs\n",
      "\n",
      "   | Name          | Type              | Params\n",
      "-----------------------------------------------------\n",
      "0  | conv1         | Conv2d            | 320   \n",
      "1  | cd1           | ConcreteDropout2D | 1     \n",
      "2  | maxpool1      | MaxPool2d         | 0     \n",
      "3  | activation1   | ReLU              | 0     \n",
      "4  | conv2         | Conv2d            | 18.5 K\n",
      "5  | cd2           | ConcreteDropout2D | 1     \n",
      "6  | maxpool2      | MaxPool2d         | 0     \n",
      "7  | activation2   | ReLU              | 0     \n",
      "8  | flatten       | Flatten           | 0     \n",
      "9  | l1            | Linear            | 31.4 K\n",
      "10 | val_accuracy  | Accuracy          | 0     \n",
      "11 | test_accuracy | Accuracy          | 0     \n",
      "-----------------------------------------------------\n",
      "50.2 K    Trainable params\n",
      "0         Non-trainable params\n",
      "50.2 K    Total params\n",
      "0.201     Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "58b73e69055b4e99ae2a2f5fd70ce65f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c1fc04aef8234cfb915c21bf5696b553",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c8d500d63f154680a795dee6e0bd9a08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b5f72a3f8674427d9f31f7fa346aff84",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02c61b9e96d244d8b8f96bf979aeb038",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.fit(mnist_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Restoring states from the checkpoint path at /mnt/d/Aure/Documenti/Github/ConcreteDropout-TF2/examples/PyTorch/lightning_logs/version_0/checkpoints/epoch=2-step=645.ckpt\n",
      "Loaded model weights from checkpoint at /mnt/d/Aure/Documenti/Github/ConcreteDropout-TF2/examples/PyTorch/lightning_logs/version_0/checkpoints/epoch=2-step=645.ckpt\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4bb72db7f4e4ca684abadc61ba26382",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Testing: 0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "       Test metric             DataLoader 0\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n",
      "        test_acc            0.9722999930381775\n",
      "        test_loss           0.0905378982424736\n",
      "────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'test_loss': 0.0905378982424736, 'test_acc': 0.9722999930381775}]"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.test(ckpt_path='best')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "learned dropout probabilities:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.00322718]\n",
      "[0.00941805]\n"
     ]
    }
   ],
   "source": [
    "for module in mnist_model.modules():\n",
    "    if isinstance(module, ConcreteDropout):\n",
    "        print(module.p.detach().numpy())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 ('pytorch')",
   "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.4"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "cecee814e95025ef937c557eba7fe28b2e5cf5b5fa19c11f76b27bd724f0265f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}