examples/gan.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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 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": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_size = (1, 28, 28)\n",
    "x_num, y_num = 5, 5\n",
    "samples_num = x_num * y_num\n",
    "margin = 15\n",
    "\n",
    "dataset, _, _, _ = load_mnist()\n",
    "dataset = dataset / 127.5 - 1  # normalization: / 255 => [0; 1]  #/ 127.5-1 => [-1; 1]\n",
    "\n",
    "noise_size = 100\n",
    "\n",
    "device = \"cuda\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "generator = nn.Sequential(\n",
    "    nn.Linear(noise_size, 256),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.BatchNorm1d(256),\n",
    "    nn.Linear(256, 512),\n",
    "    nn.Dropout(0.2),\n",
    "    nn.BatchNorm1d(512),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.Linear(512, 784),\n",
    "    nn.Tanh(),\n",
    ").to(device)\n",
    "\n",
    "discriminator = nn.Sequential(\n",
    "    nn.Linear(784, 128),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.Linear(128, 64),\n",
    "    nn.LeakyReLU(),\n",
    "    nn.Linear(64, 1),\n",
    "    nn.Sigmoid(),\n",
    ").to(device)\n",
    "\n",
    "loss_fn = nn.MSELoss()\n",
    "\n",
    "g_optimizer = Adam(generator.parameters(), lr=0.001, betas=(0.5, 0.999))\n",
    "d_optimizer = Adam(discriminator.parameters(), lr=0.001, betas=(0.5, 0.999))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch: 1/30, G loss: 0.4438888, D loss: 1.5700481: 100%|██████████| 600/600 [00:22<00:00, 26.90it/s]\n",
      "epoch: 2/30, G loss: 0.4623058, D loss: 1.4214900: 100%|██████████| 600/600 [00:19<00:00, 30.44it/s]\n",
      "epoch: 3/30, G loss: 0.3498764, D loss: 1.9643488: 100%|██████████| 600/600 [00:20<00:00, 29.10it/s]\n",
      "epoch: 4/30, G loss: 0.4362092, D loss: 1.4030291: 100%|██████████| 600/600 [00:20<00:00, 29.32it/s]\n",
      "epoch: 5/30, G loss: 0.4470197, D loss: 1.4730237: 100%|██████████| 600/600 [00:20<00:00, 28.65it/s]\n",
      "epoch: 6/30, G loss: 0.4218569, D loss: 1.4497340: 100%|██████████| 600/600 [00:20<00:00, 29.94it/s]\n",
      "epoch: 7/30, G loss: 0.4156743, D loss: 1.4897844: 100%|██████████| 600/600 [00:20<00:00, 29.47it/s]\n",
      "epoch: 8/30, G loss: 0.4309656, D loss: 1.4677904: 100%|██████████| 600/600 [00:19<00:00, 30.03it/s]\n",
      "epoch: 9/30, G loss: 0.4034271, D loss: 1.5358443: 100%|██████████| 600/600 [00:20<00:00, 29.98it/s]\n",
      "epoch: 10/30, G loss: 0.4116855, D loss: 1.4903698: 100%|██████████| 600/600 [00:20<00:00, 29.08it/s]\n",
      "epoch: 11/30, G loss: 0.4156543, D loss: 1.4935591: 100%|██████████| 600/600 [00:20<00:00, 29.20it/s]\n",
      "epoch: 12/30, G loss: 0.4110027, D loss: 1.4945952: 100%|██████████| 600/600 [00:20<00:00, 29.32it/s]\n",
      "epoch: 13/30, G loss: 0.4147280, D loss: 1.4892874: 100%|██████████| 600/600 [00:20<00:00, 29.67it/s]\n",
      "epoch: 14/30, G loss: 0.4209598, D loss: 1.4880155: 100%|██████████| 600/600 [00:20<00:00, 28.67it/s]\n",
      "epoch: 15/30, G loss: 0.4041674, D loss: 1.4957856: 100%|██████████| 600/600 [00:20<00:00, 29.56it/s]\n",
      "epoch: 16/30, G loss: 0.4122545, D loss: 1.4912832: 100%|██████████| 600/600 [00:20<00:00, 29.00it/s]\n",
      "epoch: 17/30, G loss: 0.4404696, D loss: 1.5473717: 100%|██████████| 600/600 [00:22<00:00, 26.88it/s]\n",
      "epoch: 18/30, G loss: 0.4270725, D loss: 1.4814024: 100%|██████████| 600/600 [00:22<00:00, 26.68it/s]\n",
      "epoch: 19/30, G loss: 0.4266073, D loss: 1.4784986: 100%|██████████| 600/600 [00:18<00:00, 31.60it/s]\n",
      "epoch: 20/30, G loss: 0.4196189, D loss: 1.4885912:  90%|█████████ | 541/600 [00:19<00:01, 31.55it/s]"
     ]
    }
   ],
   "source": [
    "batch_size = 100\n",
    "epochs = 30\n",
    "\n",
    "each_epoch_generated_samples = []\n",
    "const_noise = nnet.tensor(\n",
    "    np.random.normal(0, 1, (samples_num, noise_size)),\n",
    "    device=device,\n",
    ")\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    tqdm_range = tqdm(range(0, len(dataset), batch_size), desc=f\"epoch {epoch}\")\n",
    "    generator.train()\n",
    "    discriminator.train()\n",
    "    for i in tqdm_range:\n",
    "        batch = dataset[i : i + batch_size]\n",
    "        batch = nnet.tensor(batch, device=device)\n",
    "\n",
    "        d_optimizer.zero_grad()\n",
    "\n",
    "        # train discriminator on real data\n",
    "        real_data = batch\n",
    "        real_data = real_data.reshape(real_data.shape[0], -1)\n",
    "\n",
    "        real_data_prediction = discriminator(real_data)\n",
    "        real_data_loss = loss_fn(\n",
    "            real_data_prediction,\n",
    "            nnet.tensor(\n",
    "                np.ones((real_data_prediction.shape[0], 1)),\n",
    "                device=device,\n",
    "            ),\n",
    "        )\n",
    "        real_data_loss.backward()\n",
    "        d_optimizer.step()\n",
    "\n",
    "        # train discriminator on fake data\n",
    "        noise = nnet.tensor(\n",
    "            np.random.normal(0, 1, (batch_size, noise_size)),\n",
    "            device=device,\n",
    "        )\n",
    "        fake_data = generator(noise)\n",
    "        fake_data_prediction = discriminator(fake_data)\n",
    "        fake_data_loss = loss_fn(\n",
    "            fake_data_prediction,\n",
    "            nnet.tensor(\n",
    "                np.zeros((fake_data_prediction.shape[0], 1)),\n",
    "                device=device,\n",
    "            ),\n",
    "        )\n",
    "        fake_data_loss.backward()\n",
    "        d_optimizer.step()\n",
    "\n",
    "        g_optimizer.zero_grad()\n",
    "\n",
    "        noise = nnet.tensor(\n",
    "            np.random.normal(0, 1, (batch_size, noise_size)),\n",
    "            device=device,\n",
    "        )\n",
    "        fake_data = generator(noise)\n",
    "        fake_data_prediction = discriminator(fake_data)\n",
    "        fake_data_loss = loss_fn(\n",
    "            fake_data_prediction,\n",
    "            nnet.tensor(\n",
    "                np.ones((fake_data_prediction.shape[0], 1)),\n",
    "                device=device,\n",
    "            ),\n",
    "        )\n",
    "        fake_data_loss.backward()\n",
    "        g_optimizer.step()\n",
    "\n",
    "        g_loss = -np.log(fake_data_prediction.detach().cpu().numpy()).mean()\n",
    "        d_loss = (\n",
    "            -np.log(real_data_prediction.detach().cpu().numpy()).mean()\n",
    "            - np.log(1 - fake_data_prediction.detach().cpu().numpy()).mean()\n",
    "        )\n",
    "        tqdm_range.set_description(\n",
    "            f\"epoch: {epoch + 1}/{epochs}, G loss: {g_loss:.7f}, D loss: {d_loss:.7f}\"\n",
    "        )\n",
    "\n",
    "    if const_noise == None:\n",
    "        noise = nnet.tensor(\n",
    "            np.random.normal(0, 1, (batch_size, noise_size)),\n",
    "            device=device,\n",
    "        )\n",
    "    else:\n",
    "        noise = const_noise\n",
    "\n",
    "    each_epoch_generated_samples.append(generator(noise).detach().cpu().numpy().reshape(-1, 28, 28))\n",
    "\n",
    "    generator.eval()\n",
    "    discriminator.eval()\n",
    "    for i in range(samples_num):\n",
    "        image = each_epoch_generated_samples[-1][i] * 127.5 + 127.5\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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "generator.eval()\n",
    "discriminator.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "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\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_vectors_interpolation():\n",
    "    \"\"\"Create vectors create interpolation  in the latent space between two sets of noise vectors\"\"\"\n",
    "    steps = 10\n",
    "    interval = 15\n",
    "    images = []\n",
    "\n",
    "    noise_1 = nnet.tensor(np.random.normal(0, 1, (samples_num, noise_size)))\n",
    "\n",
    "    for _ in range(steps):\n",
    "        noise_2 = nnet.tensor(np.random.normal(0, 1, (samples_num, noise_size)))\n",
    "\n",
    "        noise_interp = np.linspace(noise_1.numpy(), noise_2.numpy(), interval)\n",
    "\n",
    "        noise_1 = noise_2\n",
    "\n",
    "        for vectors in noise_interp:\n",
    "            generated_images = (\n",
    "                generator(nnet.tensor(vectors, device=device))\n",
    "                .reshape(-1, 28, 28)\n",
    "                .to(\"cpu\")\n",
    "                .detach()\n",
    "                .numpy()\n",
    "                * 127.5\n",
    "                + 127.5\n",
    "            )\n",
    "\n",
    "            images.append(get_images_set(generated_images).convert(\"L\").convert(\"P\"))\n",
    "\n",
    "    images[0].save(\n",
    "        \"generated images/gan_vectors_interpolation.gif\",\n",
    "        save_all=True,\n",
    "        append_images=images[1:],\n",
    "        optimize=False,\n",
    "        duration=100,\n",
    "        loop=0,\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "create_vectors_interpolation()\n",
    "\n",
    "\n",
    "each_epoch_generated_images = []\n",
    "for epoch_samples in each_epoch_generated_samples:\n",
    "    each_epoch_generated_images.append(\n",
    "        get_images_set(epoch_samples * 127.5 + 127.5).convert(\"L\").convert(\"P\")\n",
    "    )\n",
    "\n",
    "each_epoch_generated_images[0].save(\n",
    "    \"generated images/gan_training_process.gif\",\n",
    "    save_all=True,\n",
    "    append_images=each_epoch_generated_images[1:],\n",
    "    optimize=False,\n",
    "    duration=150,\n",
    "    loop=0,\n",
    ")\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
}