Pruning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Pruning.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "hNEXyNaqATeQ"
      },
      "source": [
        "import numpy as np\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import os\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "from torch.autograd import Variable\n",
        "import torch.nn.functional as F\n",
        "import torch.nn.utils.prune as prune\n",
        "\n",
        "import copy\n",
        "import random\n",
        "\n",
        "import torchvision\n",
        "import torchvision.transforms as transforms\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from sklearn.metrics import confusion_matrix\n",
        "from torch.quantization import QuantStub, DeQuantStub"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BT6IyucYYEjS"
      },
      "source": [
        "# Глобальное прореживание нейронной сети (спарсификация)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bygS826kV7S7"
      },
      "source": [
        "Для того, чтобы начать оптимизировать размер сети, нам нужен инструментарий для удаления связей внутри нашей модели. Ниже представлена возможная реализация оберток на полносвязанный линейный и сверточный 2d слои. Идея прореживания в том, чтобы сгенирировать определенным образом бинарную маску, регулирующую какие веса мы оставляем, а какие будем отключать. Далее мы перемножим маску с весами слоя,оставляя самые важные, исходя из определенного правила при генерации этой маски.\n",
        "\n",
        "Функция `weight_sparse` реализует алгоритм генерации такой маски исходя из абсолютного значения - задавая пороговое значение, вычисляем персентиль и далее зануляем только те веса, которые меньше этого значения."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zBUzp42tEgRT"
      },
      "source": [
        "class MaskedLinear(nn.Linear):\n",
        "    def __init__(self, in_features, out_features, bias=True):\n",
        "        super(MaskedLinear, self).__init__(in_features, out_features, bias)\n",
        "    \n",
        "    def set_mask(self, mask):\n",
        "        self.mask = torch.tensor(mask, requires_grad=False)\n",
        "        self.weight.data = self.weight.data*self.mask.data\n",
        "    \n",
        "    def get_mask(self):\n",
        "        return self.mask\n",
        "    \n",
        "    def forward(self, x):\n",
        "        return F.linear(x, self.weight, self.bias)\n",
        "        \n",
        "        \n",
        "class MaskedConv2d(nn.Conv2d):\n",
        "    def __init__(self, in_channels, out_channels, kernel_size, stride=1,\n",
        "                 padding=0, dilation=1, groups=1, bias=True):\n",
        "        super(MaskedConv2d, self).__init__(in_channels, out_channels, \n",
        "            kernel_size, stride, padding, dilation, groups, bias)\n",
        "    \n",
        "    def prune(self, mask):\n",
        "        self.mask = torch.tensor(mask, requires_grad=False)\n",
        "        self.weight.data = self.weight.data*self.mask.data\n",
        "    \n",
        "    def get_mask(self):\n",
        "        return self.mask\n",
        "    \n",
        "    def forward(self, x):\n",
        "        return F.conv2d(x, self.weight, self.bias, self.stride,\n",
        "                        self.padding, self.dilation, self.groups)\n",
        "\n",
        "def weight_sparse(model, pruning_perc):    \n",
        "    all_weights = []\n",
        "    for p in model.parameters():\n",
        "        if len(p.data.size()) != 1:\n",
        "            all_weights += list(p.cpu().data.abs().numpy().flatten())\n",
        "    threshold = np.percentile(np.array(all_weights), pruning_perc)\n",
        "\n",
        "    masks = []\n",
        "    for p in model.parameters():\n",
        "        if len(p.data.size()) != 1:\n",
        "            pruned_inds = p.data.abs() > threshold\n",
        "            masks.append(pruned_inds.float())\n",
        "    return masks"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sAIpoKk0Yby9"
      },
      "source": [
        "Алгоритмы прунинга достаточно востребованы и поэтому в современных фреймворках уже представлен набор инструментов позволяющий проводить операцию прунинга. Ниже представлена архитектура нейронной сети ResNet18, адаптированная под низкое разрешение и решающая задачу классификации 10 классов. В методе `prune_unstructured` мы пробегаем по всем слоям нашей сети и применяем функцию прунинга, аналогичную представленной выше.\n",
        "\n",
        "Метод `calc_weights` позволяет посчитать количество весов с учетом сгенерированных для прунинга масок."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "u4XFIdvkBn4T"
      },
      "source": [
        "class BasicBlock(nn.Module):\n",
        "    expansion = 1\n",
        "\n",
        "    def __init__(self, in_planes, planes, stride=1):\n",
        "        super(BasicBlock, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(\n",
        "            in_planes, planes, kernel_size=3, stride=stride, padding=1, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(planes)\n",
        "        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3,\n",
        "                               stride=1, padding=1, bias=False)\n",
        "        self.bn2 = nn.BatchNorm2d(planes)\n",
        "\n",
        "        self.shortcut = nn.Sequential()\n",
        "        if stride != 1 or in_planes != self.expansion*planes:\n",
        "            self.shortcut = nn.Sequential(\n",
        "                nn.Conv2d(in_planes, self.expansion*planes,\n",
        "                          kernel_size=1, stride=stride, bias=False),\n",
        "                nn.BatchNorm2d(self.expansion*planes)\n",
        "            )\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.bn1(self.conv1(x)))\n",
        "        out = self.bn2(self.conv2(out))\n",
        "        skip_branch = self.shortcut(x)\n",
        "        out += skip_branch\n",
        "        out = F.relu(out)\n",
        "        return out\n",
        "\n",
        "class ResNet(nn.Module):\n",
        "    def __init__(self, num_classes=10):\n",
        "        super(ResNet, self).__init__()\n",
        "        block = BasicBlock\n",
        "        num_blocks = [2, 2, 2, 2]\n",
        "        self.in_planes = 64\n",
        "\n",
        "        self.conv1 = nn.Conv2d(1, 64, kernel_size=3,\n",
        "                               stride=1, padding=1, bias=False)\n",
        "        self.bn1 = nn.BatchNorm2d(64)\n",
        "        self.layer1 = self._make_layer(block, 64, num_blocks[0], stride=1)\n",
        "        self.layer2 = self._make_layer(block, 128, num_blocks[1], stride=1)\n",
        "        self.layer3 = self._make_layer(block, 256, num_blocks[2], stride=2)\n",
        "        self.layer4 = self._make_layer(block, 512, num_blocks[3], stride=2)\n",
        "        self.linear = nn.Linear(512*block.expansion, num_classes)\n",
        "\n",
        "    def _make_layer(self, block, planes, num_blocks, stride):\n",
        "        strides = [stride] + [1]*(num_blocks-1)\n",
        "        layers = []\n",
        "        for stride in strides:\n",
        "            layers.append(block(self.in_planes, planes, stride))\n",
        "            self.in_planes = planes * block.expansion\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.bn1(self.conv1(x)))\n",
        "        out = self.layer1(out)\n",
        "        out = self.layer2(out)\n",
        "        out = self.layer3(out)\n",
        "        out = self.layer4(out)\n",
        "        out = F.adaptive_avg_pool2d(out, 1)\n",
        "        out = out.view(out.size(0), -1)\n",
        "        out = self.linear(out)\n",
        "        return out\n",
        "    \n",
        "    def prune_unstructured(self, rate):\n",
        "        for name, module in self.named_modules():\n",
        "            if isinstance(module, torch.nn.Conv2d):\n",
        "                prune.l1_unstructured(module, name='weight', amount=rate)\n",
        "    \n",
        "    def prune_structured(self, rate):\n",
        "        for name, module in self.named_modules():\n",
        "            if isinstance(module, torch.nn.Conv2d):\n",
        "                prune.ln_structured(module, name='weight', n=2, amount=rate, dim=1)\n",
        "\n",
        "    def calc_weights(self):\n",
        "        result = 0\n",
        "        for name, module in self.named_modules():\n",
        "            if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear):\n",
        "                if hasattr(module, 'weight_mask'):\n",
        "                    result += int(torch.sum(module.weight_mask.reshape(-1)).item())\n",
        "                else:\n",
        "                    result += len(module.weight.reshape(-1))\n",
        "        return result"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QXlKL4szZtrb"
      },
      "source": [
        "Далее реализуем простую функцию тренировки сети на тренировочных данных и функцию для тестирования на валидационных данных. Реализация этих функций поддерживает обучение и тестирование как на ЦПУ, так и на ГПУ."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eX6wrlVyFt_d"
      },
      "source": [
        "def fit(model, train_loader, epoch_number=5, device='cuda'):\n",
        "    optimizer = torch.optim.Adam(model.parameters())\n",
        "    error = nn.CrossEntropyLoss()\n",
        "    model.train()\n",
        "    \n",
        "    for epoch in range(epoch_number):\n",
        "        correct = 0\n",
        "        \n",
        "        for batch_idx, (X_batch, y_batch) in enumerate(train_loader):\n",
        "            var_X_batch = X_batch.to(device)\n",
        "            var_y_batch = y_batch.to(device)\n",
        "            \n",
        "            optimizer.zero_grad()\n",
        "            output = model(var_X_batch)\n",
        "            loss = error(output, var_y_batch)\n",
        "            loss.backward()\n",
        "            optimizer.step()\n",
        "\n",
        "            predicted = torch.max(output.data, 1)[1] \n",
        "            correct += (predicted == var_y_batch).sum()\n",
        "            if batch_idx % 500 == 0:\n",
        "                print('Epoch : {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\\t Accuracy:{:.3f}%'.format(\n",
        "                    epoch, batch_idx*len(X_batch), len(train_loader.dataset), \n",
        "                    100.*batch_idx / len(train_loader), loss.data, \n",
        "                    float(correct*100) / float(BATCH_SIZE*(batch_idx+1))))\n",
        "                \n",
        "                \n",
        "def evaluate(model, loader, device='cuda'):\n",
        "    correct = 0\n",
        "    model.eval() \n",
        "    for test_imgs, test_labels in loader:\n",
        "        test_imgs = test_imgs.to(device)\n",
        "        test_labels = test_labels.to(device)\n",
        "        \n",
        "        output = model(test_imgs)\n",
        "        predicted = torch.max(output,1)[1]\n",
        "        correct += (predicted == test_labels).sum()\n",
        "    print(\"Test accuracy:{:.3f}% \".format( float(correct) / (len(loader)*BATCH_SIZE)))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zup2vbLwaTtN"
      },
      "source": [
        "Просчитаем размер нашей сети. Изначальная сеть, ResNet18 имеет примерно 11.16 млн параметров."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "SSEZMlH5J76g",
        "outputId": "7dbe97d9-ed52-4bca-b003-be8cc00f1d58"
      },
      "source": [
        "model = ResNet()\n",
        "model.calc_weights()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "11163200"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 5
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kc4oQNNzay3p"
      },
      "source": [
        "Далее подготовим данные для тренировки и валидации. Для наших целей будем использовать тот же набор данных, что и при реализации алгоритма дистиляции, датасет FashionMNIST, содержащий примрено 70 тысяч черно-белых изображений с разрешением 32х32"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W5jbiVJDbzjA"
      },
      "source": [
        "BATCH_SIZE = 32\n",
        "EPOCH = 10\n",
        "DEVICE = 'cuda'\n",
        "SEED = 5\n",
        "\n",
        "torch.manual_seed(SEED)\n",
        "np.random.seed(SEED)\n",
        "random.seed(SEED)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8XIUXPm8bLwT"
      },
      "source": [
        "train_data = torchvision.datasets.FashionMNIST('./', train=True, download=True, transform=transforms.Compose([transforms.ToTensor()]))\n",
        "test_data = torchvision.datasets.FashionMNIST('./', train=False, download=True, transform=transforms.Compose([transforms.ToTensor()]))\n",
        "train_loader = torch.utils.data.DataLoader(train_data, batch_size = BATCH_SIZE, shuffle = True)\n",
        "test_loader = torch.utils.data.DataLoader(test_data, batch_size = BATCH_SIZE, shuffle = True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AtkSMKxocLkP"
      },
      "source": [
        "Первоначально обучим нашу нейронную сеть на данных перед применением алгоритма прунинга"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "rhZ76it6b1t3",
        "outputId": "5afb5fff-6d99-4735-8f76-74cbfd899c9e"
      },
      "source": [
        "model.to(DEVICE)\n",
        "fit(model, train_loader, epoch_number=EPOCH, device=DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch : 0 [0/60000 (0%)]\tLoss: 2.363766\t Accuracy:9.375%\n",
            "Epoch : 0 [16000/60000 (27%)]\tLoss: 0.407048\t Accuracy:75.680%\n",
            "Epoch : 0 [32000/60000 (53%)]\tLoss: 0.433269\t Accuracy:80.045%\n",
            "Epoch : 0 [48000/60000 (80%)]\tLoss: 0.345497\t Accuracy:82.491%\n",
            "Epoch : 1 [0/60000 (0%)]\tLoss: 0.120425\t Accuracy:96.875%\n",
            "Epoch : 1 [16000/60000 (27%)]\tLoss: 0.383100\t Accuracy:89.590%\n",
            "Epoch : 1 [32000/60000 (53%)]\tLoss: 0.072533\t Accuracy:89.623%\n",
            "Epoch : 1 [48000/60000 (80%)]\tLoss: 0.173559\t Accuracy:89.740%\n",
            "Epoch : 2 [0/60000 (0%)]\tLoss: 0.269849\t Accuracy:90.625%\n",
            "Epoch : 2 [16000/60000 (27%)]\tLoss: 0.142185\t Accuracy:91.442%\n",
            "Epoch : 2 [32000/60000 (53%)]\tLoss: 0.214349\t Accuracy:91.296%\n",
            "Epoch : 2 [48000/60000 (80%)]\tLoss: 0.088656\t Accuracy:91.383%\n",
            "Epoch : 3 [0/60000 (0%)]\tLoss: 0.154574\t Accuracy:93.750%\n",
            "Epoch : 3 [16000/60000 (27%)]\tLoss: 0.042977\t Accuracy:92.633%\n",
            "Epoch : 3 [32000/60000 (53%)]\tLoss: 0.285396\t Accuracy:92.576%\n",
            "Epoch : 3 [48000/60000 (80%)]\tLoss: 0.324311\t Accuracy:92.634%\n",
            "Epoch : 4 [0/60000 (0%)]\tLoss: 0.119824\t Accuracy:93.750%\n",
            "Epoch : 4 [16000/60000 (27%)]\tLoss: 0.071797\t Accuracy:93.787%\n",
            "Epoch : 4 [32000/60000 (53%)]\tLoss: 0.251117\t Accuracy:93.697%\n",
            "Epoch : 4 [48000/60000 (80%)]\tLoss: 0.168030\t Accuracy:93.606%\n",
            "Epoch : 5 [0/60000 (0%)]\tLoss: 0.274272\t Accuracy:87.500%\n",
            "Epoch : 5 [16000/60000 (27%)]\tLoss: 0.111505\t Accuracy:94.785%\n",
            "Epoch : 5 [32000/60000 (53%)]\tLoss: 0.057063\t Accuracy:94.621%\n",
            "Epoch : 5 [48000/60000 (80%)]\tLoss: 0.342732\t Accuracy:94.447%\n",
            "Epoch : 6 [0/60000 (0%)]\tLoss: 0.077095\t Accuracy:96.875%\n",
            "Epoch : 6 [16000/60000 (27%)]\tLoss: 0.084301\t Accuracy:95.677%\n",
            "Epoch : 6 [32000/60000 (53%)]\tLoss: 0.412743\t Accuracy:95.370%\n",
            "Epoch : 6 [48000/60000 (80%)]\tLoss: 0.077073\t Accuracy:95.293%\n",
            "Epoch : 7 [0/60000 (0%)]\tLoss: 0.021903\t Accuracy:100.000%\n",
            "Epoch : 7 [16000/60000 (27%)]\tLoss: 0.191635\t Accuracy:96.158%\n",
            "Epoch : 7 [32000/60000 (53%)]\tLoss: 0.084607\t Accuracy:96.182%\n",
            "Epoch : 7 [48000/60000 (80%)]\tLoss: 0.099870\t Accuracy:96.128%\n",
            "Epoch : 8 [0/60000 (0%)]\tLoss: 0.050992\t Accuracy:100.000%\n",
            "Epoch : 8 [16000/60000 (27%)]\tLoss: 0.008592\t Accuracy:97.343%\n",
            "Epoch : 8 [32000/60000 (53%)]\tLoss: 0.020059\t Accuracy:96.966%\n",
            "Epoch : 8 [48000/60000 (80%)]\tLoss: 0.089153\t Accuracy:96.877%\n",
            "Epoch : 9 [0/60000 (0%)]\tLoss: 0.040822\t Accuracy:100.000%\n",
            "Epoch : 9 [16000/60000 (27%)]\tLoss: 0.004009\t Accuracy:97.967%\n",
            "Epoch : 9 [32000/60000 (53%)]\tLoss: 0.012826\t Accuracy:97.855%\n",
            "Epoch : 9 [48000/60000 (80%)]\tLoss: 0.163545\t Accuracy:97.712%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3qnyEwP0dPSn",
        "outputId": "6c94e081-c76a-4b28-8df4-66ecd540b8b7"
      },
      "source": [
        "evaluate(model, test_loader, device=DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Test accuracy:0.932% \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8a4Azkfcyhp"
      },
      "source": [
        "Далее попробуем применить спарсификацию на полученной модели напрямую. Посмотрим на просадку в качестве"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yep7tMP4dEGx"
      },
      "source": [
        "pruned_model = copy.deepcopy(model)\n",
        "pruned_model.prune_unstructured(0.5)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TN70TSQadW_w"
      },
      "source": [
        "Посмотрим сколько теперь параметров имеет наша сеть. Видно, что, действительно мы смогли убрать 50% весов. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8ALbI39Sdcsp",
        "outputId": "c9d373fd-fdcc-44de-b0d1-d15b76f80aaf"
      },
      "source": [
        "pruned_model.calc_weights()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "5584160"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "38sqG7bEdxsQ"
      },
      "source": [
        "Измерим качество модели после такого алгоритма спарсификации"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "GkgCrGa3d03I",
        "outputId": "9b6c901d-08f7-4be2-f36a-623880400f3e"
      },
      "source": [
        "evaluate(pruned_model, test_loader, device=DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Test accuracy:0.909% \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qpbKZ5vZwOdc"
      },
      "source": [
        "Точность модели просела на 2.3%. Результат довльно слабый. Попробуем применить другой подход, сжать сеть еще больше и при этом не допустить сильной потери в качестве"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LzMiAyBSd_Xc"
      },
      "source": [
        "# Итеративное прореживание"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RgFvQfB_eHbb"
      },
      "source": [
        "Идин из способов сжатия нейронных сетей - итеративное прореживание (Incremental Magnitude Pruning). Он достаточно ресурсоемкий, однако позволяет достаточно несложными методами добиться неплохого результата. Здесь используется более хитрый подход.\n",
        "\n",
        "Будем идти с шагом, каждый раз будем отключать внутри сети несколько десятков процентов связей. После отключения, оставшиеся веса дообучим на всех данных используя одну эпоху. Ожидается, что так как мы выкинули за один раз не очень много, то оставшиеся связи \"перехватят\" ответственность тех слабых, которые мы только что отключили.\n",
        "\n",
        "Таким образом за P таких итераций мы выкинем желаемое количество сети и не должны при этом потерять сильно в качестве.\n",
        "\n",
        "Мы составим расписание для сети в виде списка и напишем более умную функцию тренировки"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ow2fWZNZemam"
      },
      "source": [
        "def smart_prune_shed(model, train_loader, schedule, device='cpu'):\n",
        "    optimizer = torch.optim.Adam(model.parameters())\n",
        "    error = nn.CrossEntropyLoss()\n",
        "    model.train()\n",
        "    \n",
        "    for rate, epochs in schedule:  # Идем шагами, согласно тому расписанию, которое передали в функцию\n",
        "        t = rate/100  # Считаем очередное пороговое значение\n",
        "        model.prune_unstructured(t)  # Отключаем слабые связи\n",
        "        for i in range(epochs):\n",
        "            correct = 0\n",
        "            for batch_idx, (X_batch, y_batch) in enumerate(train_loader):  # Далее дообучаем модель как обычно в течение указанного количества эпох\n",
        "                var_X_batch = X_batch.to(device)\n",
        "                var_y_batch = y_batch.to(device)\n",
        "                optimizer.zero_grad()\n",
        "                output = model(var_X_batch)\n",
        "                loss = error(output, var_y_batch)\n",
        "                loss.backward()\n",
        "                optimizer.step()\n",
        "\n",
        "                predicted = torch.max(output.data, 1)[1] \n",
        "                correct += (predicted == var_y_batch).sum()\n",
        "                if batch_idx % 500 == 0:\n",
        "                    print('Rate : {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}\\t Accuracy:{:.3f}%'.format(\n",
        "                        rate, batch_idx*len(X_batch), len(train_loader.dataset), 100.*batch_idx / len(train_loader), loss.data, float(correct*100) / float(BATCH_SIZE*(batch_idx+1))))\n",
        "    return model"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zfQrjh_GfKhE",
        "outputId": "e30bec82-1e87-4ba1-ef93-d9eaedc11655"
      },
      "source": [
        "pruned_model = copy.deepcopy(model)\n",
        "pruned_model_70 = smart_prune_shed(pruned_model, train_loader, [\n",
        "    (50, 1), \n",
        "    (20, 1), \n",
        "    (10, 1), \n",
        "    (10, 1), \n",
        "    (10, 1), \n",
        "], DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Rate : 50 [0/60000 (0%)]\tLoss: 0.048585\t Accuracy:96.875%\n",
            "Rate : 50 [16000/60000 (27%)]\tLoss: 0.046688\t Accuracy:98.765%\n",
            "Rate : 50 [32000/60000 (53%)]\tLoss: 0.084925\t Accuracy:98.764%\n",
            "Rate : 50 [48000/60000 (80%)]\tLoss: 0.089002\t Accuracy:98.767%\n",
            "Rate : 20 [0/60000 (0%)]\tLoss: 0.164155\t Accuracy:96.875%\n",
            "Rate : 20 [16000/60000 (27%)]\tLoss: 0.000696\t Accuracy:99.077%\n",
            "Rate : 20 [32000/60000 (53%)]\tLoss: 0.001487\t Accuracy:99.151%\n",
            "Rate : 20 [48000/60000 (80%)]\tLoss: 0.011173\t Accuracy:99.203%\n",
            "Rate : 10 [0/60000 (0%)]\tLoss: 0.003566\t Accuracy:100.000%\n",
            "Rate : 10 [16000/60000 (27%)]\tLoss: 0.003731\t Accuracy:99.339%\n",
            "Rate : 10 [32000/60000 (53%)]\tLoss: 0.004622\t Accuracy:99.422%\n",
            "Rate : 10 [48000/60000 (80%)]\tLoss: 0.002862\t Accuracy:99.415%\n",
            "Rate : 10 [0/60000 (0%)]\tLoss: 0.101215\t Accuracy:93.750%\n",
            "Rate : 10 [16000/60000 (27%)]\tLoss: 0.010260\t Accuracy:99.495%\n",
            "Rate : 10 [32000/60000 (53%)]\tLoss: 0.041891\t Accuracy:99.529%\n",
            "Rate : 10 [48000/60000 (80%)]\tLoss: 0.007088\t Accuracy:99.563%\n",
            "Rate : 10 [0/60000 (0%)]\tLoss: 0.000233\t Accuracy:100.000%\n",
            "Rate : 10 [16000/60000 (27%)]\tLoss: 0.000239\t Accuracy:99.563%\n",
            "Rate : 10 [32000/60000 (53%)]\tLoss: 0.007374\t Accuracy:99.582%\n",
            "Rate : 10 [48000/60000 (80%)]\tLoss: 0.001345\t Accuracy:99.590%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BbcdKZ10kWu_"
      },
      "source": [
        "Подобным расписанием мы смогли сжать сеть примерно на 70%, посмотрим на полученное качество"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QA8nwQSGkcBF",
        "outputId": "671eb755-850e-4e68-d0ab-403064e6bc4b"
      },
      "source": [
        "evaluate(pruned_model_70, test_loader, device=DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Test accuracy:0.932% \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EtiCWSXOk5Qr"
      },
      "source": [
        "Попробуем сжать еще сильнее. Выкинем 90% сети."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "4PPrVHEwk9tO",
        "outputId": "cfdcbf3f-277a-4992-c321-0363d75eaded"
      },
      "source": [
        "pruned_model_90 = smart_prune_shed(pruned_model_70, train_loader, [\n",
        "    (20, 1), \n",
        "    (20, 1), \n",
        "    (20, 1), \n",
        "    (10, 1), \n",
        "    (10, 1), \n",
        "    (10, 1), \n",
        "    (5, 1), \n",
        "], DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Rate : 20 [0/60000 (0%)]\tLoss: 0.094138\t Accuracy:96.875%\n",
            "Rate : 20 [16000/60000 (27%)]\tLoss: 0.002804\t Accuracy:98.971%\n",
            "Rate : 20 [32000/60000 (53%)]\tLoss: 0.043634\t Accuracy:99.285%\n",
            "Rate : 20 [48000/60000 (80%)]\tLoss: 0.001374\t Accuracy:99.332%\n",
            "Rate : 20 [0/60000 (0%)]\tLoss: 0.005377\t Accuracy:100.000%\n",
            "Rate : 20 [16000/60000 (27%)]\tLoss: 0.000176\t Accuracy:98.933%\n",
            "Rate : 20 [32000/60000 (53%)]\tLoss: 0.000257\t Accuracy:99.298%\n",
            "Rate : 20 [48000/60000 (80%)]\tLoss: 0.006784\t Accuracy:99.375%\n",
            "Rate : 20 [0/60000 (0%)]\tLoss: 0.309464\t Accuracy:90.625%\n",
            "Rate : 20 [16000/60000 (27%)]\tLoss: 0.000751\t Accuracy:98.834%\n",
            "Rate : 20 [32000/60000 (53%)]\tLoss: 0.000462\t Accuracy:99.151%\n",
            "Rate : 20 [48000/60000 (80%)]\tLoss: 0.022798\t Accuracy:99.296%\n",
            "Rate : 10 [0/60000 (0%)]\tLoss: 0.141815\t Accuracy:93.750%\n",
            "Rate : 10 [16000/60000 (27%)]\tLoss: 0.001237\t Accuracy:99.489%\n",
            "Rate : 10 [32000/60000 (53%)]\tLoss: 0.000137\t Accuracy:99.604%\n",
            "Rate : 10 [48000/60000 (80%)]\tLoss: 0.000535\t Accuracy:99.652%\n",
            "Rate : 10 [0/60000 (0%)]\tLoss: 0.006457\t Accuracy:100.000%\n",
            "Rate : 10 [16000/60000 (27%)]\tLoss: 0.039290\t Accuracy:99.420%\n",
            "Rate : 10 [32000/60000 (53%)]\tLoss: 0.009612\t Accuracy:99.563%\n",
            "Rate : 10 [48000/60000 (80%)]\tLoss: 0.003061\t Accuracy:99.625%\n",
            "Rate : 10 [0/60000 (0%)]\tLoss: 0.019359\t Accuracy:100.000%\n",
            "Rate : 10 [16000/60000 (27%)]\tLoss: 0.018327\t Accuracy:99.295%\n",
            "Rate : 10 [32000/60000 (53%)]\tLoss: 0.002454\t Accuracy:99.522%\n",
            "Rate : 10 [48000/60000 (80%)]\tLoss: 0.000019\t Accuracy:99.604%\n",
            "Rate : 5 [0/60000 (0%)]\tLoss: 0.001032\t Accuracy:100.000%\n",
            "Rate : 5 [16000/60000 (27%)]\tLoss: 0.013508\t Accuracy:99.750%\n",
            "Rate : 5 [32000/60000 (53%)]\tLoss: 0.000324\t Accuracy:99.769%\n",
            "Rate : 5 [48000/60000 (80%)]\tLoss: 0.002301\t Accuracy:99.786%\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aoFL-jNUmexY",
        "outputId": "396d16aa-4376-435c-90e6-dd543f6c7803"
      },
      "source": [
        "evaluate(pruned_model, test_loader, device=DEVICE)"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Test accuracy:0.933% \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3N9rXHt6vwcr",
        "outputId": "30a0618f-8010-4a69-b15c-9ceccb50c910"
      },
      "source": [
        "pruned_model.calc_weights()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1158832"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 18
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BWbqwuwm1BcQ",
        "outputId": "543e5c5d-c3e0-4084-9e00-f22e854c8fa5"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "!cp \"drive/My Drive/Colab Notebooks/Pruning.ipynb\" ./\n",
        "\n",
        "!jupyter nbconvert --to latex Pruning.ipynb\n",
        "!cp Pruning.tex \"drive/My Drive/Colab Notebooks/\""
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive\n",
            "[NbConvertApp] Converting notebook Pruning.ipynb to latex\n",
            "[NbConvertApp] Writing 66068 bytes to Pruning.tex\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}