liset_implementation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LiSET\n",
    "\n",
    "Versions of libraries used:\n",
    "- `torch==1.9.0`\n",
    "- `torchvision==0.9.1+cu102`\n",
    "- `numpy==1.18.4`\n",
    "- `tqdm==4.60.0`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "from torchvision.datasets import MNIST\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "\n",
    "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Datasets\n",
    "\n",
    "- MNIST - provided by torchvision https://pytorch.org/vision/0.8/datasets.html#mnist\n",
    "- ISOLET - https://archive.ics.uci.edu/ml/datasets/ISOLET"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATABASE = 'MNIST'\n",
    "batch_size = 100\n",
    "\n",
    "if DATABASE == 'MNIST':\n",
    "    transform = transforms.Compose([transforms.ToTensor()])\n",
    "\n",
    "    train_set = MNIST(\"./\", train=True, download=True, transform=transform)\n",
    "    test_set = MNIST(\"./\", train=False, download=True, transform=transform)\n",
    "\n",
    "    train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True)\n",
    "    test_loader = torch.utils.data.DataLoader(test_set, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "elif DATABASE == \"ISOLET\":\n",
    "    class TrainLoader(Dataset):\n",
    "        def __init__(self):\n",
    "            data = pd.read_csv(\"uci-isolet/original/isolet1+2+3+4.data\", index_col=False).to_numpy(dtype=np.float)\n",
    "            self.x, self.y = torch.from_numpy(data[:, :-1]).float(), torch.from_numpy(data[:, -1] - 1).type(torch.LongTensor)\n",
    "\n",
    "        def __getitem__(self, idx):\n",
    "            return self.x[idx], self.y[idx]\n",
    "\n",
    "        def __len__(self):\n",
    "            return len(self.x)\n",
    "\n",
    "    class TestLoader(Dataset):\n",
    "        def __init__(self):\n",
    "            data = pd.read_csv(\"uci-isolet/original/isolet5.data\", index_col=False).to_numpy(dtype=np.float)\n",
    "            self.x, self.y = torch.from_numpy(data[:, :-1]).float(), torch.from_numpy(data[:, -1] - 1).type(torch.LongTensor)\n",
    "\n",
    "        def __getitem__(self, idx):\n",
    "            return self.x[idx], self.y[idx]\n",
    "\n",
    "        def __len__(self):\n",
    "            return len(self.x)\n",
    "\n",
    "    train_loader = DataLoader(TrainLoader(), batch_size=100, shuffle=True)\n",
    "    test_loader = DataLoader(TestLoader(), batch_size=100, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Function used to initialize spare matrices"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_sparse_weights(epsilon, rows, cols):\n",
    "    prob = epsilon * (rows + cols) / (rows * cols)\n",
    "    mask_weights = (torch.rand(cols, rows) < prob).float()\n",
    "    print(f\"Created {cols} x {rows} sparse matrix with {torch.sum(mask_weights)} parameters, {rows} rows, {cols} cols\")\n",
    "    return mask_weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementation of LiSET for MLP\n",
    "`PyTorch provides ambiguous support for unstructured sparsity, but it has not been used in this implementation. Sparsity was achieved by applying mask on matrices.`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LiSET_MLP(nn.Module):\n",
    "    def __init__(self, *dims, EPSILON=20, ZETA=0.3, GAMMA=0, device=\"cpu\"):\n",
    "        assert len(dims) > 2\n",
    "        super().__init__()\n",
    "        \n",
    "        self.dims = dims\n",
    "        self.ZETA = ZETA\n",
    "        self.GAMMA = GAMMA\n",
    "        self.device=device\n",
    "        \n",
    "        self.layers = nn.ModuleList([nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)]).to(device)\n",
    "        self.skip_layers = nn.ModuleList([nn.Linear(dims[i], dims[-1]) for i in range(len(dims) - 1)]).to(device)\n",
    "        self.masks = [create_sparse_weights(EPSILON, dims[i], dims[i + 1]).to(device) for i in range(len(dims) - 2)]\n",
    "        self.apply_masks()\n",
    "    \n",
    "    def forward(self, x, dropout=0.3, with_skip=True, skip_dropout=0.5):\n",
    "        self.apply_masks()\n",
    "        x = torch.flatten(x, 1)\n",
    "        \n",
    "        if with_skip:\n",
    "            result = torch.zeros(x.shape[0], self.dims[-1]).to(x)\n",
    "        for i, l in enumerate(self.layers[:-1]):\n",
    "            if with_skip and np.random.rand() > skip_dropout:\n",
    "                result += self.skip_layers[i](x)\n",
    "            \n",
    "            x = F.leaky_relu(l(x))\n",
    "            if dropout > 0:\n",
    "                x = F.dropout(x)\n",
    "        \n",
    "        x = self.layers[-1](x)\n",
    "        return F.log_softmax(x + result if with_skip else x, dim=1)\n",
    "    \n",
    "    def apply_mask(self, layer, mask):\n",
    "        assert layer.weight.data.shape == mask.shape\n",
    "        with torch.no_grad():\n",
    "            layer.weight.data *= mask\n",
    "    \n",
    "    def apply_masks(self):\n",
    "        for l, m in zip(self.layers[:-1], self.masks):\n",
    "            self.apply_mask(l, m)\n",
    "    \n",
    "    def remove_weights(self, weights):\n",
    "        sorted_weights = weights.reshape(-1).sort().values\n",
    "        \n",
    "        negative_weights = torch.masked_select(sorted_weights, sorted_weights < -0)\n",
    "        largest_negative = negative_weights[int((1 - self.ZETA) * len(negative_weights))]\n",
    "        \n",
    "        positive_weights = torch.masked_select(sorted_weights, sorted_weights > 0)\n",
    "        smallest_positive = positive_weights[int(self.ZETA * len(positive_weights))]\n",
    "        \n",
    "        weight_mask_core = ((weights < largest_negative) + (weights > smallest_positive)).float()\n",
    "        n_removed = ((torch.abs(weights) > 0).float() - weight_mask_core).nonzero().shape[0]\n",
    "        \n",
    "        return weight_mask_core, n_removed\n",
    "    \n",
    "    def add_weights(self, mask, n_to_add, skip):\n",
    "        w = torch.norm(skip, p=1, dim=0)\n",
    "        w_star = torch.sigmoid(w / torch.max(w) * self.GAMMA)\n",
    "        \n",
    "        probs = (mask == 0) * torch.rand(mask.shape).to(self.device) * w_star\n",
    "        new_weights = probs > torch.sort(probs.reshape(-1), descending=True).values[n_to_add]\n",
    "        \n",
    "        return new_weights.float()\n",
    "    \n",
    "    def evolution(self):\n",
    "        for layer, skip, mask in zip(self.layers[:-1], self.skip_layers, self.masks):\n",
    "            self.apply_mask(layer, mask)\n",
    "            mask_core, n_removed = self.remove_weights(layer.weight.data)\n",
    "            new_weights_mask = self.add_weights(mask_core, n_removed, skip.weight.data)\n",
    "            mask = mask_core + new_weights_mask\n",
    "            \n",
    "            limit = np.sqrt(1. / float(layer.weight.data.shape[0]))\n",
    "            with torch.no_grad():\n",
    "                layer.weight.data = layer.weight.data * mask_core + \\\n",
    "                                    new_weights_mask * torch.empty(mask.shape).to(self.device).uniform_(-limit, limit)\n",
    "                skip.weight.data = torch.empty(skip.weight.data.shape).to(self.device).uniform_(-limit/100, limit/100)\n",
    "    \n",
    "    def calculate_sparsity_percentage(self):\n",
    "        self.apply_masks()\n",
    "        zero_weights = 0\n",
    "        total_weights = 0\n",
    "        for l in self.layers[:-1]:\n",
    "            vals = l.weight.data.reshape(-1)\n",
    "            zero_weights += torch.sum(vals == 0).item()\n",
    "            total_weights += vals.shape[0]\n",
    "        return round(zero_weights / total_weights * 100, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = LiSET_MLP(28 * 28, 300, 100, 10, EPSILON=20, GAMMA=0, device=device)\n",
    "net.calculate_sparsity_percentage()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "EPOCHS = 30\n",
    "\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.01)\n",
    "\n",
    "pbar = tqdm(range(EPOCHS))\n",
    "for epoch in pbar:\n",
    "    for data in train_loader:\n",
    "        X, y = data[0].to(device), data[1].to(device)\n",
    "        net.zero_grad()\n",
    "        output = net(X)\n",
    "        loss = F.nll_loss(output, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    if epoch < (EPOCHS - 1):\n",
    "        net.evolution()\n",
    "    \n",
    "    pbar.set_description(f\"Loss {round(loss.item(), 4)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "correct = 0\n",
    "total = 0\n",
    "with torch.no_grad():\n",
    "    for data in test_loader:\n",
    "        X, y = data[0].to(device), data[1].to(device)\n",
    "        output = net(X, dropout=0, with_skip=False)\n",
    "        for idx, i in enumerate(output):\n",
    "            if torch.argmax(i) == y[idx]:\n",
    "                correct += 1\n",
    "            total += 1\n",
    "\n",
    "print(f\"Accuracy: {round(correct / total * 100, 3)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## Function of proposed evaluation metric to calculate speed up in convergence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_ara(accuracy_list, b=0.5):\n",
    "    n = len(accuracy_list)\n",
    "    s = 0\n",
    "    \n",
    "    a_const = (n - n * b) / (n - 1)\n",
    "    b_const = n - a_const * n\n",
    "    for i in range(1, n + 1):\n",
    "        s += accuracy_list[i - 1] * (i * a_const + b_const)\n",
    "    return s / (n ** 2 * (b + 1) / 2)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.7 64-bit",
   "language": "python",
   "name": "python37764bit2a47edd992914d51b00bb3d1afbed8ae"
  },
  "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}