pruning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "pruning.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOB6vekC4J9sqoWTqa09QVL"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "893cWqVUMkK7"
      },
      "source": [
        "import math\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "from torch.nn import Parameter\n",
        "from torch.nn.modules.module import Module\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class PruningModule(Module):\n",
        "    def prune_by_percentile(self, q=5.0, **kwargs):\n",
        "        \"\"\"\n",
        "        Note:\n",
        "             The pruning percentile is based on all layer's parameters concatenated\n",
        "        Args:\n",
        "            q (float): percentile in float\n",
        "            **kwargs: may contain `cuda`\n",
        "        \"\"\"\n",
        "        # Calculate percentile value\n",
        "        alive_parameters = []\n",
        "        for name, p in self.named_parameters():\n",
        "            # We do not prune bias term\n",
        "            if 'bias' in name or 'mask' in name:\n",
        "                continue\n",
        "            tensor = p.data.cpu().numpy()\n",
        "            alive = tensor[np.nonzero(tensor)] # flattened array of nonzero values\n",
        "            alive_parameters.append(alive)\n",
        "\n",
        "        all_alives = np.concatenate(alive_parameters)\n",
        "        percentile_value = np.percentile(abs(all_alives), q)\n",
        "        print(f'Pruning with threshold : {percentile_value}')\n",
        "\n",
        "        # Prune the weights and mask\n",
        "        # Note that module here is the layer\n",
        "        # ex) fc1, fc2, fc3\n",
        "        for name, module in self.named_modules():\n",
        "            if name in ['fc1', 'fc2', 'fc3']:\n",
        "                module.prune(threshold=percentile_value)\n",
        "\n",
        "    def prune_by_std(self, s=0.25):\n",
        "        \"\"\"\n",
        "        Note that `s` is a quality parameter / sensitivity value according to the paper.\n",
        "        According to Song Han's previous paper (Learning both Weights and Connections for Efficient Neural Networks),\n",
        "        'The pruning threshold is chosen as a quality parameter multiplied by the standard deviation of a layer’s weights'\n",
        "        I tried multiple values and empirically, 0.25 matches the paper's compression rate and number of parameters.\n",
        "        Note : In the paper, the authors used different sensitivity values for different layers.\n",
        "        \"\"\"\n",
        "        for name, module in self.named_modules():\n",
        "            if name in ['fc1', 'fc2', 'fc3']:\n",
        "                threshold = np.std(module.weight.data.cpu().numpy()) * s\n",
        "                print(f'Pruning with threshold : {threshold} for layer {name}')\n",
        "                module.prune(threshold)\n",
        "\n",
        "\n",
        "class MaskedLinear(Module):\n",
        "    r\"\"\"Applies a masked linear transformation to the incoming data: :math:`y = (A * M)x + b`\n",
        "    Args:\n",
        "        in_features: size of each input sample\n",
        "        out_features: size of each output sample\n",
        "        bias: If set to False, the layer will not learn an additive bias.\n",
        "            Default: ``True``\n",
        "    Shape:\n",
        "        - Input: :math:`(N, *, in\\_features)` where `*` means any number of\n",
        "          additional dimensions\n",
        "        - Output: :math:`(N, *, out\\_features)` where all but the last dimension\n",
        "          are the same shape as the input.\n",
        "    Attributes:\n",
        "        weight: the learnable weights of the module of shape\n",
        "            (out_features x in_features)\n",
        "        bias:   the learnable bias of the module of shape (out_features)\n",
        "        mask: the unlearnable mask for the weight.\n",
        "            It has the same shape as weight (out_features x in_features)\n",
        "    \"\"\"\n",
        "    def __init__(self, in_features, out_features, bias=True):\n",
        "        super(MaskedLinear, self).__init__()\n",
        "        self.in_features = in_features\n",
        "        self.out_features = out_features\n",
        "        self.weight = Parameter(torch.Tensor(out_features, in_features))\n",
        "        # Initialize the mask with 1\n",
        "        self.mask = Parameter(torch.ones([out_features, in_features]), requires_grad=False)\n",
        "        if bias:\n",
        "            self.bias = Parameter(torch.Tensor(out_features))\n",
        "        else:\n",
        "            self.register_parameter('bias', None)\n",
        "        self.reset_parameters()\n",
        "\n",
        "    def reset_parameters(self):\n",
        "        stdv = 1. / math.sqrt(self.weight.size(1))\n",
        "        self.weight.data.uniform_(-stdv, stdv)\n",
        "        if self.bias is not None:\n",
        "            self.bias.data.uniform_(-stdv, stdv)\n",
        "\n",
        "    def forward(self, input):\n",
        "        return F.linear(input, self.weight * self.mask, self.bias)\n",
        "\n",
        "    def __repr__(self):\n",
        "        return self.__class__.__name__ + '(' \\\n",
        "            + 'in_features=' + str(self.in_features) \\\n",
        "            + ', out_features=' + str(self.out_features) \\\n",
        "            + ', bias=' + str(self.bias is not None) + ')'\n",
        "\n",
        "    def prune(self, threshold):\n",
        "        weight_dev = self.weight.device\n",
        "        mask_dev = self.mask.device\n",
        "        # Convert Tensors to numpy and calculate\n",
        "        tensor = self.weight.data.cpu().numpy()\n",
        "        mask = self.mask.data.cpu().numpy()\n",
        "        new_mask = np.where(abs(tensor) < threshold, 0, mask)\n",
        "        # Apply new weight and mask\n",
        "        self.weight.data = torch.from_numpy(tensor * new_mask).to(weight_dev)\n",
        "        self.mask.data = torch.from_numpy(new_mask).to(mask_dev)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Epp7dwI1RNsR"
      },
      "source": [
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "\n",
        "class LeNet(PruningModule):\n",
        "    def __init__(self, mask=False):\n",
        "        super(LeNet, self).__init__()\n",
        "        linear = MaskedLinear if mask else nn.Linear\n",
        "        self.fc1 = linear(784, 300)\n",
        "        self.fc2 = linear(300, 100)\n",
        "        self.fc3 = linear(100, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = x.view(-1, 784)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = F.relu(self.fc2(x))\n",
        "        x = F.log_softmax(self.fc3(x), dim=1)\n",
        "        return x\n",
        "\n",
        "\n",
        "class LeNet_5(PruningModule):\n",
        "    def __init__(self, mask=False):\n",
        "        super(LeNet_5, self).__init__()\n",
        "        linear = MaskedLinear if mask else Linear\n",
        "        self.conv1 = nn.Conv2d(1, 6, kernel_size=(5, 5))\n",
        "        self.conv2 = nn.Conv2d(6, 16, kernel_size=(5, 5))\n",
        "        self.conv3 = nn.Conv2d(16, 120, kernel_size=(5,5))\n",
        "        self.fc1 = linear(120, 84)\n",
        "        self.fc2 = linear(84, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        # Conv1\n",
        "        x = self.conv1(x)\n",
        "        x = F.relu(x)\n",
        "        x = F.max_pool2d(x, kernel_size=(2, 2), stride=2)\n",
        "\n",
        "        # Conv2\n",
        "        x = self.conv2(x)\n",
        "        x = F.relu(x)\n",
        "        x = F.max_pool2d(x, kernel_size=(2, 2), stride=2)\n",
        "\n",
        "        # Conv3\n",
        "        x = self.conv3(x)\n",
        "        x = F.relu(x)\n",
        "\n",
        "        # Fully-connected\n",
        "        x = x.view(-1, 120)\n",
        "        x = self.fc1(x)\n",
        "        x = F.relu(x)\n",
        "        x = self.fc2(x)\n",
        "        x = F.log_softmax(x, dim=1)\n",
        "\n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s2JNqe52RPbv"
      },
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "from sklearn.cluster import KMeans\n",
        "from scipy.sparse import csc_matrix, csr_matrix\n",
        "\n",
        "\n",
        "def apply_weight_sharing(model, bits=5):\n",
        "    \"\"\"\n",
        "    Applies weight sharing to the given model\n",
        "    \"\"\"\n",
        "    for module in model.children():\n",
        "        dev = module.weight.device\n",
        "        weight = module.weight.data.cpu().numpy()\n",
        "        shape = weight.shape\n",
        "        mat = csr_matrix(weight) if shape[0] < shape[1] else csc_matrix(weight)\n",
        "        min_ = min(mat.data)\n",
        "        max_ = max(mat.data)\n",
        "        space = np.linspace(min_, max_, num=2**bits)\n",
        "        kmeans = KMeans(n_clusters=len(space), init=space.reshape(-1,1), n_init=1, precompute_distances=True, algorithm=\"full\")\n",
        "        kmeans.fit(mat.data.reshape(-1,1))\n",
        "        new_weight = kmeans.cluster_centers_[kmeans.labels_].reshape(-1)\n",
        "        mat.data = new_weight\n",
        "        module.weight.data = torch.from_numpy(mat.toarray()).to(dev)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fr_AGMpFSgHS"
      },
      "source": [
        "import os\n",
        "import torch\n",
        "import math\n",
        "import numpy as np\n",
        "from torch.nn import Parameter\n",
        "from torch.nn.modules.module import Module\n",
        "import torch.nn.functional as F\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "def log(filename, content):\n",
        "    with open(filename, 'a') as f:\n",
        "        content += \"\\n\"\n",
        "        f.write(content)\n",
        "\n",
        "\n",
        "def print_model_parameters(model, with_values=False):\n",
        "    print(f\"{'Param name':20} {'Shape':30} {'Type':15}\")\n",
        "    print('-'*70)\n",
        "    for name, param in model.named_parameters():\n",
        "        print(f'{name:20} {str(param.shape):30} {str(param.dtype):15}')\n",
        "        if with_values:\n",
        "            print(param)\n",
        "\n",
        "\n",
        "def print_nonzeros(model):\n",
        "    nonzero = total = 0\n",
        "    for name, p in model.named_parameters():\n",
        "        if 'mask' in name:\n",
        "            continue\n",
        "        tensor = p.data.cpu().numpy()\n",
        "        nz_count = np.count_nonzero(tensor)\n",
        "        total_params = np.prod(tensor.shape)\n",
        "        nonzero += nz_count\n",
        "        total += total_params\n",
        "        print(f'{name:20} | nonzeros = {nz_count:7} / {total_params:7} ({100 * nz_count / total_params:6.2f}%) | total_pruned = {total_params - nz_count :7} | shape = {tensor.shape}')\n",
        "    print(f'alive: {nonzero}, pruned : {total - nonzero}, total: {total}, Compression rate : {total/nonzero:10.2f}x  ({100 * (total-nonzero) / total:6.2f}% pruned)')\n",
        "\n",
        "\n",
        "def test(model, use_cuda=True):\n",
        "    kwargs = {'num_workers': 5, 'pin_memory': True} if use_cuda else {}\n",
        "    device = torch.device(\"cuda\" if use_cuda else 'cpu')\n",
        "    test_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('data', train=False, transform=transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       transforms.Normalize((0.1307,), (0.3081,))\n",
        "                   ])),\n",
        "    batch_size=1000, shuffle=False, **kwargs)\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "            test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss\n",
        "            pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability\n",
        "            correct += pred.eq(target.data.view_as(pred)).sum().item()\n",
        "\n",
        "        test_loss /= len(test_loader.dataset)\n",
        "        accuracy = 100. * correct / len(test_loader.dataset)\n",
        "        print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')\n",
        "    return accuracy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TKcz00XlRybB",
        "outputId": "5e0fa4a3-7c28-4f0e-e0a7-55f38e2802d2",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "import argparse\n",
        "import os\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torch.optim as optim\n",
        "from torchvision import datasets, transforms\n",
        "from tqdm import tqdm\n",
        "\n",
        "\n",
        "os.makedirs('saves', exist_ok=True)\n",
        "\n",
        "# Training settings\n",
        "parser = argparse.ArgumentParser(description='PyTorch MNIST pruning from deep compression paper')\n",
        "parser.add_argument('-f')\n",
        "parser.add_argument('--batch-size', type=int, default=50, metavar='N',\n",
        "                    help='input batch size for training (default: 50)')\n",
        "parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',\n",
        "                    help='input batch size for testing (default: 1000)')\n",
        "parser.add_argument('--epochs', type=int, default=15, metavar='N',\n",
        "                    help='number of epochs to train (default: 100)')\n",
        "parser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n",
        "                    help='learning rate (default: 0.01)')\n",
        "parser.add_argument('--no-cuda', action='store_true', default=False,\n",
        "                    help='disables CUDA training')\n",
        "parser.add_argument('--seed', type=int, default=42, metavar='S',\n",
        "                    help='random seed (default: 42)')\n",
        "parser.add_argument('--log-interval', type=int, default=10, metavar='N',\n",
        "                    help='how many batches to wait before logging training status')\n",
        "parser.add_argument('--log', type=str, default='log.txt',\n",
        "                    help='log file name')\n",
        "parser.add_argument('--sensitivity', type=float, default=2,\n",
        "                    help=\"sensitivity value that is multiplied to layer's std in order to get threshold value\")\n",
        "args = parser.parse_args()\n",
        "\n",
        "# Control Seed\n",
        "torch.manual_seed(args.seed)\n",
        "\n",
        "# Select Device\n",
        "use_cuda = not args.no_cuda and torch.cuda.is_available()\n",
        "device = torch.device(\"cuda\" if use_cuda else 'cpu')\n",
        "if use_cuda:\n",
        "    print(\"Using CUDA!\")\n",
        "    torch.cuda.manual_seed(args.seed)\n",
        "else:\n",
        "    print('Not using CUDA!!!')\n",
        "\n",
        "# Loader\n",
        "kwargs = {'num_workers': 5, 'pin_memory': True} if use_cuda else {}\n",
        "train_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('data', train=True, download=True,\n",
        "                   transform=transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       transforms.Normalize((0.1307,), (0.3081,))\n",
        "                   ])),\n",
        "    batch_size=args.batch_size, shuffle=True, **kwargs)\n",
        "test_loader = torch.utils.data.DataLoader(\n",
        "    datasets.MNIST('data', train=False, transform=transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       transforms.Normalize((0.1307,), (0.3081,))\n",
        "                   ])),\n",
        "    batch_size=args.test_batch_size, shuffle=False, **kwargs)\n",
        "\n",
        "\n",
        "# Define which model to use\n",
        "model = LeNet(mask=True).to(device)\n",
        "\n",
        "print(model)\n",
        "print_model_parameters(model)\n",
        "\n",
        "# NOTE : `weight_decay` term denotes L2 regularization loss term\n",
        "optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=0.0001)\n",
        "initial_optimizer_state_dict = optimizer.state_dict()\n",
        "\n",
        "def train(epochs):\n",
        "    model.train()\n",
        "    for epoch in range(epochs):\n",
        "        pbar = tqdm(enumerate(train_loader), total=len(train_loader))\n",
        "        for batch_idx, (data, target) in pbar:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            optimizer.zero_grad()\n",
        "            output = model(data)\n",
        "            loss = F.nll_loss(output, target)\n",
        "            loss.backward()\n",
        "\n",
        "            # zero-out all the gradients corresponding to the pruned connections\n",
        "            for name, p in model.named_parameters():\n",
        "                if 'mask' in name:\n",
        "                    continue\n",
        "                tensor = p.data.cpu().numpy()\n",
        "                grad_tensor = p.grad.data.cpu().numpy()\n",
        "                grad_tensor = np.where(tensor==0, 0, grad_tensor)\n",
        "                p.grad.data = torch.from_numpy(grad_tensor).to(device)\n",
        "\n",
        "            optimizer.step()\n",
        "            if batch_idx % args.log_interval == 0:\n",
        "                done = batch_idx * len(data)\n",
        "                percentage = 100. * batch_idx / len(train_loader)\n",
        "                pbar.set_description(f'Train Epoch: {epoch} [{done:5}/{len(train_loader.dataset)} ({percentage:3.0f}%)]  Loss: {loss.item():.6f}')\n",
        "\n",
        "\n",
        "def test():\n",
        "    model.eval()\n",
        "    test_loss = 0\n",
        "    correct = 0\n",
        "    with torch.no_grad():\n",
        "        for data, target in test_loader:\n",
        "            data, target = data.to(device), target.to(device)\n",
        "            output = model(data)\n",
        "            test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss\n",
        "            pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability\n",
        "            correct += pred.eq(target.data.view_as(pred)).sum().item()\n",
        "\n",
        "        test_loss /= len(test_loader.dataset)\n",
        "        accuracy = 100. * correct / len(test_loader.dataset)\n",
        "        print(f'Test set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')\n",
        "    return accuracy\n",
        "\n",
        "\n",
        "# Initial training\n",
        "print(\"--- Initial training ---\")\n",
        "train(args.epochs)\n",
        "accuracy = test()\n",
        "log(args.log, f\"initial_accuracy {accuracy}\")\n",
        "torch.save(model, f\"saves/initial_model.ptmodel\")\n",
        "print(\"--- Before pruning ---\")\n",
        "print_nonzeros(model)\n",
        "\n",
        "# Pruning\n",
        "model.prune_by_std(args.sensitivity)\n",
        "accuracy = test()\n",
        "log(args.log, f\"accuracy_after_pruning {accuracy}\")\n",
        "print(\"--- After pruning ---\")\n",
        "print_nonzeros(model)\n",
        "\n",
        "# Retrain\n",
        "print(\"--- Retraining ---\")\n",
        "optimizer.load_state_dict(initial_optimizer_state_dict) # Reset the optimizer\n",
        "train(args.epochs)\n",
        "torch.save(model, f\"saves/model_after_retraining.ptmodel\")\n",
        "accuracy = test()\n",
        "log(args.log, f\"accuracy_after_retraining {accuracy}\")\n",
        "\n",
        "print(\"--- After Retraining ---\")\n",
        "print_nonzeros(model)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 0 [    0/60000 (  0%)]  Loss: 2.275712:   0%|          | 6/1200 [00:00<00:21, 55.46it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Not using CUDA!!!\n",
            "LeNet(\n",
            "  (fc1): MaskedLinear(in_features=784, out_features=300, bias=True)\n",
            "  (fc2): MaskedLinear(in_features=300, out_features=100, bias=True)\n",
            "  (fc3): MaskedLinear(in_features=100, out_features=10, bias=True)\n",
            ")\n",
            "Param name           Shape                          Type           \n",
            "----------------------------------------------------------------------\n",
            "fc1.weight           torch.Size([300, 784])         torch.float32  \n",
            "fc1.mask             torch.Size([300, 784])         torch.float32  \n",
            "fc1.bias             torch.Size([300])              torch.float32  \n",
            "fc2.weight           torch.Size([100, 300])         torch.float32  \n",
            "fc2.mask             torch.Size([100, 300])         torch.float32  \n",
            "fc2.bias             torch.Size([100])              torch.float32  \n",
            "fc3.weight           torch.Size([10, 100])          torch.float32  \n",
            "fc3.mask             torch.Size([10, 100])          torch.float32  \n",
            "fc3.bias             torch.Size([10])               torch.float32  \n",
            "--- Initial training ---\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 0 [59500/60000 ( 99%)]  Loss: 0.349057: 100%|██████████| 1200/1200 [00:18<00:00, 63.45it/s]\n",
            "Train Epoch: 1 [59500/60000 ( 99%)]  Loss: 0.542076: 100%|██████████| 1200/1200 [00:19<00:00, 61.72it/s]\n",
            "Train Epoch: 2 [59500/60000 ( 99%)]  Loss: 0.314782: 100%|██████████| 1200/1200 [00:19<00:00, 60.96it/s]\n",
            "Train Epoch: 3 [59500/60000 ( 99%)]  Loss: 0.144507: 100%|██████████| 1200/1200 [00:19<00:00, 61.17it/s]\n",
            "Train Epoch: 4 [59500/60000 ( 99%)]  Loss: 0.309548: 100%|██████████| 1200/1200 [00:20<00:00, 59.07it/s]\n",
            "Train Epoch: 5 [59500/60000 ( 99%)]  Loss: 0.265050: 100%|██████████| 1200/1200 [00:20<00:00, 59.50it/s]\n",
            "Train Epoch: 6 [59500/60000 ( 99%)]  Loss: 0.228112: 100%|██████████| 1200/1200 [00:20<00:00, 59.10it/s]\n",
            "Train Epoch: 7 [59500/60000 ( 99%)]  Loss: 0.179873: 100%|██████████| 1200/1200 [00:20<00:00, 59.22it/s]\n",
            "Train Epoch: 8 [59500/60000 ( 99%)]  Loss: 0.540372: 100%|██████████| 1200/1200 [00:20<00:00, 58.53it/s]\n",
            "Train Epoch: 9 [59500/60000 ( 99%)]  Loss: 0.117685: 100%|██████████| 1200/1200 [00:20<00:00, 57.83it/s]\n",
            "Train Epoch: 10 [59500/60000 ( 99%)]  Loss: 0.281032: 100%|██████████| 1200/1200 [00:20<00:00, 57.65it/s]\n",
            "Train Epoch: 11 [59500/60000 ( 99%)]  Loss: 0.132364: 100%|██████████| 1200/1200 [00:20<00:00, 57.62it/s]\n",
            "Train Epoch: 12 [59500/60000 ( 99%)]  Loss: 0.046662: 100%|██████████| 1200/1200 [00:20<00:00, 57.37it/s]\n",
            "Train Epoch: 13 [59500/60000 ( 99%)]  Loss: 0.220705: 100%|██████████| 1200/1200 [00:20<00:00, 57.33it/s]\n",
            "Train Epoch: 14 [59500/60000 ( 99%)]  Loss: 0.112033: 100%|██████████| 1200/1200 [00:20<00:00, 57.19it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Test set: Average loss: 0.2075, Accuracy: 9513/10000 (95.13%)\n",
            "--- Before pruning ---\n",
            "fc1.weight           | nonzeros =  235200 /  235200 (100.00%) | total_pruned =       0 | shape = (300, 784)\n",
            "fc1.bias             | nonzeros =     300 /     300 (100.00%) | total_pruned =       0 | shape = (300,)\n",
            "fc2.weight           | nonzeros =   30000 /   30000 (100.00%) | total_pruned =       0 | shape = (100, 300)\n",
            "fc2.bias             | nonzeros =     100 /     100 (100.00%) | total_pruned =       0 | shape = (100,)\n",
            "fc3.weight           | nonzeros =    1000 /    1000 (100.00%) | total_pruned =       0 | shape = (10, 100)\n",
            "fc3.bias             | nonzeros =      10 /      10 (100.00%) | total_pruned =       0 | shape = (10,)\n",
            "alive: 266610, pruned : 0, total: 266610, Compression rate :       1.00x  (  0.00% pruned)\n",
            "Pruning with threshold : 0.24249455332756042 for layer fc1\n",
            "Pruning with threshold : 0.19876055419445038 for layer fc2\n",
            "Pruning with threshold : 0.2157256156206131 for layer fc3\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 0 [    0/60000 (  0%)]  Loss: 1.161682:   0%|          | 6/1200 [00:00<00:22, 52.53it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Test set: Average loss: 1.3665, Accuracy: 6123/10000 (61.23%)\n",
            "--- After pruning ---\n",
            "fc1.weight           | nonzeros =   10964 /  235200 (  4.66%) | total_pruned =  224236 | shape = (300, 784)\n",
            "fc1.bias             | nonzeros =     300 /     300 (100.00%) | total_pruned =       0 | shape = (300,)\n",
            "fc2.weight           | nonzeros =    1419 /   30000 (  4.73%) | total_pruned =   28581 | shape = (100, 300)\n",
            "fc2.bias             | nonzeros =     100 /     100 (100.00%) | total_pruned =       0 | shape = (100,)\n",
            "fc3.weight           | nonzeros =      80 /    1000 (  8.00%) | total_pruned =     920 | shape = (10, 100)\n",
            "fc3.bias             | nonzeros =      10 /      10 (100.00%) | total_pruned =       0 | shape = (10,)\n",
            "alive: 12873, pruned : 253737, total: 266610, Compression rate :      20.71x  ( 95.17% pruned)\n",
            "--- Retraining ---\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Train Epoch: 0 [59500/60000 ( 99%)]  Loss: 0.162916: 100%|██████████| 1200/1200 [00:20<00:00, 57.25it/s]\n",
            "Train Epoch: 1 [59500/60000 ( 99%)]  Loss: 0.138832: 100%|██████████| 1200/1200 [00:22<00:00, 53.81it/s]\n",
            "Train Epoch: 2 [59500/60000 ( 99%)]  Loss: 0.243298: 100%|██████████| 1200/1200 [00:23<00:00, 52.15it/s]\n",
            "Train Epoch: 3 [59500/60000 ( 99%)]  Loss: 0.077512: 100%|██████████| 1200/1200 [00:22<00:00, 52.51it/s]\n",
            "Train Epoch: 4 [59500/60000 ( 99%)]  Loss: 0.162076: 100%|██████████| 1200/1200 [00:23<00:00, 52.14it/s]\n",
            "Train Epoch: 5 [59500/60000 ( 99%)]  Loss: 0.130403: 100%|██████████| 1200/1200 [00:22<00:00, 54.15it/s]\n",
            "Train Epoch: 6 [59500/60000 ( 99%)]  Loss: 0.103247: 100%|██████████| 1200/1200 [00:22<00:00, 53.34it/s]\n",
            "Train Epoch: 7 [59500/60000 ( 99%)]  Loss: 0.231975: 100%|██████████| 1200/1200 [00:22<00:00, 53.71it/s]\n",
            "Train Epoch: 8 [59500/60000 ( 99%)]  Loss: 0.065441: 100%|██████████| 1200/1200 [00:22<00:00, 54.06it/s]\n",
            "Train Epoch: 9 [59500/60000 ( 99%)]  Loss: 0.410031: 100%|██████████| 1200/1200 [00:22<00:00, 54.35it/s]\n",
            "Train Epoch: 10 [59500/60000 ( 99%)]  Loss: 0.250777: 100%|██████████| 1200/1200 [00:22<00:00, 54.01it/s]\n",
            "Train Epoch: 11 [59500/60000 ( 99%)]  Loss: 0.144500: 100%|██████████| 1200/1200 [00:22<00:00, 54.27it/s]\n",
            "Train Epoch: 12 [59500/60000 ( 99%)]  Loss: 0.085857: 100%|██████████| 1200/1200 [00:22<00:00, 54.02it/s]\n",
            "Train Epoch: 13 [59500/60000 ( 99%)]  Loss: 0.336460: 100%|██████████| 1200/1200 [00:22<00:00, 53.13it/s]\n",
            "Train Epoch: 14 [59500/60000 ( 99%)]  Loss: 0.075711: 100%|██████████| 1200/1200 [00:22<00:00, 54.32it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Test set: Average loss: 0.1313, Accuracy: 9631/10000 (96.31%)\n",
            "--- After Retraining ---\n",
            "fc1.weight           | nonzeros =   10964 /  235200 (  4.66%) | total_pruned =  224236 | shape = (300, 784)\n",
            "fc1.bias             | nonzeros =     300 /     300 (100.00%) | total_pruned =       0 | shape = (300,)\n",
            "fc2.weight           | nonzeros =    1419 /   30000 (  4.73%) | total_pruned =   28581 | shape = (100, 300)\n",
            "fc2.bias             | nonzeros =     100 /     100 (100.00%) | total_pruned =       0 | shape = (100,)\n",
            "fc3.weight           | nonzeros =      80 /    1000 (  8.00%) | total_pruned =     920 | shape = (10, 100)\n",
            "fc3.bias             | nonzeros =      10 /      10 (100.00%) | total_pruned =       0 | shape = (10,)\n",
            "alive: 12873, pruned : 253737, total: 266610, Compression rate :      20.71x  ( 95.17% pruned)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U42SRJuHVjgE"
      },
      "source": [
        "You can control other values such as\n",
        "\n",
        "random seed,\n",
        "epochs,\n",
        "sensitivity,\n",
        "batch size,\n",
        "learning rate.\n"
      ]
    }
  ]
}