sem16/main.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "celltoolbar": "Slideshow",
    "colab": {
      "name": "main.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.7.3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OIQJ9bM4MSPL"
      },
      "source": [
        "#  Нейронные сети: Автокодировщик, Генеративно-Состязательные сети"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s3qWY0M5LA6r"
      },
      "source": [
        "## Библиотеки"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K2_VhyWeteMB"
      },
      "source": [
        "from copy import deepcopy\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib.image import imread\n",
        "from mpl_toolkits import mplot3d\n",
        "from matplotlib import gridspec\n",
        "from PIL import Image\n",
        "import io\n",
        "import os\n",
        "from urllib.request import urlopen\n",
        "from skimage.segmentation import mark_boundaries\n",
        "\n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "import requests\n",
        "from scipy.stats import norm\n",
        "import torch\n",
        "\n",
        "from sklearn.metrics import classification_report\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "\n",
        "from torchvision import datasets, transforms\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UCn8xDPhteMB"
      },
      "source": [
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WF45BXo5MWuU"
      },
      "source": [
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "device"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GolP2cBBXfeN"
      },
      "source": [
        "# Tensorboard"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S_Ns3K7pXiP1"
      },
      "source": [
        "%load_ext tensorboard\n",
        "%tensorboard --logdir ./"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YJCc8xi09uWI"
      },
      "source": [
        "## Код для обучени модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R3BclNuZMSPT"
      },
      "source": [
        "def train_epoch(train_generator, model, loss_function, optimizer, callback = None):\n",
        "    epoch_loss = 0\n",
        "    total = 0\n",
        "    for it, (batch_of_x, batch_of_y) in enumerate(train_generator):\n",
        "        batch_loss = train_on_batch(model, batch_of_x, batch_of_y, optimizer, loss_function)\n",
        "        \n",
        "        if callback is not None:\n",
        "            with torch.no_grad():\n",
        "                callback(model, batch_loss)\n",
        "            \n",
        "        epoch_loss += batch_loss*len(batch_of_x)\n",
        "        total += len(batch_of_x)\n",
        "    \n",
        "    return epoch_loss/total\n",
        "        "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8VG64px1MSPT"
      },
      "source": [
        "def trainer(count_of_epoch, \n",
        "            batch_size, \n",
        "            dataset,\n",
        "            model, \n",
        "            loss_function,\n",
        "            optimizer,\n",
        "            lr = 0.001,\n",
        "            callback = None):\n",
        "\n",
        "    optima = optimizer(model.parameters(), lr=lr)\n",
        "    \n",
        "    iterations = tqdm(range(count_of_epoch), desc='epoch')\n",
        "    iterations.set_postfix({'train epoch loss': np.nan})\n",
        "    for it in iterations:\n",
        "        batch_generator = tqdm(\n",
        "            torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True), \n",
        "            leave=False, total=len(dataset)//batch_size+(len(dataset)%batch_size> 0))\n",
        "        \n",
        "        epoch_loss = train_epoch(train_generator=batch_generator, \n",
        "                    model=model, \n",
        "                    loss_function=loss_function, \n",
        "                    optimizer=optima, \n",
        "                    callback=callback)\n",
        "        \n",
        "        iterations.set_postfix({'train epoch loss': epoch_loss})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "C83NQZtcMSPP"
      },
      "source": [
        "## Автокодировщик"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A7NcDFLef7XZ"
      },
      "source": [
        "digit_size = (28, 28)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7BMdudN9MSPT"
      },
      "source": [
        "### Код для обучения модели автокодировщика"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XWEv5oI9MSPT"
      },
      "source": [
        "def train_on_batch(model, x_batch, y_batch, optimizer, loss_function):\n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    output = model(x_batch.to(model.device))\n",
        "    \n",
        "    loss = loss_function(output, x_batch.to(model.device))\n",
        "    loss.backward()\n",
        "\n",
        "    optimizer.step()\n",
        "    return loss.cpu().item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YJTqzA44NX4A"
      },
      "source": [
        "class callback():\n",
        "    def __init__(self, writer, dataset, loss_function, delimeter = 100, batch_size=64):\n",
        "        self.step = 0\n",
        "        self.writer = writer\n",
        "        self.delimeter = delimeter\n",
        "        self.loss_function = loss_function\n",
        "        self.batch_size = batch_size\n",
        "\n",
        "        self.dataset = dataset\n",
        "\n",
        "    def forward(self, model, loss):\n",
        "        self.step += 1\n",
        "        self.writer.add_scalar('LOSS/train', loss, self.step)\n",
        "        \n",
        "        if self.step % self.delimeter == 0:\n",
        "            \n",
        "            batch_generator = torch.utils.data.DataLoader(dataset = self.dataset, \n",
        "                                                          batch_size=self.batch_size)\n",
        "            \n",
        "            pred = []\n",
        "            real = []\n",
        "            test_loss = 0\n",
        "            model.eval()\n",
        "            for it, (x_batch, _) in enumerate(batch_generator):\n",
        "                x_batch = x_batch.to(model.device)\n",
        "\n",
        "                output = model(x_batch)\n",
        "\n",
        "                test_loss += self.loss_function(output, x_batch).cpu().item()*len(x_batch)\n",
        "\n",
        "                pred.extend(torch.argmax(output, dim=-1).cpu().numpy().tolist())\n",
        "            \n",
        "            test_loss /= len(self.dataset)\n",
        "            \n",
        "            self.writer.add_scalar('LOSS/test', test_loss, self.step)\n",
        "\n",
        "            x = x_batch[-10:]\n",
        "            \n",
        "            fig = plt.figure(figsize=(12, 12 / 10 * (x.shape[0] // 10 + 1)))\n",
        "            gs = gridspec.GridSpec(1, x.shape[0])\n",
        "            for i in range(x.shape[0]):\n",
        "                ax = fig.add_subplot(gs[i])\n",
        "                ax.imshow(model(x.to(model.device)).cpu().data[i].view(*digit_size).numpy(), cmap='Greys_r', interpolation='lanczos')\n",
        "                ax.axis('off')\n",
        "                    \n",
        "            self.writer.add_figure('VISUAL/decoded', fig, self.step)\n",
        "          \n",
        "    def __call__(self, model, loss):\n",
        "        return self.forward(model, loss)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_sqA2WHaMSPQ"
      },
      "source": [
        "### Линейный Автокодировщик"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cNmrueoeVrZw"
      },
      "source": [
        "#### Данные"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E-FbiZB-MSPQ"
      },
      "source": [
        "preprocess = transforms.Compose([transforms.ToTensor(), \n",
        "                                 torch.FloatTensor, \n",
        "                                 lambda x: x.view(784)])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JOCjz2TUMSPQ"
      },
      "source": [
        "MNIST_train = datasets.MNIST('./mnist', train=True, download=True, \n",
        "                             transform=preprocess)\n",
        "\n",
        "MNIST_test = datasets.MNIST('./mnist', train=False, download=True,\n",
        "                            transform=preprocess)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MuJgVfmwMSPQ"
      },
      "source": [
        "for x_batch, y_batch in torch.utils.data.DataLoader(dataset=MNIST_train, batch_size=64, shuffle=True):\n",
        "    pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SoAsJHYVMSPR"
      },
      "source": [
        "x = x_batch[-10:]\n",
        "plt.figure(figsize=(12, 12 / 10 * (x.shape[0] // 10 + 1)))\n",
        "for i in range(x.shape[0]):\n",
        "    plt.subplot(x.shape[0] // 10 + 1, 10, i + 1)\n",
        "    plt.imshow(x.data[i].view(*digit_size).numpy(), cmap='Greys_r', interpolation='lanczos')\n",
        "    plt.axis('off')\n",
        "        \n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-_FGHGKYe0DM"
      },
      "source": [
        "#### Модель"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SXunTUISMSPS"
      },
      "source": [
        "class Encoder(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, object_dim, d):\n",
        "        super(type(self), self).__init__()\n",
        "        self.object_dim = object_dim\n",
        "        self.d = d\n",
        "        \n",
        "        self.linear = torch.nn.Linear(object_dim, d)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.linear(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I-P40h8TMSPS"
      },
      "source": [
        "class Decoder(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, d, object_dim):\n",
        "        super(type(self), self).__init__()\n",
        "        self.object_dim = object_dim\n",
        "        self.d = d\n",
        "        \n",
        "        self.linear = torch.nn.Linear(d, object_dim)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.linear(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kpwQG80LMSPS"
      },
      "source": [
        "class Autoencoder(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, object_dim, d):\n",
        "        super(type(self), self).__init__()\n",
        "        self.object_dim = object_dim\n",
        "        self.d = d\n",
        "        \n",
        "        self.encoder = Encoder(np.prod(digit_size), d)\n",
        "        self.decoder = Decoder(d, np.prod(digit_size))\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.decode(self.encode(x))\n",
        "    \n",
        "    def encode(self, x):\n",
        "        return self.encoder(x)\n",
        "    \n",
        "    def decode(self, z):\n",
        "        return self.decoder(z)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rkaIDHwAMSPU"
      },
      "source": [
        "#### Определение модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nIs7TFlCMSPU"
      },
      "source": [
        "loss_function = torch.nn.MSELoss()\n",
        "optimizer = torch.optim.Adam"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EMVySnTbRcvI"
      },
      "source": [
        "#### Перебор разного d: обучение модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "E2SnLdSzRa9l"
      },
      "source": [
        "for d in [2, 4, 8, 16, 32, 64]:\n",
        "    autoencoder = Autoencoder(np.prod(digit_size), d)\n",
        "    autoencoder.to(device)\n",
        "\n",
        "    writer = SummaryWriter(log_dir = 'autoencoder-linear/{}'.format(d))\n",
        "\n",
        "    call = callback(writer, MNIST_test, loss_function, delimeter = 100)\n",
        "\n",
        "    trainer(count_of_epoch=5, \n",
        "            batch_size=64, \n",
        "            dataset=MNIST_train,\n",
        "            model=autoencoder, \n",
        "            loss_function=loss_function,\n",
        "            optimizer = optimizer,\n",
        "            lr = 0.001,\n",
        "            callback = call)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k2yKXsqNMSPW"
      },
      "source": [
        "### Нейросетевой автокодировщик "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t05EFVAafbvi"
      },
      "source": [
        "#### Данные"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xEBAXqo5Yzxn"
      },
      "source": [
        "preprocess = transforms.Compose([transforms.ToTensor(), \n",
        "                                 torch.FloatTensor])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s1xXXsVKYzxo"
      },
      "source": [
        "MNIST_train = datasets.MNIST('./mnist', train=True, download=True, \n",
        "                             transform=preprocess)\n",
        "\n",
        "MNIST_test = datasets.MNIST('./mnist', train=False, download=True,\n",
        "                            transform=preprocess)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kbW2bUIwfezO"
      },
      "source": [
        "#### Модель"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9-NoOA7MVxDg"
      },
      "source": [
        "class EncoderCNN(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, d):\n",
        "        super(type(self), self).__init__()\n",
        "        self.d = d\n",
        "\n",
        "        self.layers = torch.nn.Sequential()\n",
        "        self.layers.add_module('conv1', torch.nn.Conv2d(1, 6, kernel_size = 5))\n",
        "        self.layers.add_module('bnorm1', torch.nn.BatchNorm2d(6))\n",
        "        self.layers.add_module('relu1', torch.nn.ReLU())\n",
        "        self.layers.add_module('pool1', torch.nn.MaxPool2d(kernel_size = 2))\n",
        "        \n",
        "        self.layers.add_module('conv2', torch.nn.Conv2d(6, 16, kernel_size = 5))\n",
        "        self.layers.add_module('bnorm2', torch.nn.BatchNorm2d(16))\n",
        "        self.layers.add_module('relu2', torch.nn.ReLU())\n",
        "        self.layers.add_module('pool2', torch.nn.MaxPool2d(kernel_size = 2))\n",
        "        \n",
        "        self.layers.add_module('conv3', torch.nn.Conv2d(16, self.d, kernel_size = 3))\n",
        "        self.layers.add_module('relu3', torch.nn.ReLU())\n",
        "        self.layers.add_module('pool3', torch.nn.MaxPool2d(kernel_size = 2))\n",
        "\n",
        "        self.layers.add_module('flatten', torch.nn.Flatten())\n",
        "\n",
        "    def forward(self, input):\n",
        "        return self.layers(input)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YxkWhXiGadOh"
      },
      "source": [
        "class Reshape(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, *args):\n",
        "        super(type(self), self).__init__()\n",
        "        self.dims = args\n",
        "\n",
        "    def forward(self, input):\n",
        "        return input.view(input.size(0), *self.dims)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YgGfbx9SXyxw"
      },
      "source": [
        "class DecoderCNN(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, d):\n",
        "        super(type(self), self).__init__()\n",
        "        self.d = d\n",
        "\n",
        "        self.layers = torch.nn.Sequential()\n",
        "        self.layers.add_module('reshape1', Reshape(self.d, 1, 1))\n",
        "        \n",
        "        self.layers.add_module('conv1', torch.nn.ConvTranspose2d(self.d, 128, 4, 1, 0, 0, bias=False))\n",
        "        self.layers.add_module('bnorm1', torch.nn.BatchNorm2d(128))\n",
        "        self.layers.add_module('relu1', torch.nn.ReLU())\n",
        "        \n",
        "        self.layers.add_module('conv2', torch.nn.ConvTranspose2d(128, 64, 3, 2, 1, 0, bias=False))\n",
        "        self.layers.add_module('bnorm2', torch.nn.BatchNorm2d(64))\n",
        "        self.layers.add_module('relu2', torch.nn.ReLU())\n",
        "        \n",
        "        self.layers.add_module('conv3', torch.nn.ConvTranspose2d(64, 32, 3, 2, 1, 1, bias=False))\n",
        "        self.layers.add_module('bnorm3', torch.nn.BatchNorm2d(32))\n",
        "        self.layers.add_module('relu3', torch.nn.ReLU())\n",
        "\n",
        "        self.layers.add_module('conv4', torch.nn.ConvTranspose2d(32, 1, 4, 2, 1, 0, bias=False))\n",
        "        self.layers.add_module('relu4', torch.nn.ReLU())\n",
        "        \n",
        "        self.layers.add_module('sigmoid', torch.nn.Sigmoid())\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.layers(x)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mMvQhigRbwCa"
      },
      "source": [
        "class AutoencoderCNN(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, d):\n",
        "        super(type(self), self).__init__()\n",
        "        self.d = d\n",
        "        \n",
        "        self.encoder = EncoderCNN(d)\n",
        "        self.decoder = DecoderCNN(d)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        return self.decode(self.encode(x))\n",
        "    \n",
        "    def encode(self, x):\n",
        "        return self.encoder(x)\n",
        "    \n",
        "    def decode(self, z):\n",
        "        return self.decoder(z)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L6GydQ8xcI5S"
      },
      "source": [
        "#### Определение модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VAXkjComcI5V"
      },
      "source": [
        "loss_function = torch.nn.MSELoss()\n",
        "optimizer = torch.optim.Adam"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gx6KwOxWcI5W"
      },
      "source": [
        "#### Перебор разного d: обучение модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1nxhq5Tcbn7p"
      },
      "source": [
        "for d in [2, 4, 8, 16, 32, 64]:\n",
        "    autoencoder = AutoencoderCNN(d)\n",
        "    autoencoder.to(device)\n",
        "\n",
        "    writer = SummaryWriter(log_dir = 'autoencoder-cnn/{}'.format(d))\n",
        "    call = callback(writer, MNIST_test, loss_function, delimeter = 100)\n",
        "\n",
        "    trainer(count_of_epoch=5, \n",
        "            batch_size=64, \n",
        "            dataset=MNIST_train,\n",
        "            model=autoencoder, \n",
        "            loss_function=loss_function,\n",
        "            optimizer = optimizer,\n",
        "            lr = 0.001,\n",
        "            callback = call)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wM3gKcgGMSPW"
      },
      "source": [
        "### Вариационный автокодировщик"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kPnQm_lEluQU"
      },
      "source": [
        "#### Модель"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0DSQY6qfke2U"
      },
      "source": [
        "class VAE(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, latent_dim, input_dim, hidden_dim=200):\n",
        "        \"\"\"\n",
        "        Standart model of VAE with ELBO optimization.\n",
        "        Args:\n",
        "            latent_dim: int - the dimension of latent space.\n",
        "            input_dim: int - the dimension of input space.\n",
        "            hidden_dim: int - the size of hidden_dim neural layer.\n",
        "        Returns:\n",
        "            None\n",
        "        Example:\n",
        "            >>> model = VAE(2, 10)\n",
        "        \"\"\"\n",
        "        super(VAE, self).__init__()\n",
        "        self.latent_dim = latent_dim\n",
        "        self.input_dim = input_dim\n",
        "        self.hidden_dim = hidden_dim\n",
        "\n",
        "        self.proposal_z = torch.nn.Sequential(\n",
        "            torch.nn.Linear(self.input_dim, hidden_dim),\n",
        "            torch.nn.LeakyReLU(),\n",
        "        )\n",
        "        self.proposal_mu = torch.nn.Linear(hidden_dim, self.latent_dim)\n",
        "        self.proposal_sigma = torch.nn.Linear(hidden_dim, self.latent_dim)\n",
        "\n",
        "        self.generative_network = torch.nn.Sequential(\n",
        "            torch.nn.Linear(self.latent_dim, hidden_dim),\n",
        "            torch.nn.LeakyReLU(),\n",
        "            torch.nn.Linear(hidden_dim, self.input_dim),\n",
        "            torch.nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "    def q_z(self, x):\n",
        "        \"\"\"\n",
        "        Generates distribution of z provided x.\n",
        "        Args:\n",
        "            x: Tensor - the matrix of shape batch_size x input_dim.\n",
        "        Returns:\n",
        "            tuple(Tensor, Tensor) - the normal distribution parameters.\n",
        "            mu: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "            sigma: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        x = x.to(self.device)\n",
        "\n",
        "        proposal = self.proposal_z(x)\n",
        "        mu = self.proposal_mu(proposal)\n",
        "        sigma = torch.nn.Softplus()(self.proposal_sigma(proposal))\n",
        "        return mu, sigma\n",
        "\n",
        "    def p_z(self, num_samples):\n",
        "        \"\"\"\n",
        "        Generetes prior distribution of z.\n",
        "        Args:\n",
        "            num_samples: int - the number of samples.\n",
        "        Returns:\n",
        "            tuple(Tensor, Tensor) - the normal distribution parameters.\n",
        "                mu: Tensor - the matrix of shape num_samples x latent_dim.\n",
        "            \tsigma: Tensor - the matrix of shape num_samples x latent_dim.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        mu = torch.zeros([num_samples, self.latent_dim], device=self.device)\n",
        "        sigma = torch.ones([num_samples, self.latent_dim], device=self.device)\n",
        "        return mu, sigma\n",
        "\n",
        "    def sample_z(self, distr, num_samples=1):\n",
        "        \"\"\"\n",
        "        Generates samples from normal distribution q(z|x).\n",
        "        Args:\n",
        "            distr = (mu, sigma): tuple(Tensor, Tensor) - the normal distribution parameters.\n",
        "                mu: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "                sigma: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "            num_samples: int - the number of samples for each element.\n",
        "        Returns:\n",
        "            Tensor - the tensor of shape batch_size x num_samples x latent_dim - samples from normal distribution in latent space.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        mu, sigma = distr\n",
        "        mu = mu.to(self.device)\n",
        "        sigma = sigma.to(self.device)\n",
        "\n",
        "        batch_size = mu.shape[0]\n",
        "\n",
        "        bias = mu.view([batch_size, 1, self.latent_dim])\n",
        "\n",
        "        epsilon = torch.randn([batch_size, num_samples, self.latent_dim],\n",
        "                              requires_grad=True,\n",
        "                              device=self.device)\n",
        "        scale = sigma.view([batch_size, 1, self.latent_dim])\n",
        "\n",
        "        return bias + epsilon * scale\n",
        "\n",
        "    def q_x(self, z):\n",
        "        \"\"\"\n",
        "        Given the latent representation matrix z, returns the matrix of Bernoulli distribution parameters for sampling x objects.\n",
        "        Args:\n",
        "            z: Tensor - the tensor of shape batch_size x num_samples x latent_dim, samples from latent space.\n",
        "        Returns:\n",
        "            Tensor - the tensor of shape batch_size x num_samples x input_dim, Bernoulli distribution parameters.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        z = z.to(self.device)\n",
        "        out = self.generative_network(z)\n",
        "\n",
        "        return torch.clamp(out, 0.01, 0.99)\n",
        "\n",
        "    def loss(self, batch_x, batch_y):\n",
        "        \"\"\"\n",
        "        Calculate ELBO approximation of log likelihood for given batch with negative sign.\n",
        "        Args:\n",
        "            batch_x: FloatTensor - the matrix of shape batch_size x input_dim.\n",
        "            batch_y: FloatTensor - dont uses parameter in this model.\n",
        "        Returns:\n",
        "            Tensor - scalar, ELBO approximation of log likelihood for given batch with negative sign.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        batch_x = batch_x.to(self.device)\n",
        "        batch_y = batch_y.to(self.device)\n",
        "\n",
        "        batch_size = batch_x.shape[0]\n",
        "\n",
        "        propos_distr = self.q_z(batch_x)\n",
        "        pri_distr = self.p_z(batch_size)\n",
        "\n",
        "        x_distr = self.q_x(self.sample_z(propos_distr))\n",
        "\n",
        "        expectation = torch.mean(\n",
        "            self.log_mean_exp(\n",
        "                self.log_likelihood(\n",
        "                    batch_x, x_distr)), dim=0)\n",
        "\n",
        "        divergence = self.divergence_KL_normal(propos_distr, pri_distr)\n",
        "\n",
        "        return -1 * torch.mean(expectation - divergence, dim=0)\n",
        "\n",
        "    def generate_samples(self, num_samples):\n",
        "        \"\"\"\n",
        "        Generate samples of object x from noises in latent space.\n",
        "        Args:\n",
        "            num_samples: int - the number of samples, witch need to generate.\n",
        "        Returns:\n",
        "            Tensor - the matrix of shape num_samples x input_dim.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        distr_z = self.p_z(num_samples=1)\n",
        "\n",
        "        z = self.sample_z(distr, num_samples=num_samples)\n",
        "\n",
        "        distr_x = self.q_x(z).view([num_samples, -1])\n",
        "\n",
        "        return torch.bernoulli(distr_x, device=self.device)\n",
        "\n",
        "    @staticmethod\n",
        "    def log_pdf_normal(distr, samples):\n",
        "        \"\"\"\n",
        "        The function calculates the logarithm of the probability density at a point relative to the corresponding normal distribution given componentwise by its mean and standard deviation.\n",
        "        Args:\n",
        "            distr = (mu, sigma): tuple(Tensor, Tensor) - the normal distribution parameters.\n",
        "                mu: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "                sigma: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "            samples: Tensor - the tensor of shape batch_size x num_samples x latent_dim, samples in latent space.\n",
        "        Returns:\n",
        "            Tensor - the matrix of shape batch_size x num_samples, each element of which is the logarithm of the probability density of a point relative to the corresponding distribution.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        mu, sigma = distr\n",
        "\n",
        "        batch_size = mu.shape[0]\n",
        "        latent_dim = mu.shape[1]\n",
        "\n",
        "        f1 = torch.sum(((samples -\n",
        "                         mu.view([batch_size, 1, latent_dim]))**2) /\n",
        "                       sigma.view([batch_size, 1, latent_dim])**2, dim=2)\n",
        "        f2 = mu.shape[1] * (math.log(2) + math.log(math.pi))\n",
        "        f3 = torch.sum(torch.log(sigma), dim=1).view(batch_size, 1)\n",
        "        return -0.5 * (f1 + f2) - f3\n",
        "\n",
        "    @staticmethod\n",
        "    def log_likelihood(x_true, x_distr):\n",
        "        \"\"\"\n",
        "        Calculate log likelihood between x_true and x_distr.\n",
        "        Args:\n",
        "            x_true:  Tensor - the matrix of shape batch_size x input_dim.\n",
        "            x_distr: Tensor - the tensor of shape batch_size x num_samples x input_dim, Bernoulli distribution parameters.\n",
        "        Returns:\n",
        "            Tensor - the matrix of shape batch_size x num_samples - log likelihood for each sample.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        batch_size = x_distr.shape[0]\n",
        "        input_dim = x_distr.shape[2]\n",
        "\n",
        "        bernoulli_log_likelihood = torch.log(\n",
        "            x_distr) * x_true.view([batch_size, 1, input_dim])\n",
        "        bernoulli_log_likelihood += torch.log(1 - x_distr) * (\n",
        "            1 - x_true).view([batch_size, 1, input_dim])\n",
        "\n",
        "        return torch.sum(bernoulli_log_likelihood, dim=2)\n",
        "\n",
        "    @staticmethod\n",
        "    def log_mean_exp(data):\n",
        "        \"\"\"\n",
        "        Args:\n",
        "            data: Tensor - the tensor of shape n_1 x n_2 x ... x n_K.\n",
        "        Returns:\n",
        "            Tensor - the tensor of shape n_1 x n_2 x ,,, x n_{K - 1}.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "\n",
        "        return torch.logsumexp(data, dim=-1) - \\\n",
        "            torch.log(torch.Tensor([data.shape[-1]]).to(data.device))\n",
        "\n",
        "    @staticmethod\n",
        "    def divergence_KL_normal(q_distr, p_distr):\n",
        "        \"\"\"\n",
        "        Calculate KL-divergence KL(q||p) between n-pairs of normal distribution.\n",
        "        Args:\n",
        "            q_distr=(mu, sigma): tuple(Tensor, Tensor) - the normal distribution parameters.\n",
        "                mu: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "                sigma: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "            p_distr=(mu, sigma): tuple(Tensor, Tensor) - the normal distribution parameters.\n",
        "                mu: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "                sigma: Tensor - the matrix of shape batch_size x latent_dim.\n",
        "        Returns:\n",
        "            Tensor - the vector of shape n, each value of which is a KL-divergence between pair of normal distribution.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "        q_mu, q_sigma = q_distr\n",
        "        p_mu, p_sigma = p_distr\n",
        "\n",
        "        D_KL = torch.sum((q_sigma / p_sigma)**2, dim=1)\n",
        "        D_KL -= p_mu.shape[1]\n",
        "        D_KL += 2 * torch.sum(torch.log(p_sigma), dim=1) - \\\n",
        "            2 * torch.sum(torch.log(q_sigma), dim=1)\n",
        "        D_KL += torch.sum((p_mu - q_mu) * (p_mu - q_mu) / (p_sigma**2), dim=1)\n",
        "        return 0.5 * D_KL\n",
        "\n",
        "    def forward(self, x):\n",
        "      \"\"\"\n",
        "        Generate decoded sample after encoding.\n",
        "        Args:\n",
        "            x: Tensor - the matrix of shape batch_size x input_dim.\n",
        "        Returns:\n",
        "            Tensor - the matrix of shape batch_size x input_dim.\n",
        "        Example:\n",
        "            >>>\n",
        "        \"\"\"\n",
        "      return self.q_x(self.sample_z(self.q_z(x))).view_as(x)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9WYJ_mwgmDeG"
      },
      "source": [
        "#### Скрипты для обучение VAE"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y5nqzt8fmN4M"
      },
      "source": [
        "def train_on_batch(model, x_batch, y_batch, optimizer, loss_function):\n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    loss = model.loss(x_batch.to(model.device), y_batch.to(model.device))\n",
        "    loss.backward()\n",
        "\n",
        "    optimizer.step()\n",
        "    return loss.cpu().item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "219UdCzZoE-4"
      },
      "source": [
        "#### Данные"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X-5KlkdwoQwM"
      },
      "source": [
        "preprocess = transforms.Compose([transforms.ToTensor(), \n",
        "                                 torch.FloatTensor, \n",
        "                                 lambda x: x.view(784)])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Iq4ZRNcroQwN"
      },
      "source": [
        "MNIST_train = datasets.MNIST('./mnist', train=True, download=True, \n",
        "                             transform=preprocess)\n",
        "\n",
        "MNIST_test = datasets.MNIST('./mnist', train=False, download=True,\n",
        "                            transform=preprocess)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S3rpGy8Ul0tG"
      },
      "source": [
        "#### Инициализация модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "m8mIbkWTkltO"
      },
      "source": [
        "optimizer = torch.optim.Adam\n",
        "loss_function = torch.nn.MSELoss()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Pgw8VSP7LRU"
      },
      "source": [
        "#### Перебор разного размера латентного пространства"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HNQrh5nNtOtq"
      },
      "source": [
        "for d in [64, 32, 16, 8, 4, 2]:\n",
        "    autoencoder = VAE(d, 28*28)\n",
        "    autoencoder.to(device)\n",
        "\n",
        "    writer = SummaryWriter(log_dir = 'autoencoder-vae/{}'.format(d))\n",
        "    call = callback(writer, MNIST_test, loss_function, delimeter = 100)\n",
        "\n",
        "    trainer(count_of_epoch=10, \n",
        "            batch_size=64, \n",
        "            dataset=MNIST_train,\n",
        "            model=autoencoder, \n",
        "            loss_function=None,\n",
        "            optimizer = optimizer,\n",
        "            lr = 0.001,\n",
        "            callback = call)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RWJ-dhwSuKNI"
      },
      "source": [
        "#### Отрисовка изменения (возмущаем скрытые вектора)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gtG7ELKfs6f0"
      },
      "source": [
        "def draw_samples_grid_vae(model,\n",
        "                          num_row=15,\n",
        "                          num_colum=15,\n",
        "                          images_size=(28, 28)):\n",
        "    \"\"\"\n",
        "    Illustrate how change digits x where change point in latent space z.\n",
        "    Args:\n",
        "        model: nn.Module - model VAE or IWAE.\n",
        "        num_row: int - the number of row.\n",
        "        num_colum: int - the number of column.\n",
        "        images_size = (x_size, y_size): tuple(int, int) - a size of input image.\n",
        "    Returns:\n",
        "        figure: float - the picture\n",
        "    Example:\n",
        "        >>>\n",
        "    \"\"\"\n",
        "    grid_x = norm.ppf(np.linspace(0.05, 0.95, num_colum))\n",
        "    grid_y = norm.ppf(np.linspace(0.05, 0.95, num_row))\n",
        "\n",
        "    figure = np.zeros((images_size[0] * num_colum, images_size[1] * num_row))\n",
        "    for i, y_i in enumerate(grid_x):\n",
        "        for j, x_i in enumerate(grid_y):\n",
        "            z_sample = np.array([[x_i, y_i]])\n",
        "\n",
        "            x_sample = model.q_x(torch.from_numpy(z_sample).float()).view(\n",
        "                images_size).cpu().data.numpy()\n",
        "\n",
        "            image = x_sample\n",
        "            figure[i * images_size[0]: (i + 1) * images_size[0],\n",
        "                   j * images_size[1]: (j + 1) * images_size[1]] = image\n",
        "\n",
        "    return figure"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OizKj4aDs5Qm"
      },
      "source": [
        "figure = draw_samples_grid_vae(autoencoder)\n",
        "\n",
        "plt.figure(figsize=(10, 10))\n",
        "plt.imshow(figure, cmap='Greys_r', interpolation='lanczos')\n",
        "plt.axis('off')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aDCSUeWyMSPW"
      },
      "source": [
        "## Transfer Learning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SZcwXF6zMSPX"
      },
      "source": [
        "### Предобученые модели"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oa63o8K14729"
      },
      "source": [
        "#### Загрузим данные котиков и собачек"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gucf9x_O1TY_"
      },
      "source": [
        "!wget -q https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip\n",
        "!unzip -qq cats_and_dogs_filtered.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qk19Hyx92uk2"
      },
      "source": [
        "preprocess = transforms.Compose([\n",
        "    transforms.Resize(256),\n",
        "    transforms.CenterCrop(224),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
        "])\n",
        "\n",
        "CatDogs_train = datasets.ImageFolder(root = 'cats_and_dogs_filtered/train/', \n",
        "                                     transform = preprocess)\n",
        "CatDogs_test = datasets.ImageFolder(root = 'cats_and_dogs_filtered/validation/', \n",
        "                                     transform = preprocess)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cphDN4B-5GgX"
      },
      "source": [
        "#### Загрузим предобученую модель"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SKz0b2eY3wxm"
      },
      "source": [
        "model = torch.hub.load('pytorch/vision:v0.6.0', 'resnet18', pretrained=True, verbose=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zJEvdqlE3_9v"
      },
      "source": [
        "for param in model.parameters():\n",
        "    param.requires_grad = False\n",
        "\n",
        "model.fc = torch.nn.Linear(512, 2)\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MAf5wHRn5Md5"
      },
      "source": [
        "#### Скрипты для обучения"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ziCx4rr543zh"
      },
      "source": [
        "def train_on_batch(model, x_batch, y_batch, optimizer, loss_function):\n",
        "    model.train()\n",
        "    optimizer.zero_grad()\n",
        "    \n",
        "    output = model(x_batch.to(device))\n",
        "    \n",
        "    loss = loss_function(output, y_batch.to(device))\n",
        "    loss.backward()\n",
        "\n",
        "    optimizer.step()\n",
        "    return loss.cpu().item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mIrNQNaK8aDn"
      },
      "source": [
        "class callback():\n",
        "    def __init__(self, writer, dataset, loss_function, delimeter = 100, batch_size=64):\n",
        "        self.step = 0\n",
        "        self.writer = writer\n",
        "        self.delimeter = delimeter\n",
        "        self.loss_function = loss_function\n",
        "        self.batch_size = batch_size\n",
        "\n",
        "        self.dataset = dataset\n",
        "\n",
        "    def forward(self, model, loss):\n",
        "        self.step += 1\n",
        "        self.writer.add_scalar('LOSS/train', loss, self.step)\n",
        "        \n",
        "        if self.step % self.delimeter == 0:\n",
        "            \n",
        "            batch_generator = torch.utils.data.DataLoader(dataset = self.dataset, \n",
        "                                                          batch_size=self.batch_size, \n",
        "                                                          pin_memory=True)\n",
        "            \n",
        "            pred = []\n",
        "            real = []\n",
        "            test_loss = 0\n",
        "            model.eval()\n",
        "            for it, (x_batch, y_batch) in enumerate(batch_generator):\n",
        "                x_batch = x_batch.to(device)\n",
        "\n",
        "                output = model(x_batch)\n",
        "\n",
        "                test_loss += self.loss_function(output, y_batch.to(device)).cpu().item()*len(x_batch)\n",
        "\n",
        "                pred.extend(torch.argmax(output, dim=-1).cpu().numpy().tolist())\n",
        "                real.extend(y_batch.cpu().numpy().tolist())\n",
        "            \n",
        "            test_loss /= len(self.dataset)\n",
        "            \n",
        "            self.writer.add_scalar('LOSS/test', test_loss, self.step)\n",
        "            self.writer.add_text('REPORT/test', str(classification_report(real, pred)), self.step)\n",
        "          \n",
        "    def __call__(self, model, loss):\n",
        "        return self.forward(model, loss)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ENuUwKwS5ZXt"
      },
      "source": [
        "#### Инициализация модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "699WSDhF5YW4"
      },
      "source": [
        "loss_function = torch.nn.CrossEntropyLoss()\n",
        "optimizer = torch.optim.Adam"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "afyLgrcz78Zq"
      },
      "source": [
        "#### Качество до дообучения"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K5NsO1Gu5mMg"
      },
      "source": [
        "real = []\n",
        "pred = []\n",
        "for x_batch, y_batch in torch.utils.data.DataLoader(CatDogs_test, \n",
        "                                                    batch_size = 64, \n",
        "                                                    shuffle = True, \n",
        "                                                    pin_memory=True):\n",
        "  \n",
        "    with torch.no_grad():\n",
        "        output = model(x_batch.to(device))\n",
        "    pred.extend(output.argmax(dim=-1).cpu().numpy().tolist())\n",
        "    real.extend(y_batch.cpu().numpy().tolist())\n",
        "\n",
        "print(classification_report(real, pred))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TyE1UxYo7Vlo"
      },
      "source": [
        "#### Обучение модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FLpm0m6f566f"
      },
      "source": [
        "writer = SummaryWriter(log_dir = 'transfer-resnet18')\n",
        "call = callback(writer, CatDogs_test, loss_function, delimeter = 10)\n",
        "\n",
        "trainer(count_of_epoch=10, \n",
        "        batch_size=64, \n",
        "        dataset=CatDogs_train,\n",
        "        model=model,\n",
        "        loss_function=loss_function,\n",
        "        optimizer = optimizer,\n",
        "        lr = 0.001,\n",
        "        callback = call)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hZhr6ml974XK"
      },
      "source": [
        "#### Качество после дообучения"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "x9FPVrYV76gO"
      },
      "source": [
        "real = []\n",
        "pred = []\n",
        "for x_batch, y_batch in torch.utils.data.DataLoader(CatDogs_test, \n",
        "                                                    batch_size = 64, \n",
        "                                                    shuffle = True, \n",
        "                                                    pin_memory=True):\n",
        "  \n",
        "    with torch.no_grad():\n",
        "        output = model(x_batch.to(device))\n",
        "    pred.extend(output.argmax(dim=-1).cpu().numpy().tolist())\n",
        "    real.extend(y_batch.cpu().numpy().tolist())\n",
        "\n",
        "print(classification_report(real, pred))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V1COxOqzrHds"
      },
      "source": [
        "## Генеративно состязательные сети"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YJyO77mxMSPX"
      },
      "source": [
        "### GAN's для MNIST"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u8W32KDH_GBu"
      },
      "source": [
        "#### Полезный код для обучения"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0UiLjhK2TmwH"
      },
      "source": [
        "def train_on_batch(model, x_batch, y_batch, optimizer, loss_function):\n",
        "    optima_generator, optima_discriminator = optimizer\n",
        "    \n",
        "    optima_generator.zero_grad()\n",
        "    optima_discriminator.zero_grad()\n",
        "    discriminator_loss = model.discriminator_loss(x_batch)\n",
        "    discriminator_loss.backward()\n",
        "    optima_discriminator.step()\n",
        "\n",
        "    optima_generator.zero_grad()\n",
        "    optima_discriminator.zero_grad()\n",
        "    generator_loss = model.generator_loss(len(x_batch))\n",
        "    generator_loss.backward()\n",
        "    optima_generator.step()\n",
        "    \n",
        "    return discriminator_loss.cpu().item()+generator_loss.cpu().item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "e9zK81LZTmwH"
      },
      "source": [
        "def train_epoch(train_generator, model, loss_function, optimizer, callback = None):\n",
        "    epoch_loss = 0\n",
        "    total = 0\n",
        "    for it, (batch_of_x, batch_of_y) in enumerate(train_generator):\n",
        "        batch_loss = train_on_batch(model, batch_of_x.to(model.device), batch_of_y.to(model.device), optimizer, loss_function)\n",
        "        \n",
        "        if callback is not None:\n",
        "            callback(model, batch_loss)\n",
        "            \n",
        "        epoch_loss += batch_loss*len(batch_of_x)\n",
        "        total += len(batch_of_x)\n",
        "    \n",
        "    return epoch_loss/total"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Eub4ku9lTmwH"
      },
      "source": [
        "def trainer(count_of_epoch, \n",
        "            batch_size, \n",
        "            dataset,\n",
        "            model, \n",
        "            loss_function,\n",
        "            optimizer,\n",
        "            lr = 0.001,\n",
        "            callback = None):\n",
        "\n",
        "    optima_generator = optimizer(model.generator.parameters(), lr=lr*10, betas=(0.5, 0.999))\n",
        "    optima_discriminator = optimizer(model.discriminator.parameters(), lr=lr, betas=(0.5, 0.999))\n",
        "    \n",
        "    iterations = tqdm(range(count_of_epoch), desc='epoch')\n",
        "    iterations.set_postfix({'epoch loss': np.nan})\n",
        "    for it in iterations:        \n",
        "        batch_generator = tqdm(\n",
        "            torch.utils.data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True), \n",
        "            leave=False, total=len(dataset)//batch_size+(len(dataset)%batch_size> 0))\n",
        "        \n",
        "        epoch_loss = train_epoch(\n",
        "            train_generator=batch_generator, \n",
        "            model=model, \n",
        "            loss_function=loss_function, \n",
        "            optimizer=(optima_generator, optima_discriminator), \n",
        "            callback=callback)\n",
        "        \n",
        "        iterations.set_postfix({'epoch loss': epoch_loss})"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t16nU2d7_GBu"
      },
      "source": [
        "#### Модель нейросети"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CCukR2naTmwH"
      },
      "source": [
        "class Reshape(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, *args):\n",
        "        super(type(self), self).__init__()\n",
        "        self.dims = args\n",
        "\n",
        "    def forward(self, input):\n",
        "        return input.view(input.size(0), *self.dims)\n",
        "\n",
        "class GAN(torch.nn.Module):\n",
        "    @property\n",
        "    def device(self):\n",
        "        return next(self.parameters()).device\n",
        "\n",
        "    def __init__(self, d, D):\n",
        "        super(type(self), self).__init__()\n",
        "        self.d = d\n",
        "        self.discriminator = torch.nn.Sequential(\n",
        "            Reshape(1, D, D),\n",
        "            torch.nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=2, padding=1, bias=False),\n",
        "            torch.nn.BatchNorm2d(num_features=64),\n",
        "            torch.nn.LeakyReLU(negative_slope=0.2),\n",
        "            torch.nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, stride=2, padding=0, bias=False),\n",
        "            torch.nn.BatchNorm2d(num_features=32),\n",
        "            torch.nn.LeakyReLU(negative_slope=0.2),\n",
        "            torch.nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, stride=1, padding=0, bias=False),\n",
        "            torch.nn.BatchNorm2d(num_features=16),\n",
        "            torch.nn.LeakyReLU(negative_slope=0.2),\n",
        "            Reshape(-1),\n",
        "            torch.nn.Linear(in_features=256, out_features=1),\n",
        "            torch.nn.Sigmoid()\n",
        "        )\n",
        "        self.generator = torch.nn.Sequential(\n",
        "            Reshape(self.d, 1, 1),\n",
        "            torch.nn.ConvTranspose2d(self.d, 128, 4, 1, 0, 0, bias=False),\n",
        "            torch.nn.BatchNorm2d(128),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.ConvTranspose2d(128, 64, 3, 2, 1, 0, bias=False),\n",
        "            torch.nn.BatchNorm2d(64),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.ConvTranspose2d(64, 32, 3, 2, 1, 1, bias=False),\n",
        "            torch.nn.BatchNorm2d(32),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.ConvTranspose2d(32, 1, 4, 2, 1, 0, bias=False),\n",
        "            Reshape(-1),\n",
        "            torch.nn.Sigmoid()\n",
        "        )\n",
        "\n",
        "    def generate_noise(self, num_samples):\n",
        "        return torch.randn([num_samples, self.d], requires_grad=True).to(self.device)\n",
        "\n",
        "    def generate_samples(self, num_samples):\n",
        "        return self.generator(self.generate_noise(num_samples))\n",
        "\n",
        "    def discriminator_loss(self, batch):\n",
        "        loss = torch.nn.BCELoss()\n",
        "        \n",
        "        batch_size = batch.shape[0]\n",
        "        \n",
        "        fake_batch = self.discriminator(self.generate_samples(batch_size).detach())\n",
        "        real_batch = self.discriminator(batch)\n",
        "        \n",
        "        fake_answ = torch.zeros(batch_size, 1).to(self.device)\n",
        "        real_answ = torch.ones(batch.shape[0], 1).to(self.device)\n",
        "        \n",
        "        return 0.5*loss(fake_batch, fake_answ) + 0.5*loss(real_batch, real_answ)\n",
        "\n",
        "    def generator_loss(self, batch_size):\n",
        "        loss = torch.nn.BCELoss()\n",
        "        \n",
        "        fake_batch = self.discriminator(self.generate_samples(batch_size))\n",
        "        \n",
        "        real_answ = torch.ones(batch_size, 1).to(self.device)\n",
        "            \n",
        "        return loss(fake_batch, real_answ)\n",
        "    \n",
        "    def init_weight(self, mean, std):\n",
        "        for m in self._modules:\n",
        "            if isinstance(self._modules[m], nn.ConvTranspose2d) or isinstance(self._modules[m], nn.Conv2d):\n",
        "                self._modules[m].weight.data.normal_(mean, std)\n",
        "                self._modules[m].bias.data.zero_()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1mjMCVWT_GBv"
      },
      "source": [
        "#### Инициализация модели"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VUcXGn0JTmwI"
      },
      "source": [
        "loss_function = None\n",
        "optimizer = torch.optim.Adam\n",
        "\n",
        "model = GAN(2, 28)\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Rx_4Yqf_GBv"
      },
      "source": [
        "#### Выборка"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NBR6IpTZUYzm"
      },
      "source": [
        "MNIST_train = datasets.MNIST('./mnist', train=True, download=True, \n",
        "                             transform=transforms.ToTensor())\n",
        "\n",
        "MNIST_test = datasets.MNIST('./mnist', train=False, download=True,\n",
        "                            transform=transforms.ToTensor())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9P7mM9_z_GBv"
      },
      "source": [
        "#### Вспомагательные функции"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UWs6Ph14TmwI"
      },
      "source": [
        "def show_images(x, digit_size=28):\n",
        "    plt.figure(figsize=(12, 12 / 10 * (x.shape[0] // 10 + 1)))\n",
        "    x = x.view(-1, digit_size, digit_size)\n",
        "    for i in range(x.shape[0]):\n",
        "        plt.subplot(x.shape[0] // 10 + 1, 10, i + 1)\n",
        "        plt.imshow(x.data[i].numpy(), cmap='Greys_r', interpolation='lanczos')\n",
        "        plt.axis('off')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "K0EiwEJNTmwI"
      },
      "source": [
        "def generate_samples(model, num_samples, batch_size):\n",
        "    size = 0\n",
        "    res = []\n",
        "    while size < num_samples:\n",
        "        res.append(\n",
        "            model.generate_samples(\n",
        "                min(batch_size, num_samples - size)))\n",
        "        size += batch_size\n",
        "    return torch.cat(res, 0).cpu().detach()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IUicTyHD_GBw"
      },
      "source": [
        "#### Как генерируются данные до обучения"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6L9QVf-9TmwI"
      },
      "source": [
        "show_images(generate_many_samples(model, 10, 10))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IuOPesh__GBw"
      },
      "source": [
        "#### Обучаем модель"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g40_ilnoTmwI"
      },
      "source": [
        "trainer(count_of_epoch=10, \n",
        "        batch_size=64, \n",
        "        dataset=MNIST_train,\n",
        "        model=model, \n",
        "        loss_function=loss_function,\n",
        "        optimizer=optimizer,\n",
        "        lr = 2e-4,\n",
        "        callback = None)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RL_EWILf_GBx"
      },
      "source": [
        "#### Как генерируются данные после обучения"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6W7NhdM_Uc7n"
      },
      "source": [
        "show_images(generate_many_samples(model, 10, 50))"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}