tutorial_T7_ICIAP21.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "te_Wwy106o9q"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "66bVc8lM6o9u"
      },
      "outputs": [],
      "source": [
        "import time\n",
        "import numpy as np\n",
        "import torch\n",
        "import torchvision\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from torch.utils.data import random_split\n",
        "from torchvision.transforms import transforms\n",
        "from IPython.display import clear_output\n",
        "from copy import deepcopy\n",
        "from matplotlib import pyplot as plt\n",
        "from torch import Generator\n",
        "from torch import nn\n",
        "from torch.utils.data import random_split, Dataset\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "qkQ6S94_6o9y"
      },
      "source": [
        "# Helper functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "izEa_7Qn6o90"
      },
      "source": [
        "## Dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "RkjfYRQB6o91"
      },
      "outputs": [],
      "source": [
        "class MapDataset(Dataset):\n",
        "    def __init__(self, dataset, map_fn):\n",
        "        self.dataset = dataset\n",
        "        self.map = map_fn\n",
        "    \n",
        "    def __getitem__(self, index):\n",
        "        return self.map(self.dataset[index][0]), self.dataset[index][1]\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.dataset)\n",
        "\n",
        "\n",
        "def split_dataset(dataset, percentage, random_seed=0):\n",
        "    dataset_length = len(dataset)\n",
        "    valid_length = int(np.floor(percentage * dataset_length))\n",
        "    train_length = dataset_length - valid_length\n",
        "    train_dataset, valid_dataset = random_split(dataset, [train_length, valid_length],\n",
        "                                                generator=Generator().manual_seed(random_seed))\n",
        "    \n",
        "    return train_dataset, valid_dataset\n",
        "\n",
        "\n",
        "transf = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])\n",
        "\n",
        "train_dataset = torchvision.datasets.MNIST(\"./data\", train=True, transform=None, download=True)\n",
        "train, validation = split_dataset(train_dataset, 0.1)\n",
        "\n",
        "train, validation = MapDataset(train, transf), MapDataset(validation, transf)\n",
        "train_dataloader = torch.utils.data.DataLoader(train, batch_size=100, shuffle=True, num_workers=2, pin_memory=True,\n",
        "                                               persistent_workers=True)\n",
        "valid_dataloader = torch.utils.data.DataLoader(validation, batch_size=1000, shuffle=False, num_workers=2,\n",
        "                                               pin_memory=True, persistent_workers=True)\n",
        "\n",
        "test_dataset = torchvision.datasets.MNIST(\"./data\", train=False, transform=transf, download=True)\n",
        "test_dataloader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=False, num_workers=2,\n",
        "                                              pin_memory=True, persistent_workers=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "VQC2CFZ66o93"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "HmrQcmyC6o94"
      },
      "outputs": [],
      "source": [
        "def topk_accuracy(outputs, labels, topk=1):\n",
        "    outputs = torch.softmax(outputs, dim=1)\n",
        "    _, preds = outputs.topk(topk, dim=1)\n",
        "    preds = preds.t()\n",
        "    correct = preds.eq(labels.view(1, -1).expand_as(preds)).sum()\n",
        "    return (correct / float(len(outputs))).item()\n",
        "\n",
        "\n",
        "def run(model, dataloader, criterion, optimizer, device, l1_lmbda):\n",
        "    train = optimizer is not None\n",
        "    \n",
        "    tot_loss = 0.\n",
        "    outputs = []\n",
        "    targets = []\n",
        "    \n",
        "    model.train(train)\n",
        "    for data, target in tqdm(dataloader):\n",
        "        data, target = data.to(device), target.to(device)\n",
        "        \n",
        "        with torch.set_grad_enabled(train):\n",
        "            output = model(data)\n",
        "            loss = criterion(output, target)\n",
        "            if l1_lmbda != 0:\n",
        "                loss += l1_lmbda * sum(torch.linalg.vector_norm(p, 1) for p in model.parameters())\n",
        "        \n",
        "        if train:\n",
        "            optimizer.zero_grad()\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "        \n",
        "        tot_loss += loss.item()\n",
        "        outputs.append(output.detach().float())\n",
        "        targets.append(target)\n",
        "    \n",
        "    outputs = torch.cat(outputs, dim=0)\n",
        "    targets = torch.cat(targets, dim=0)\n",
        "    \n",
        "    accs = {\n",
        "        'top1': topk_accuracy(outputs, targets, topk=1) * 100,\n",
        "        'top5': topk_accuracy(outputs, targets, topk=5) * 100\n",
        "    }\n",
        "    return {'loss': tot_loss / len(dataloader.dataset), 'accuracy': accs}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "XnF07VoK6o99"
      },
      "source": [
        "## Pruning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "oOWhy5gj6o9_"
      },
      "outputs": [],
      "source": [
        "def unstructured_magnitude_pruning(model, amount):\n",
        "    for n, m in model.named_modules():\n",
        "        if isinstance(m, (nn.Linear, nn.Conv2d)):\n",
        "            prune.l1_unstructured(m, \"weight\", amount)\n",
        "\n",
        "def structured_magnitude_pruning(model, amount):\n",
        "    layers = [(n, m) for n, m in model.named_modules() if isinstance(m, (nn.Linear, nn.Conv2d))]\n",
        "    for i, (n, m) in enumerate(layers):\n",
        "        if i+1 < len(layers):\n",
        "            prune.ln_structured(m, \"weight\", amount, 1, 0)\n",
        "\n",
        "\n",
        "def finalize_pruning(model):\n",
        "    for n, m in model.named_modules():\n",
        "        if hasattr(m, \"weight_mask\"):\n",
        "            prune.remove(m, \"weight\")\n",
        "\n",
        "\n",
        "def get_pruning_statistics(model):\n",
        "    stats = {}\n",
        "    \n",
        "    for n, m in model.named_modules():\n",
        "        if isinstance(m, (nn.Linear, nn.Conv2d)):\n",
        "            total_weights = m.weight.numel()\n",
        "            total_neurons = m.weight.shape[0]\n",
        "            if isinstance(m, nn.Linear):\n",
        "                non_zero_weights = torch.count_nonzero(m.weight)\n",
        "                non_zero_neurons = torch.sum(torch.abs(m.weight), dim=1).count_nonzero()\n",
        "            elif isinstance(m, nn.Conv2d):\n",
        "                non_zero_weights = torch.count_nonzero(m.weight)\n",
        "                non_zero_neurons = torch.sum(torch.abs(m.weight), dim=(1, 2, 3)).count_nonzero()\n",
        "            \n",
        "            stats[n] = {\"total_weights\":     total_weights,\n",
        "                        \"non_zero_weights\":  non_zero_weights.item(),\n",
        "                        \"zero_weights\":      (total_weights - non_zero_weights).item(),\n",
        "                        \"zero_weights_perc\": (100 - (non_zero_weights / total_weights * 100)).item(),\n",
        "                        \"total_neurons\":     total_neurons,\n",
        "                        \"non_zero_neurons\": non_zero_neurons.item(),\n",
        "                        \"zero_neurons\":      (total_neurons - non_zero_neurons).item(),\n",
        "                        \"zero_neurons_perc\": (100 - (non_zero_neurons / total_neurons * 100)).item()}\n",
        "\n",
        "    return stats"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "KGcp3KZ66o-B"
      },
      "source": [
        "## Plots"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EJWzCjQGpFvK",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "def plot(data, fnc, xlabel, ylabel, title):\n",
        "    legend = False\n",
        "    for x, y, label in data:\n",
        "\n",
        "        if label != \"\":\n",
        "            legend = True\n",
        "\n",
        "        fnc(x, y, label=label)\n",
        "        plt.xlabel(xlabel)\n",
        "        plt.ylabel(ylabel)\n",
        "    plt.title(title)\n",
        "    \n",
        "    if legend:\n",
        "        plt.legend()\n",
        "\n",
        "    plt.show()\n",
        "    \n",
        "def plot_layer_sparsity(layer):\n",
        "    bin_weight = torch.where(layer.weight == 0, torch.zeros_like(layer.weight), torch.ones_like(layer.weight))\n",
        "    \n",
        "    plt.pcolormesh(bin_weight, linewidth=0.5, cmap=\"gray\", vmin=0, vmax=1)\n",
        "    ax = plt.gca()\n",
        "    ax.set_aspect('equal')\n",
        "    plt.xlabel(\"Connections\")\n",
        "    plt.ylabel(\"Neurons\")\n",
        "    plt.tight_layout()\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5fnW594VreaY",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "# Making deep neural networks efficient with PyTorch"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6hxYUxUnoaac",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "State-of-the-art deep learning techniques rely on over-parametrized models that are hard to deploy. On the contrary, biological neural networks are known to use efficient sparse connectivity. Identifying optimal techniques to compress models by reducing the number of parameters in them is important in order to reduce memory, battery, and hardware consumption without sacrificing accuracy, deploy lightweight models on device, and guarantee privacy with private on-device computation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "AXvTcMWz6o-a"
      },
      "source": [
        "![1%20nicFUkeUpWMW1w_hUVtZiw.png](attachment:1%20nicFUkeUpWMW1w_hUVtZiw.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a56FiFMwpJQl",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "In this notebook we will leverage on the PyTorch's `nn.utils.prune` package to show the difference between unstructured and structured sparsification in a one-shot pruning procedure.\n",
        "\n",
        "We will:\n",
        "\n",
        "1. Cover the basics of `torch.nn.utils.prune`.\n",
        "2. Apply unstructured and structured pruning procedures (with and without fine-tuning).\n",
        "3. Introduce `simplify` to exploit the introduced structured sparsity."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CvU6yebLd9cq",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "## Introduction - PyTorch pruning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wZGxUBeIefzm",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "As a first example, we will prune a `Linear` module. Note that all considerations also apply to other module type, e.g. `Conv2d`.\n",
        "\n",
        "Let's define a simple `Linear` module with 10 `in_features` and 2 `out_features`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UOzeTp4AeBoU",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "module = nn.Linear(10, 4)\n",
        "print(module)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cX8OfSkE4nXD",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "We can inspect the module and see that it contains two parameters: `weight` and `bias`. And no buffers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jcAI7E6LebqL",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "for n, p in module.named_parameters():\n",
        "    print(f\"{n}\\n{p}\\n{p.shape}\\n\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WEbknuGffaqN",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "print(list(module.named_buffers()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pvIZNTVngCO8",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Now that we defined our module, let's see how can we prune it using the `nn.utils.prune` package. In this example we will perform a magnitude-based pruning, removing 30% of the module parameters.\n",
        "\n",
        "PyTorch provides different, already implemented, functions. In this case, we will employ the `random_unstructured` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Yvx_t2tQf9XH",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "from torch.nn.utils import prune\n",
        "\n",
        "prune.random_unstructured(module, name=\"weight\", amount=0.3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kYWvLOQMgUN9",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "The pruning procedures preserve the original parameter values, removing the target parameter and replacing it appending `\"_orig\"` to the initial parameter's `name`. In this case, `weight_orig` stores the unpruned version of the tensor. The bias was not pruned, so it will remain intact."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aolFYX3ZgUos",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "for n, p in module.named_parameters():\n",
        "    print(f\"{n}\\n{p}\\n{p.shape}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dGWc3v29g42y",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "The procedure generates a pruning mask (a binary tensor) which is saved as a buffer in the module.\n",
        "\n",
        "The name of the buffer is defined by appending `\"_mask\"` to the initial parameter's `name`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cWtWWtysgglc",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "for n, p in module.named_buffers():\n",
        "    print(f\"{n}\\n{p}\\n{p.shape}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kXJVN4SumbZH",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "For the forward pass to work without modification, the `weight` attribute needs to exist.\n",
        "\n",
        "For this reason, during any forward pass, the pruned `weight` is defined by combining the mask and the corresponding unpruned parameter. This is done using a `forward_pre_hook`. This also means that the pruning is not lost due to the update of the weight.\n",
        "\n",
        "Note that `weight` is no longer a parameter but a simple attribute of the module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ar7FBKhxH6M9",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "print(module.weight)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2igGzcTs3wsW",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "To make the pruning permanent, remove the re-parametrization in terms of `weight_orig` and `weight_mask`, and remove the `forward_pre_hook`, we can use the `remove` functionality from `nn.utils.prune`.\n",
        "\n",
        "Note that this doesn’t undo the pruning, as if it never happened. It simply makes it permanent, instead, by reassigning the parameter `weight` to the model parameters, in its pruned version."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DgFQ2nqh4Wkp",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "prune.remove(module, 'weight')\n",
        "\n",
        "for n, p in module.named_parameters():\n",
        "    print(f\"{n}\\n{p}\\n{p.shape}\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dLgRlSG76o-r"
      },
      "source": [
        "### A visual representation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H9kMUAJ-6o-s"
      },
      "outputs": [],
      "source": [
        "plot_layer_sparsity(module)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ub0Xsgt9TBgY",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "As with most of PyTorch functionalities, we are not limited by the provided functions but we can implement custom procedures by subclassing the `BasePruningMethod` class, implementing the `__init__` and `compute_mask()` functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "54oFMdHi5tIl",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "## Neural network pruning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4UcA_jKHrpur",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Now that we have seen the functionalities of `nn.utils.prune` on a single module, let's use it to prune an entire neural network model.\n",
        "\n",
        "for simplicity's sake, in this tutorial, we will use the [LeNet-5 Caffe architecture](https://caffe.berkeleyvision.org/gathered/examples/mnist.html) and the [MNIST dataset](http://yann.lecun.com/exdb/mnist/).\n",
        "\n",
        "Following a simple one-shot strategy, we will:\n",
        "\n",
        "1. Train a baseline, dense, model.\n",
        "2. Prune the trained models and compare the sparse models accuracy with the baseline obtained in the first step.\n",
        "3. Fine-tune the pruned models and evaluate the (eventual) recover of performance lost during due to the pruning step.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PURU9N3XK4Oy",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "### Define the network architecture"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FRJdlvCPzHza",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "As mentioned, for this example we will use the Caffe version of the LeNet-5 architecture.\n",
        "\n",
        "This model is composed of four layers: two `Conv2d` of 20 and 50 channels respectively, and two `Linear` of 500 and 10 neurons."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "POY0g79ZrebT",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "class LeNet5(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(LeNet5, self).__init__()\n",
        "        \n",
        "        self.conv1 = nn.Conv2d(in_channels=1, out_channels=20, kernel_size=5, stride=1, bias=True)\n",
        "        self.conv2 = nn.Conv2d(in_channels=20, out_channels=50, kernel_size=5, stride=1, bias=True)\n",
        "        self.fc1 = nn.Linear(800, 500, bias=True)\n",
        "        self.fc2 = nn.Linear(500, 10, bias=True)\n",
        "        self.relu = nn.ReLU()\n",
        "    \n",
        "    def forward(self, img):\n",
        "        output = self.relu(self.conv1(img))\n",
        "        output = F.max_pool2d(output, 2)\n",
        "        output = self.relu(self.conv2(output))\n",
        "        output = F.max_pool2d(output, 2)\n",
        "        output = output.view(img.size(0), -1)\n",
        "        output = self.relu(self.fc1(output))\n",
        "        output = self.fc2(output)\n",
        "        \n",
        "        return output"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "AnJlIL6j6o-y"
      },
      "source": [
        "![Immagine.png](attachment:Immagine.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JLz_T0-k9rMi",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "### Training the networks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "89K2fsK9zhoh",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "The trainin procedure is standard: we will train the model using `SGD` with a learning rate of 0.1 and a `cross_entropy` loss."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GHsMxOyjGCsv",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "epochs = 5\n",
        "learning_rate = 0.1\n",
        "\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "model = LeNet5().to(device)\n",
        "\n",
        "criterion = F.cross_entropy\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7d3BHCRZ1MBT",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Let's start with training the baseline model. During training we will keep track of the best weight configuration (highest accuracy), which will be used in the following phases.\n",
        "\n",
        "The network should be able to achieve an accuracy of 99.2% on the test set in around 10 epochs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RAHObfR86bUZ",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "test_acc = []\n",
        "\n",
        "dense_trained_model_test_acc = 0\n",
        "\n",
        "for epoch in range(epochs):\n",
        "    train = run(model, train_dataloader, criterion, optimizer, device, 0)\n",
        "    test = run(model, test_dataloader, criterion, None, device, 0)\n",
        "    test_acc.append(test[\"accuracy\"][\"top1\"])\n",
        "\n",
        "    if test[\"accuracy\"][\"top1\"] > dense_trained_model_test_acc:\n",
        "        dense_trained_model_test_acc = test[\"accuracy\"][\"top1\"]\n",
        "        dense_trained_model = deepcopy(model.state_dict())\n",
        "\n",
        "    clear_output(wait=True)\n",
        "    plot([(range(0, epoch+1), test_acc, \"\")],\n",
        "         plt.plot, \"epochs\", \"accuracy\", \"\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "erUsvkAc6f9c",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "### Unstructured pruning without fine-tuning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oaf5oEKt8cMw",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Now that we obtained our trained network, we can move on to the pruning stage. Here we will evaluate the performance of the model for different percentages of pruned parameters.\n",
        "\n",
        "To prune the network we will use the magnitude-based approach implemented in the `l1_unstructured` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1qJHArMKfJVJ",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "test_acc = []\n",
        "amounts = [0.5, 0.6, 0.7, 0.8, 0.81, 0.82, 0.83, 0.84, 0.85, 0.86, 0.87, 0.88, 0.89, 0.9,]\n",
        "\n",
        "for i, amount in enumerate(amounts):\n",
        "    model.load_state_dict(dense_trained_model)\n",
        "\n",
        "    unstructured_magnitude_pruning(model, amount)\n",
        "    finalize_pruning(model)\n",
        "    \n",
        "    test = run(model, test_dataloader, criterion, None, device, 0)\n",
        "    test_acc.append(test[\"accuracy\"][\"top1\"] - dense_trained_model_test_acc)\n",
        "\n",
        "    clear_output(wait=True)\n",
        "    plot([(amounts[:i+1], test_acc, \"\")],\n",
        "         plt.scatter, \"pruning %\", \"accuracy difference\", \"\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Q6mjNaTXmGl",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "We can see that the performance degrades as we remove paramameters from the network. The more parameters we remove the more the accuracy drops."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w1BMHsDe6kyw",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "### Unstructured pruning with fine-tuning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XRCUee2k8v8v",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Let's introduce a fine-tuning step after the pruning stage. Thanks to this new step, we can recover some of the accuracy lost when pruning the parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EdjUpR9S6rMK",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "test_acc_ft = []\n",
        "fine_tune_epochs = 1\n",
        "fine_tune_lr = 0.1\n",
        "\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=fine_tune_lr)\n",
        "for i, amount in enumerate(amounts):\n",
        "    model.load_state_dict(dense_trained_model)\n",
        "\n",
        "    unstructured_magnitude_pruning(model, amount)\n",
        "\n",
        "    for epoch in range(fine_tune_epochs):\n",
        "        run(model, train_dataloader, criterion, optimizer, device, 0)\n",
        "\n",
        "    finalize_pruning(model)\n",
        "    \n",
        "    test = run(model, test_dataloader, criterion, None, device, 0)\n",
        "    test_acc_ft.append(test[\"accuracy\"][\"top1\"] - dense_trained_model_test_acc)\n",
        "\n",
        "    clear_output(wait=True)\n",
        "    plot([(amounts[:i+1], test_acc_ft, \"W/ fine-tuning\"), (amounts[:i+1], test_acc[:i+1], \"W/o fine-tuning\")],\n",
        "            plt.scatter, \"pruning %\", \"accuracy difference\", \"W/o vs W/ fine-tuning\")\n",
        "\n",
        "clear_output(wait=True) \n",
        "plot([(amounts, test_acc_ft, \"W/ fine-tuning\"), (amounts, test_acc, \"W/o fine-tuning\")],\n",
        "        plt.scatter, \"pruning %\", \"accuracy difference\", \"W/o vs W/ fine-tuning\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5045h8GOaC1F",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Comparing the accuracy lost before and after the fine-tuning step, we can see that this step is quite beneficial: the drop in accuracy after fine-tuning is way less severe."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m-lKC2fg9IHr",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "### Structured pruning (with fine-tuning)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-X7BShsG9a2M",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Up untill now, we focused on unstructured pruning procedure, in which we remove parameters independently.\n",
        "\n",
        "Now we will introduce, in our network, a structured sparsification. To this end, we will use the provided `ln_unstructured` function, with the L1-norm along the 0-th dimension (the one representing the neurons of the module).\n",
        "\n",
        "Structured pruning poses a stronger constraint on the network and the loss in performance is usually grater, even if the amount of pruned weights is roughly the same (pruning 50% of the neurons prunes around 50% of the parameters)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nVGhFPvj6o-9"
      },
      "source": [
        "### A visual representation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dV29abnQr56G",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "unstructured = nn.Linear(128, 64)\n",
        "structured = nn.Linear(128, 64)\n",
        "\n",
        "prune.random_unstructured(unstructured, name=\"weight\", amount=0.5)\n",
        "prune.random_structured(structured, name=\"weight\", amount=0.5, dim=0)\n",
        "\n",
        "print(\"Unstructured\")\n",
        "plot_layer_sparsity(unstructured)\n",
        "print(\"Structured\")\n",
        "plot_layer_sparsity(structured)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qeGs2LCT6o--"
      },
      "source": [
        "### Performance comparison"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vWDN1I8i9KYk",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "test_acc_ft_st = []\n",
        "\n",
        "fine_tune_epochs = 1\n",
        "fine_tune_lr = 0.1\n",
        "\n",
        "model = LeNet5().to(device)\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=fine_tune_lr)\n",
        "\n",
        "for i, amount in enumerate(amounts):\n",
        "    model.load_state_dict(dense_trained_model)\n",
        "\n",
        "    structured_magnitude_pruning(model, amount)\n",
        "\n",
        "    for epoch in range(fine_tune_epochs):\n",
        "        run(model, train_dataloader, criterion, optimizer, device, 0)\n",
        "\n",
        "    finalize_pruning(model)\n",
        "    \n",
        "    test = run(model, test_dataloader, criterion, None, device, 0)\n",
        "    test_acc_ft_st.append(test[\"accuracy\"][\"top1\"] - dense_trained_model_test_acc)\n",
        "\n",
        "    clear_output(wait=True)\n",
        "    plot([(amounts[:i+1], test_acc_ft_st, \"Structured\"), (amounts[:i+1], test_acc_ft[:i+1], \"Unstructured\")],\n",
        "         plt.scatter, \"pruning %\", \"accuracy difference\", \"Structured vs Unstructured pruning\")\n",
        "\n",
        "\n",
        "clear_output(wait=True)\n",
        "plot([(amounts, test_acc_ft_st, \"Structured\"), (amounts, test_acc_ft, \"Unstructured\")],\n",
        "     plt.scatter, \"pruning %\", \"accuracy difference\", \"Structured vs Unstructured pruning\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9F2a_nmPLoqk",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "### Exploiting structured sparsity"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pbEwzot0cLBs",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "If structured sparsity is harder to achive, why are we interested in it?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lLT2A-yWAd_j",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "The strenght of structured sparsity lays in the ability to actually shrink the resulting network, removing enitre rows (or columns) from the matrices that represent the layers. This allows us to produce smaller and faster models without the need for ad hoc software or hardware able to run sparse tensors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "79a-g-sIL2UI",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "[Simplify](https://github.com/EIDOSlab/simplify) is an open source PyTorch-compatible library. It allows us to exploit the introduced structured sparsity by removing zeroed neurons from the network architecture.\n",
        "\n",
        "It can be installed via pip:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ug9UU69MxatT",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "!pip3 install torch-simplify"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7way3nu26o_F"
      },
      "source": [
        "#### Why simplify?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QsmQkfSG6o_G"
      },
      "source": [
        "![Immagine.png](attachment:Immagine.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SFoiexux6o_G"
      },
      "source": [
        "**Bias propagation**: pruned neurons may still have non-zero bias, they result in a costant output value and cannot be removed. Such value can be accumulated into the biases of the next layers, allowing us to simplify the neurons. Linear layer and convolutional layers without padding are trivial, but padded convolution lead to non-scalar biases, simplify handles them.\n",
        "\n",
        "**Residual connections**: such connection consists in the sum of the outputs of two branched layers. The sum pose a constraint: both tensors not only must have the same dimensions but also channels must be pruned at the same index. Pruning procedures do not usually ensure such conditions. Simplify handles residual connections by re-expanding the layer output before returning it (the operation is still performed on smaller matrices).\n",
        "\n",
        "![Immagine.png](attachment:Immagine.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "AGI9NaOx6o_I"
      },
      "source": [
        "#### Simplification visualized"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "Mc7eh-G86o_I"
      },
      "outputs": [],
      "source": [
        "import simplify\n",
        "\n",
        "module = nn.Linear(128, 64)\n",
        "data = torch.randn(1, 128)\n",
        "\n",
        "print(\"Dense module\")\n",
        "plot_layer_sparsity(module)\n",
        "\n",
        "unstructured = deepcopy(module)\n",
        "structured = deepcopy(module)\n",
        "\n",
        "print(\"Unstructured sparity\")\n",
        "prune.l1_unstructured(unstructured, \"weight\", 0.5)\n",
        "prune.remove(unstructured, \"weight\")\n",
        "plot_layer_sparsity(unstructured)\n",
        "\n",
        "print(\"Structured sparity\")\n",
        "prune.ln_structured(structured, \"weight\", 0.5, 1, 0)\n",
        "prune.remove(structured, \"weight\")    \n",
        "plot_layer_sparsity(structured)\n",
        "\n",
        "print(\"Simplified module\")\n",
        "simplified_module = simplify.simplify(structured, data, [], False, [])\n",
        "plot_layer_sparsity(simplified_module)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "5RB_TSgU6o_J"
      },
      "source": [
        "#### Simplification applications"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sLKbu101jMba",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "Let's see how we can use the structured sparsity of a `nn.Conv2d` module in which we prune 50% of the parameters.\n",
        "\n",
        "In this example we will compare the inference speed of the dense module, the pruned module (both unstructured and structured), and the simplified module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UTLnE8849NN5",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "module = nn.Conv2d(20, 50, 5)\n",
        "data = torch.randn(10, 20, 24, 24)\n",
        "\n",
        "unstructured = deepcopy(module)\n",
        "structured = deepcopy(module)\n",
        "\n",
        "prune.l1_unstructured(unstructured, \"weight\", 0.5)\n",
        "prune.remove(unstructured, \"weight\")\n",
        "\n",
        "prune.ln_structured(structured, \"weight\", 0.5, 1, 0)\n",
        "prune.remove(structured, \"weight\")\n",
        "\n",
        "dense_times = []\n",
        "for i in tqdm(range(10000)):\n",
        "    start = time.perf_counter_ns()\n",
        "    module(data)\n",
        "    dense_times.append((time.perf_counter_ns() - start) / 1e+6)\n",
        "\n",
        "unstructured_times = []\n",
        "for i in tqdm(range(10000)):\n",
        "    start = time.perf_counter_ns()\n",
        "    unstructured(data)\n",
        "    unstructured_times.append((time.perf_counter_ns() - start) / 1e+6)\n",
        "\n",
        "structured_times = []\n",
        "for i in tqdm(range(10000)):\n",
        "    start = time.perf_counter_ns()\n",
        "    structured(data)\n",
        "    structured_times.append((time.perf_counter_ns() - start) / 1e+6)\n",
        "\n",
        "simplified_module = simplify.simplify(structured, torch.randn(1, 20, 24, 24), [], False, [])\n",
        "\n",
        "simplified_times = []\n",
        "for i in tqdm(range(10000)):\n",
        "    start = time.perf_counter_ns()\n",
        "    simplified_module(data)\n",
        "    simplified_times.append((time.perf_counter_ns() - start) / 1e+6)\n",
        "    \n",
        "print(f\"Dense module inference time {np.mean(dense_times[1:])}ms pm {np.std(dense_times[1:])}\")\n",
        "print(f\"Unstructured module inference time {np.mean(unstructured_times[1:])}ms pm {np.std(unstructured_times[1:])}\")\n",
        "print(f\"Structured module inference time {np.mean(structured_times[1:])}ms pm {np.std(structured_times[1:])}\")\n",
        "print(f\"Simplified module inference time {np.mean(simplified_times[1:])}ms pm {np.std(simplified_times[1:])}\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ovu8fCu1jrax",
        "pycharm": {
          "name": "#%% md\n"
        }
      },
      "source": [
        "What about and entire neural network? Again, we will use the LeNet-5 model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KU3JbRMtj3Zm",
        "pycharm": {
          "name": "#%%\n"
        }
      },
      "outputs": [],
      "source": [
        "data = torch.randn(10, 1, 28, 28)\n",
        "\n",
        "model = LeNet5()\n",
        "structured_magnitude_pruning(model, 0.5)\n",
        "finalize_pruning(model)\n",
        "\n",
        "simplified_model = simplify.simplify(deepcopy(model), torch.randn(1, 1, 28, 28))\n",
        "\n",
        "print(model)\n",
        "print(simplified_model)\n",
        "\n",
        "pruned_times = []\n",
        "for i in tqdm(range(10000)):\n",
        "    start = time.perf_counter_ns()\n",
        "    model(data)\n",
        "    pruned_times.append((time.perf_counter_ns() - start) / 1e+6)\n",
        "\n",
        "simplified_times = []\n",
        "for i in tqdm(range(10000)):\n",
        "    start = time.perf_counter_ns()\n",
        "    simplified_model(data)\n",
        "    simplified_times.append((time.perf_counter_ns() - start) / 1e+6)\n",
        "\n",
        "print(f\"Pruned model inference time {np.mean(pruned_times[1:])}ms pm {np.std(pruned_times[1:])}\")\n",
        "print(f\"Simplified model inference time {np.mean(simplified_times[1:])}ms pm {np.std(simplified_times[1:])}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vi67IQA26o_N"
      },
      "source": [
        "## Going further"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZGuvfyO6o_N"
      },
      "source": [
        "With this notebook we have seen how to prune neural networks with PyTorch.\n",
        "\n",
        "Here we provide some ideas to delve deeper into the topic:\n",
        "- Harder tasks (e.g. [CIFAR-10](https://pytorch.org/vision/stable/generated/torchvision.datasets.CIFAR10.html)), deeper models (e.g. [ResNet-32](https://github.com/akamaster/pytorch_resnet_cifar10))\n",
        "- Evaluate the effect of regularization on the pruning procedure (e.g. L2 or L1 norms)\n",
        "- Improve the pruning procedure (e.g. find the best threshold with a validation set)\n",
        "- Create your own pruning policy, subclassing the `BasePruningMethod` class, implementing the `__init__` and `compute_mask()` functions (more info [here](https://pytorch.org/tutorials/intermediate/pruning_tutorial.html#extending-torch-nn-utils-prune-with-custom-pruning-functions))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "tutorial_T7_ICIAP21.ipynb",
      "provenance": []
    },
    "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.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}