examples/vae.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "if 'has_changed_dir' not in globals():\n",
    "    repo_path = os.path.abspath(os.path.join('..'))\n",
    "    \n",
    "    if repo_path not in sys.path:\n",
    "        sys.path.append(repo_path)\n",
    "    \n",
    "    os.chdir(repo_path)\n",
    "    \n",
    "    globals()['has_changed_dir'] = True\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from tqdm import tqdm\n",
    "\n",
    "import neunet as nnet\n",
    "import neunet.nn as nn\n",
    "from data_loader import load_mnist\n",
    "from neunet.optim import Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "noisy_inputs = False\n",
    "\n",
    "image_size = (1, 28, 28)\n",
    "x_num, y_num = 5, 5\n",
    "samples_num = x_num * y_num\n",
    "margin = 15\n",
    "\n",
    "\n",
    "def add_noise(data):\n",
    "    noise_factor = 0.5\n",
    "\n",
    "    noisy_data = data + noise_factor * np.random.normal(0, 1, (data.shape))\n",
    "\n",
    "    return np.clip(noisy_data, 0, 1)\n",
    "\n",
    "\n",
    "training_data, test_data, training_labels, test_labels = load_mnist()\n",
    "training_data = training_data / 255  # normalization: / 255 => [0; 1]  #/ 127.5-1 => [-1; 1]\n",
    "test_data = test_data / 255  # normalization: / 255 => [0; 1]  #/ 127.5-1 => [-1; 1]\n",
    "\n",
    "latent_size = 2\n",
    "\n",
    "device = \"cpu\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class VAE(nn.Module):\n",
    "    def __init__(self, input_size, latent_size):\n",
    "        super().__init__()\n",
    "        self.input_size = input_size\n",
    "        self.latent_size = latent_size\n",
    "\n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.Linear(input_size, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(512),\n",
    "            nn.Linear(512, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(256),\n",
    "            nn.Linear(256, latent_size),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(latent_size),\n",
    "        )\n",
    "\n",
    "        self.decoder = nn.Sequential(\n",
    "            nn.Linear(latent_size, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(256),\n",
    "            nn.Linear(256, 512),\n",
    "            nn.ReLU(),\n",
    "            nn.BatchNorm1d(512),\n",
    "            nn.Linear(512, input_size),\n",
    "            nn.Sigmoid(),\n",
    "        )\n",
    "        self.mu_encoder = nn.Linear(latent_size, latent_size)\n",
    "        self.logvar_encoder = nn.Linear(latent_size, latent_size)\n",
    "\n",
    "        self.loss_fn = nn.BCELoss(reduction=\"sum\")\n",
    "\n",
    "    def reparameterize(self, mu, logvar):\n",
    "        std = logvar.mul(0.5).exp()\n",
    "        eps = nnet.tensor(np.random.normal(0, 1, size=std.shape), device=device)\n",
    "        z = mu + eps * std\n",
    "        return z\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.encoder(x)\n",
    "\n",
    "        # mu, logvar = x_enc[:, :self.latent_size], x_enc[:, self.latent_size:] #\n",
    "        mu = self.mu_encoder(x)\n",
    "        logvar = self.logvar_encoder(x)\n",
    "\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "\n",
    "        return self.decoder(z), mu, logvar\n",
    "\n",
    "    def loss_function(self, x, x_recon, mu, logvar):\n",
    "        BCE = self.loss_fn(x_recon, x)\n",
    "        KLD = -0.5 * nnet.sum(1 + logvar - mu.power(2) - logvar.exp())\n",
    "        return BCE + KLD\n",
    "\n",
    "    def train_step(self, in_x, out_x, optimizer):\n",
    "        x_recon, mu, logvar = self.forward(in_x)\n",
    "\n",
    "        loss = self.loss_function(out_x, x_recon, mu, logvar)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        return loss\n",
    "\n",
    "    def encode(self, x):\n",
    "        x = self.encoder(x)\n",
    "\n",
    "        mu = self.mu_encoder(x)\n",
    "        logvar = self.logvar_encoder(x)\n",
    "\n",
    "        z = self.reparameterize(mu, logvar)\n",
    "\n",
    "        return z\n",
    "\n",
    "    def decode(self, z):\n",
    "        return self.decoder(z)\n",
    "\n",
    "    def reconstruct(self, x):\n",
    "        return self.forward(x)[0]\n",
    "\n",
    "\n",
    "vae = VAE(28 * 28, latent_size).to(device)\n",
    "optimizer = Adam(vae.parameters(), lr=0.0005)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 100\n",
    "epochs = 20\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    tqdm_range = tqdm(range(0, len(training_data), batch_size), desc=\"epoch %d\" % epoch)\n",
    "    vae.train()\n",
    "    for i in tqdm_range:\n",
    "        batch = training_data[i : i + batch_size]\n",
    "\n",
    "        in_batch = nnet.tensor(batch, device=device).reshape(-1, 28 * 28)\n",
    "        if noisy_inputs:\n",
    "            in_batch = nnet.tensor(add_noise(in_batch.cpu().numpy()), device=device)\n",
    "\n",
    "        out_batch = nnet.tensor(batch, device=device).reshape(-1, 28 * 28)\n",
    "\n",
    "        loss = vae.train_step(in_batch, out_batch, optimizer)\n",
    "\n",
    "        tqdm_range.set_description(f\"epoch: {epoch + 1}/{epochs}, loss: {loss.item():.7f}\")\n",
    "\n",
    "    generated = (\n",
    "        vae.decode(\n",
    "            nnet.tensor(\n",
    "                np.random.normal(0, 1, size=(samples_num, latent_size)),\n",
    "                device=device,\n",
    "            )\n",
    "        )\n",
    "        .to(\"cpu\")\n",
    "        .detach()\n",
    "        .numpy()\n",
    "    )\n",
    "\n",
    "    # samples = training_data[np.random.randint(0, len(training_data), samples_num)]\n",
    "    # if noisy_inputs:\n",
    "    #     samples = add_noise(samples)\n",
    "    # generated = vae.reconstruct(nnet.tensor(samples, requires_grad=False).reshape(-1, 28 * 28)).data\n",
    "    vae.eval()\n",
    "    for i in range(25):\n",
    "        image = generated[i] * 255\n",
    "        image = image.astype(np.uint8)\n",
    "        image = image.reshape(28, 28)\n",
    "        image = Image.fromarray(image)\n",
    "        image.save(f\"generated images/{i}.png\")\n",
    "\n",
    "vae.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_images_set(images):\n",
    "    images_array = np.full(\n",
    "        (y_num * (margin + image_size[1]), x_num * (margin + image_size[2])),\n",
    "        255,\n",
    "        dtype=np.uint8,\n",
    "    )\n",
    "    num = 0\n",
    "    for i in range(y_num):\n",
    "        for j in range(x_num):\n",
    "            y = i * (margin + image_size[1])\n",
    "            x = j * (margin + image_size[2])\n",
    "\n",
    "            images_array[y : y + image_size[1], x : x + image_size[2]] = images[num]\n",
    "            num += 1\n",
    "\n",
    "    images_array = images_array[\n",
    "        : (y_num - 1) * (image_size[1] + margin) + image_size[1],\n",
    "        : (x_num - 1) * (image_size[2] + margin) + image_size[2],\n",
    "    ]\n",
    "\n",
    "    return Image.fromarray(images_array).convert(\"L\")\n",
    "\n",
    "\n",
    "samples = test_data[np.random.randint(0, len(test_data), samples_num)]\n",
    "if noisy_inputs:\n",
    "    samples = add_noise(samples)\n",
    "generated = (\n",
    "    vae.reconstruct(nnet.tensor(samples, device=device).reshape(-1, 28 * 28))\n",
    "    .to(\"cpu\")\n",
    "    .detach()\n",
    "    .numpy()\n",
    ")\n",
    "\n",
    "get_images_set(samples.reshape(-1, 28, 28) * 255).save(\"generated images/vae_in_samples.jpeg\")\n",
    "get_images_set(generated.reshape(-1, 28, 28) * 255).save(\"generated images/vae_out_samples.jpeg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Visualize latent space only with latent_dim = 2\"\"\"\n",
    "\n",
    "\n",
    "def plot_latent_space_digits(n=30, figsize=15):\n",
    "    if latent_size != 2:\n",
    "        print(\"Can`t plot 2d latent space for non-2d latent space\")\n",
    "        return\n",
    "    # display a n*n 2D manifold of digits\n",
    "    digit_size = 28\n",
    "    scale = 1.0\n",
    "    figure = np.zeros((digit_size * n, digit_size * n))\n",
    "    # linearly spaced coordinates corresponding to the 2D plot\n",
    "    # of digit classes in the latent space\n",
    "    grid_x = np.linspace(-scale, scale, n)\n",
    "    grid_y = np.linspace(-scale, scale, n)[::-1]\n",
    "\n",
    "    for i, yi in enumerate(grid_y):\n",
    "        for j, xi in enumerate(grid_x):\n",
    "            z_sample = np.array([[xi, yi]])\n",
    "            x_decoded = vae.decode(nnet.tensor(z_sample, device=device)).to(\"cpu\").detach().numpy()\n",
    "            digit = x_decoded[0].reshape(digit_size, digit_size)\n",
    "            figure[\n",
    "                i * digit_size : (i + 1) * digit_size,\n",
    "                j * digit_size : (j + 1) * digit_size,\n",
    "            ] = digit\n",
    "\n",
    "    plt.figure(figsize=(figsize, figsize))\n",
    "    start_range = digit_size // 2\n",
    "    end_range = n * digit_size + start_range\n",
    "    pixel_range = np.arange(start_range, end_range, digit_size)\n",
    "    sample_range_x = np.round(grid_x, 1)\n",
    "    sample_range_y = np.round(grid_y, 1)\n",
    "    plt.xticks(pixel_range, sample_range_x)\n",
    "    plt.yticks(pixel_range, sample_range_y)\n",
    "    plt.xlabel(\"z[0]\")\n",
    "    plt.ylabel(\"z[1]\")\n",
    "    plt.imshow(figure, cmap=\"Greys_r\")\n",
    "\n",
    "    plt.savefig(\"generated images/vae 2d latent space.jpeg\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "plot_latent_space_digits()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib widget\n",
    "\"\"\"Visualize latent space of labels only with latent_dim = 2\"\"\"\n",
    "\n",
    "\n",
    "def plot_label_clusters(data, labels):\n",
    "    if latent_size != 2:\n",
    "        print(\"Can`t plot 2d latent space for non-2d latent space\")\n",
    "        return\n",
    "    # display a 2D plot of the digit classes in the latent space\n",
    "    z_mean = (\n",
    "        vae.encode(nnet.tensor(data, device=device).reshape(-1, 28 * 28)).to(\"cpu\").detach().numpy()\n",
    "    )\n",
    "    plt.figure(figsize=(12, 10))\n",
    "    plt.scatter(z_mean[:, 0], z_mean[:, 1], c=labels)\n",
    "    plt.colorbar()\n",
    "    plt.xlabel(\"z[0]\")\n",
    "    plt.ylabel(\"z[1]\")\n",
    "\n",
    "    plt.savefig(\"generated images/vae 2d latent space labels.jpeg\")\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "plot_label_clusters(add_noise(training_data) if noisy_inputs else training_data, training_labels)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}