{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### VAE",
   "id": "dfa25c27711ef214"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Создадим датасет с кругами, как в лабораторной 1",
   "id": "47036a180392df32"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "def make_circles(\n",
    "    n_per_class: int = 500,\n",
    "    r_inner: float = 2.0,\n",
    "    r_outer: float = 5.0,\n",
    "    noise_std: float = 0.25,\n",
    "    seed: int = 42,\n",
    ") -> tuple[np.ndarray, np.ndarray]:\n",
    "    \"\"\"\n",
    "    Генерирует два класса точек: внутренний круг (класс 0) и внешний круг (класс 1), каждый со своим радиусом и гауссовым шумом.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    n_per_class : Число точек на каждый класс.\n",
    "    r_inner :     Радиус внутреннего круга.\n",
    "    r_outer :     Радиус внешнего круга.\n",
    "    noise_std :   Стандартное отклонение аддитивного шума по каждой оси.\n",
    "    seed :        Зерно ГСЧ для воспроизводимости.\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    X : np.ndarray, shape (2 * n_per_class, 2)\n",
    "        Координаты точек.\n",
    "    y : np.ndarray, shape (2 * n_per_class,)\n",
    "    \"\"\"\n",
    "    rng = np.random.default_rng(seed)\n",
    "\n",
    "    # углы равномерно на [0, 2\\pi)\n",
    "    theta = rng.uniform(0, 2 * np.pi, size=n_per_class)\n",
    "\n",
    "    # радиусы с шумом\n",
    "    r1 = r_inner + rng.normal(0, noise_std, size=n_per_class)\n",
    "    r2 = r_outer + rng.normal(0, noise_std, size=n_per_class)\n",
    "\n",
    "    # полярные → декартовы координаты\n",
    "    X1 = np.column_stack([r1 * np.cos(theta), r1 * np.sin(theta)])\n",
    "    X2 = np.column_stack([r2 * np.cos(theta), r2 * np.sin(theta)])\n",
    "\n",
    "    X = np.vstack([X1, X2])\n",
    "    y = np.array([0] * n_per_class + [1] * n_per_class)\n",
    "\n",
    "    # перемешиваем\n",
    "    perm = rng.permutation(len(X))\n",
    "    return X[perm], y[perm]"
   ],
   "id": "5d1fd1d27d69c5a2",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "X, y = make_circles(n_per_class=1000, noise_std=0.2, seed=42)\n",
    "\n",
    "plt.scatter(X[:, 0], X[:, 1], c=y, alpha=0.7, s=12, edgecolors=\"k\", linewidths=0.3)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ],
   "id": "91e482bd6313867a",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "X_tensor = torch.tensor(X, dtype=torch.float32)          # (N, 2)\n",
    "y_tensor = torch.tensor(y, dtype=torch.float32).unsqueeze(1)  # (N, 1)"
   ],
   "id": "2f5ffa1c1b41f4a6",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Определим VAE. Опишите, что происходит в каждом блоке кода.",
   "id": "6f48eeeb2ff66d7d"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "class VAE(nn.Module):\n",
    "    def __init__(self, input_dim, hidden_dim=64, latent_dim=2):\n",
    "        super().__init__()\n",
    "\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(input_dim, hidden_dim),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        self.mu = nn.Linear(hidden_dim, latent_dim)\n",
    "        self.logvar = nn.Linear(hidden_dim, latent_dim)\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(latent_dim, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, input_dim)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        h = self.encoder(x)\n",
    "\n",
    "        mu = self.mu(h)\n",
    "        logvar = self.logvar(h)\n",
    "\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        z = mu + std * torch.randn_like(std)\n",
    "\n",
    "        recon = self.decoder(z)\n",
    "        return recon, mu, logvar"
   ],
   "id": "242def3011983d5e",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Напишем обучающий скрипт для простого VAE",
   "id": "94da5afae7e2e7d6"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "epochs = 500\n",
    "batch_size = 64\n",
    "\n",
    "model = VAE(\n",
    "    input_dim=X_tensor.shape[1],\n",
    "    hidden_dim=64,\n",
    "    latent_dim=2\n",
    ").to(DEVICE)\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "criterion = nn.MSELoss(reduction=\"sum\")\n",
    "\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "    # TRAIN\n",
    "    model.train()\n",
    "    epoch_loss = 0.0\n",
    "\n",
    "    for i in range(0, X_tensor.shape[0], batch_size):\n",
    "        x_batch = X_tensor[i:i + batch_size].to(DEVICE)\n",
    "\n",
    "        recon, mu, logvar = model(x_batch)\n",
    "        recon_loss = criterion(recon, x_batch)\n",
    "\n",
    "        kl_loss = -0.5 * torch.sum(\n",
    "            1 + logvar - mu.pow(2) - logvar.exp()\n",
    "        )  # Что здесь происходит?\n",
    "\n",
    "        loss = (recon_loss + kl_loss) / x_batch.size(0)  # Что такое recon_loss и kl_loss\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        epoch_loss += loss.item() * x_batch.size(0)\n",
    "\n",
    "    if (epoch + 1) % 50 == 0:\n",
    "        print(\n",
    "            f\"Epoch {epoch+1:3d} | \"\n",
    "            f\"loss = {epoch_loss / X_tensor.shape[0]:.4f}\"\n",
    "        )"
   ],
   "id": "3ecc826f2249ef3f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Сгенерируйте и отобразите полученные точки вместе с оригинальными. Получилось ли воссоздать изначальное распределение? Почему?",
   "id": "8b65fbea276f013f"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "fd7a8c097b3c8d50",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### CVAE",
   "id": "ab461086c30203bb"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Зададим Conditional VAE, чтобы генерировать точки каждого класса отдельно. Расскажите, что изменилось?",
   "id": "ad140dcbcbe34f1b"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "class CVAE(nn.Module):\n",
    "    def __init__(self, input_dim=2, hidden_dim=64, latent_dim=2, num_classes=2):\n",
    "        super().__init__()\n",
    "        self.num_classes = num_classes\n",
    "        self.latent_dim = latent_dim\n",
    "\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(input_dim + num_classes, hidden_dim),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        self.mu = nn.Linear(hidden_dim, latent_dim)\n",
    "        self.logvar = nn.Linear(hidden_dim, latent_dim)\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(latent_dim + num_classes, hidden_dim),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_dim, input_dim)\n",
    "        )\n",
    "\n",
    "    def condition(self, y):\n",
    "        return F.one_hot(\n",
    "            y.long().reshape(-1), num_classes=self.num_classes\n",
    "        ).float()\n",
    "\n",
    "    def decode(self, z, y):\n",
    "        condition = self.condition(y)\n",
    "        return self.decoder(torch.cat([z, condition], dim=1))\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        condition = self.condition(y)\n",
    "        h = self.encoder(torch.cat([x, condition], dim=1))\n",
    "\n",
    "        mu = self.mu(h)\n",
    "        logvar = self.logvar(h)\n",
    "\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        z = mu + std * torch.randn_like(std)\n",
    "\n",
    "        recon = self.decode(z, y)\n",
    "        return recon, mu, logvar"
   ],
   "id": "23ba6fc5a693321c",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Напишите обучающий скрипт на основе того, что выше.",
   "id": "7239f7a458921ff8"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "97101c4112100cec",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Снова отобразите сгенерированные точки и оригинальные. Какие есть проблемы? Попробуйте изменить архитектуру таким образом, чтобы решить их",
   "id": "401b86412bb1c31e"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "b86da1195efbf5e",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### VAE MNIST",
   "id": "98097d960528ad33"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Загрузим уже известный нам датасет MNIST",
   "id": "8d7004db543d31a9"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "from torchvision import datasets, transforms\n",
    "\n",
    "dataset = datasets.MNIST(\n",
    "    root=\"./data\",\n",
    "    train=True,\n",
    "    download=True,\n",
    "    transform=transforms.ToTensor()\n",
    ")\n",
    "\n",
    "loader = DataLoader(dataset, batch_size=128, shuffle=True)"
   ],
   "id": "440be934303f2980",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Модифицируйте простейший VAE, определенный выше, для задачи MNIST. И обучите его",
   "id": "9f3647069bbf0a21"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "3946059ed5b3d6cb",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "71ef45e0cee08cf4",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Теперь выведите несколько генераций для каждой метки. Оцените визуально, насколько эти генерации приемлемы?",
   "id": "b0f7b506c57271e5"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "bb5a518560512eb3",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Теперь давайте изменим модель MNIST, заменив линейные слои на сверточные. Ниже приведен пример такой сети. Вы можете изменять слои по своему усмотрению. Обучите эту сеть и выведите результат",
   "id": "897ac90bcea4d5c2"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "class CVAE_MNIST(nn.Module):\n",
    "    def __init__(self, latent_dim=32, num_classes=10):\n",
    "        super().__init__()\n",
    "        self.latent_dim = latent_dim\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "        # (B, 1, 28, 28) → (B, 128, 7, 7)\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Conv2d(1, 32, kernel_size=3, padding=1),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Flatten()\n",
    "        )\n",
    "\n",
    "        self.encoder_fc = nn.Sequential(\n",
    "            nn.Linear(128 * 7 * 7 + num_classes, 256),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        self.mu = nn.Linear(256, latent_dim)\n",
    "        self.logvar = nn.Linear(256, latent_dim)\n",
    "\n",
    "        self.decoder_fc = nn.Sequential(\n",
    "            nn.Linear(latent_dim + num_classes, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 128 * 7 * 7),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.ConvTranspose2d(\n",
    "                128, 64, kernel_size=4, stride=2, padding=1\n",
    "            ),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.ConvTranspose2d(\n",
    "                64, 32, kernel_size=4, stride=2, padding=1\n",
    "            ),\n",
    "            nn.ReLU(),\n",
    "\n",
    "            nn.Conv2d(32, 1, kernel_size=3, padding=1),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def condition(self, y):\n",
    "        return F.one_hot(\n",
    "            y.long(), num_classes=self.num_classes\n",
    "        ).float()\n",
    "\n",
    "    def decode(self, z, y):\n",
    "        condition = self.condition(y)\n",
    "\n",
    "        h = self.decoder_fc(torch.cat([z, condition], dim=1))\n",
    "        h = h.reshape(-1, 128, 7, 7)\n",
    "\n",
    "        return self.decoder(h).flatten(start_dim=1)\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        x = x.reshape(-1, 1, 28, 28)\n",
    "        condition = self.condition(y)\n",
    "\n",
    "        h = self.encoder(x)\n",
    "        h = self.encoder_fc(torch.cat([h, condition], dim=1))\n",
    "\n",
    "        mu = self.mu(h)\n",
    "        logvar = self.logvar(h)\n",
    "\n",
    "        std = torch.exp(0.5 * logvar)\n",
    "        z = mu + std * torch.randn_like(std)\n",
    "\n",
    "        recon = self.decode(z, y)\n",
    "        return recon, mu, logvar"
   ],
   "id": "b4ecddc22f56896f",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Проверьте, насколько хорошо VAE справляется с шумом. Добавьте шум в изображения цифр и проведите цифру через обученный VAE. Насколько полученная цифра похожа на исходную? Измерьте метрику MSE",
   "id": "ae4db27cdafcf58d"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "dceb0235b99bb502"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Теперь возьмите две цифры с одинаковой меткой из датасета и выведите 8 изображений, находящихся на прямой между ними в пиксельном пространстве. То есть рассматривайте изображения как векторы в 28x28-мерном пространстве. Сохраняется ли метка при таком переходе?",
   "id": "8e34740bb30ed9e8"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "64f4f27aa147f3e1",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Теперь сначала закодируйте эти две цифры с помощью encoder-а VAE. И перейдите от многомерного 28х28-мерного пространства в латентное пространство меньшей размерности. Найдите в нем 8 векторов, лежащих на прямой между выбранными элементами и декодируйте их через VAE decoder. Насколько сгенерированные изображения похожи на изначальные цифры?",
   "id": "84db4e3130bfc486"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "6c77248085e5f9d8",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "Теперь осуществите перенос почерка (стиля). Для этого закодируйте цифру с одной меткой и из полученного вектора с помощью decoder-а сгенерируйте цифры с другими метками. Насколько похож почерк? Какие особенности есть у сгенерированных цифр? С чем они связаны?",
   "id": "21875a5b3637790b"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "",
   "id": "24574e1ffcb15588",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}