PokeGAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PokeGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute this to save new versions of the notebook\n",
    "import jovian\n",
    "jovian.commit(filename=\"PokeGAN.ipynb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Github Repository"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I have a full GitHub repository available for this project. I found my workflow to be quite a bit smoother using the run.py script available in the repository, but this notebook offers the same functionality if you prefer. \n",
    "\n",
    "To clone the repository to your machine use the following bash command: \n",
    "\n",
    "`$ git clone https://github.com/Kodlak15/PokeGAN`\n",
    "\n",
    "Then navigate to the local repository in your terminal, and enter the following command:\n",
    "\n",
    "`$ python3 run.py`\n",
    "\n",
    "You will be prompted for the number of epochs you would like to train for as well as the learning rate. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Device Management"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_default_device():\n",
    "    \"\"\" Pick GPU if available, else CPU \"\"\"\n",
    "    if torch.cuda.is_available():\n",
    "        return torch.device('cuda')\n",
    "    \n",
    "    else:\n",
    "        return torch.device('cpu')\n",
    "\n",
    "def to_device(data, device):\n",
    "    \"\"\" Move tensor(s) to chosen device \"\"\"\n",
    "    if isinstance(data, (list,tuple)):\n",
    "        return [to_device(x, device) for x in data]\n",
    "    \n",
    "    if isinstance(data, dict):\n",
    "        return {k: to_device(t, device) for k, t in data.items()}\n",
    "    \n",
    "    return data.to(device, non_blocking=True)\n",
    "\n",
    "class DeviceDataLoader():\n",
    "    \"\"\" Wrap a dataloader to move data to a device \"\"\"\n",
    "    def __init__(self, dl, device):\n",
    "        self.dl = dl\n",
    "        self.device = device\n",
    "        \n",
    "    def __iter__(self):\n",
    "        \"\"\" Yield a batch of data after moving it to device \"\"\"\n",
    "        for b in self.dl: \n",
    "            yield to_device(b, self.device)\n",
    "\n",
    "    def __len__(self):\n",
    "        \"\"\" Number of batches \"\"\"\n",
    "        return len(self.dl)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Utilities "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Tuple, Union\n",
    "from torch import Tensor\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.utils import make_grid, save_image\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image\n",
    "import os\n",
    "from os.path import join\n",
    "import cv2\n",
    "import json\n",
    "import numpy as np\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_path(path: str):\n",
    "    return path.replace(\"\\\\\", '/')\n",
    "\n",
    "def prepare_paths():\n",
    "    root_dir = os.getcwd()\n",
    "    img_dir = format_path(join(root_dir, \"images\"))\n",
    "    train_dir = format_path(join(img_dir, \"train-images\"))\n",
    "    fake_dir = join(root_dir, \"fakes\")\n",
    "    weights_dir = join(root_dir, \"weights\")\n",
    "    history_dir = join(root_dir, \"history\")\n",
    "\n",
    "    with open(\"paths.txt\", 'w') as f:\n",
    "        f.write(root_dir + '\\n')\n",
    "        f.write(img_dir + '\\n')\n",
    "        f.write(train_dir + '\\n')\n",
    "        f.write(fake_dir + '\\n')\n",
    "        f.write(weights_dir + '\\n')\n",
    "        f.write(history_dir)\n",
    "\n",
    "    if not \"images\" in os.listdir(root_dir):\n",
    "        os.mkdir(img_dir)\n",
    "\n",
    "    if not \"train-images\" in os.listdir(img_dir):\n",
    "        os.mkdir(join(img_dir, \"train-images\"))\n",
    "\n",
    "    if not \"fakes\" in os.listdir(root_dir):\n",
    "        os.mkdir(fake_dir)\n",
    "\n",
    "    if not \"weights\" in os.listdir(root_dir):\n",
    "        os.mkdir(weights_dir)\n",
    "\n",
    "    if not \"history\" in os.listdir(root_dir):\n",
    "        os.mkdir(history_dir)\n",
    "        with open(join(history_dir, \"history.json\"), 'w') as f:\n",
    "            json.dump({\"history\": []}, f)\n",
    "\n",
    "def get_paths():\n",
    "    prepare_paths()\n",
    "    with open(\"paths.txt\", 'r') as f:\n",
    "        return [path.replace('\\n', '') for path in f.readlines()]\n",
    "\n",
    "root_dir, img_dir, train_dir, fake_dir, weights_dir, history_dir = get_paths()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_stats = [0.1874, 0.1779, 0.1681], [1.0, 1.0, 1.0]\n",
    "\n",
    "@torch.no_grad()\n",
    "def show_images(batch: Union[DataLoader, Tensor]):\n",
    "    \"\"\"\n",
    "    Takes a tensor (B, C, W, H) or dataloader as input and displays a batch of training images\n",
    "    \"\"\"\n",
    "    for images in batch:\n",
    "        fig, ax = plt.subplots(figsize=(32,32))\n",
    "        ax.set_xticks([]); ax.set_yticks([])\n",
    "        images = images.to(\"cpu\")\n",
    "        images = denormalize(images, *train_stats)\n",
    "        ax.imshow(make_grid(images, nrow=8).permute(1, 2, 0))\n",
    "        break\n",
    "\n",
    "@torch.no_grad()\n",
    "def show_fakes(images, num_to_show=64):\n",
    "    \"\"\"\n",
    "    Displays a collection of fake images\n",
    "    \"\"\"\n",
    "    fig, ax = plt.subplots(figsize=(32,32))\n",
    "    ax.set_xticks([]); ax.set_yticks([])\n",
    "    images = images.to(\"cpu\")\n",
    "    images = denormalize(images[:num_to_show], *train_stats)\n",
    "    ax.imshow(make_grid(images, nrow=8).permute(1, 2, 0))\n",
    "\n",
    "def transform_image(img: Image):\n",
    "    \"\"\"\n",
    "    Intended to transform images from CMYK -> RGB\n",
    "    Overlays image on a plain, black background to remove transparent pixels\n",
    "    \"\"\"\n",
    "    new_img = Image.new(\"RGBA\", img.size, \"BLACK\")\n",
    "    new_img.paste(img, (0, 0), img)\n",
    "    new_img = new_img.convert(\"RGB\")\n",
    "    return new_img\n",
    "\n",
    "def save_samples(G: nn.Module, index: int, x: Tensor):\n",
    "    fake_images = denormalize(to_device(G(x), device=\"cpu\"), *train_stats)\n",
    "    filename = join(fake_dir, \"generated-images-{0:0=4d}.png\".format(index))\n",
    "    print(f\"Saving {filename}\")\n",
    "    save_image(fake_images[:64], filename, nrow=8)\n",
    "\n",
    "def make_video(fps=30):\n",
    "    vid_fname = \"pokeGAN.avi\"\n",
    "\n",
    "    files = [join(fake_dir, f) for f in os.listdir(fake_dir) if 'generated' in f]\n",
    "    files.sort()\n",
    "\n",
    "    out = cv2.VideoWriter(vid_fname, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), fps, (1042, 1042))\n",
    "    [out.write(cv2.imread(fname)) for fname in files]\n",
    "    out.release()\n",
    "\n",
    "def denormalize(images, means, stds):\n",
    "    means = torch.tensor(means).reshape(1, 3, 1, 1)\n",
    "    stds = torch.tensor(stds).reshape(1, 3, 1, 1)\n",
    "    return images * stds + means"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def parse_history():\n",
    "    with open(join(os.getcwd(), \"history\", \"history.json\"), 'r') as f:\n",
    "        history = json.load(f)[\"history\"]\n",
    "\n",
    "    epochs = np.arange(len(history))\n",
    "    losses_d = []\n",
    "    losses_g = []\n",
    "    real_scores = []\n",
    "    fake_scores = []\n",
    "\n",
    "    for epoch in history:\n",
    "        loss_g, loss_d, real_score, fake_score = re.findall(r\"[-+]?(?:\\d*\\.\\d+|\\d+)\", epoch)\n",
    "        losses_g.append(float(loss_g))\n",
    "        losses_d.append(float(loss_d))\n",
    "        real_scores.append(float(real_score))\n",
    "        fake_scores.append(float(fake_score))\n",
    "\n",
    "    return epochs, losses_d, losses_g, real_scores, fake_scores\n",
    "\n",
    "def plot_results():\n",
    "    epochs, losses_d, losses_g, real_scores, fake_scores = parse_history()\n",
    "\n",
    "    plt.style.use(\"seaborn\")\n",
    "    fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(16, 4))\n",
    "    ax[0][0].title.set_text(\"Losses\")\n",
    "    ax[0][1].title.set_text(\"Scores\")\n",
    "    # Plot losses\n",
    "    ax[0][0].plot(epochs, losses_d)\n",
    "    ax[0][0].plot(epochs, np.poly1d(np.polyfit(epochs, losses_d, deg=2))(epochs), color=\"red\", linestyle=\"--\")\n",
    "    ax[0][0].set_ylim([-0.2, 5])\n",
    "    ax[0][0].set_ylabel(\"Discriminator\")\n",
    "    ax[1][0].plot(epochs, losses_g)\n",
    "    ax[1][0].plot(epochs, np.poly1d(np.polyfit(epochs, losses_g, deg=2))(epochs), color=\"red\", linestyle=\"--\")\n",
    "    ax[1][0].set_ylim([-0.2, 5])\n",
    "    ax[1][0].set_ylabel(\"Generator\")\n",
    "    # Plot scores\n",
    "    ax[0][1].plot(epochs, real_scores)\n",
    "    ax[0][1].plot(epochs, np.poly1d(np.polyfit(epochs, real_scores, deg=2))(epochs), color=\"red\", linestyle=\"--\")\n",
    "    ax[0][1].set_ylabel(\"Real\")\n",
    "    ax[1][1].plot(epochs, fake_scores)\n",
    "    ax[1][1].plot(epochs, np.poly1d(np.polyfit(epochs, fake_scores, deg=2))(epochs), color=\"red\", linestyle=\"--\")\n",
    "    ax[1][1].set_ylabel(\"Fake\")\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "from bs4 import BeautifulSoup\n",
    "from urllib.parse import urljoin\n",
    "import re\n",
    "from PIL import Image\n",
    "import os\n",
    "from os.path import join\n",
    "from tqdm import tqdm\n",
    "import random\n",
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_image_urls(refresh=False, seed=15):\n",
    "    url = \"https://www.pokemon.com/us/pokedex/\"\n",
    "    r = requests.get(url)\n",
    "    soup = BeautifulSoup(r.content, \"html.parser\")\n",
    "    body = soup.find(\"body\")\n",
    "    pattern = \"\\/us\\/pokedex\\/[a-zA-z]+\"\n",
    "    pokemon = body.find_all('a', {\"href\": re.compile(pattern)})\n",
    "    \n",
    "    if refresh or \"img-urls.json\" not in os.listdir(img_dir):\n",
    "        img_urls = []\n",
    "        for p in tqdm(pokemon):\n",
    "            p_url = urljoin(url, p[\"href\"])\n",
    "            r = requests.get(p_url)\n",
    "            soup = BeautifulSoup(r.content, \"html.parser\")\n",
    "            body = soup.find(\"body\")\n",
    "            img_url = body.find(\"img\", {\"class\": \"active\"})[\"src\"]\n",
    "            img_urls.append(img_url)\n",
    "\n",
    "        print(f\"{len(img_urls)} Pokemon images found.\")\n",
    "        random.seed(seed)\n",
    "        random.shuffle(img_urls)\n",
    "\n",
    "        filename = join(img_dir, \"img-urls.json\")\n",
    "        with open(filename, 'w') as f:\n",
    "            json.dump(img_urls, f)\n",
    "\n",
    "    filename = join(img_dir, \"img-urls.json\")\n",
    "    with open(filename, 'r') as f:\n",
    "        img_urls = json.load(f)\n",
    "\n",
    "    return img_urls\n",
    "\n",
    "def get_images(refresh=False, seed=15):\n",
    "    if not \"img-urls.json\" in os.listdir(img_dir):\n",
    "        img_urls = get_image_urls(refresh=True, seed=seed)\n",
    "\n",
    "    else:\n",
    "        img_urls = get_image_urls(seed=seed)\n",
    "\n",
    "    if len(os.listdir(train_dir)) == 0:\n",
    "        for url in tqdm(img_urls):\n",
    "            \"Creating training set...\"\n",
    "            pID = url.split('/')[-1]\n",
    "            r = requests.get(url)\n",
    "            if r.status_code == 200:\n",
    "                r.raw.decode_content = True\n",
    "\n",
    "                if pID not in os.listdir(join(img_dir, \"train-images\")):\n",
    "                    filename = join(img_dir, \"train-images\", pID)\n",
    "                    with open(filename, 'wb') as f:\n",
    "                        f.write(r.content)\n",
    "\n",
    "            else:\n",
    "                print(f\"Image {pID} could not be retrieved\")\n",
    "\n",
    "    print(\"Finished!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "refresh = False\n",
    "seed = 15 \n",
    "get_images()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PokemonDataset(Dataset):\n",
    "    \"\"\" Pokemon images dataset \"\"\"\n",
    "    def __init__(self, directory: str, transform=None):\n",
    "        self.directory = directory\n",
    "        self.transform = transform\n",
    "        self.image_paths = [format_path(join(directory, img_name)) for img_name in sorted(os.listdir(directory))]\n",
    "\n",
    "    def __len__(self):\n",
    "        return len([f for f in os.listdir(self.directory) if \".png\" in f])\n",
    "\n",
    "    def __getitem__(self, idx: int):\n",
    "        path = self.image_paths[idx]\n",
    "        with Image.open(path) as img:\n",
    "\n",
    "            if self.transform:\n",
    "                img = self.transform(img)\n",
    "            \n",
    "            return img"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        latent_size = 128\n",
    "\n",
    "        self.net = nn.Sequential(\n",
    "            # in: 3 x 128 x 128\n",
    "\n",
    "            nn.Conv2d(3, 128, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # out: 128 x 64 x 64\n",
    "\n",
    "            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # out: 256 x 32 x 32\n",
    "\n",
    "            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # out: 512 x 16 x 16\n",
    "\n",
    "            nn.Conv2d(512, 1024, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # out: 1024 x 8 x 8\n",
    "\n",
    "            nn.Conv2d(1024, 2048, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(2048),\n",
    "            nn.LeakyReLU(0.2, inplace=True),\n",
    "            # out: 2048 x 4 x 4\n",
    "\n",
    "            nn.Conv2d(2048, 1, kernel_size=4, stride=1, padding=0, bias=False),\n",
    "            # out: 1 x 1 x 1\n",
    "\n",
    "            nn.Flatten(),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "        if \"Discriminator.pth\" in os.listdir(weights_dir):\n",
    "            self.load()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.net(x)\n",
    "        return out\n",
    "\n",
    "    def save(self):\n",
    "        torch.save(self.state_dict(), join(weights_dir, \"Discriminator.pth\"))\n",
    "\n",
    "    def load(self):\n",
    "        assert \"Discriminator.pth\" in os.listdir(weights_dir), \"No discriminator weights found\"\n",
    "        self.load_state_dict(torch.load(join(weights_dir, \"Discriminator.pth\")))\n",
    "        print(\"Discriminator weights loaded successfully!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        latent_size = 128\n",
    "\n",
    "        self.net = nn.Sequential(\n",
    "            # in: latent_size x 1 x 1\n",
    "\n",
    "            nn.ConvTranspose2d(latent_size, 2048, kernel_size=4, stride=1, padding=0, bias=False),\n",
    "            nn.BatchNorm2d(2048),\n",
    "            nn.ReLU(True),\n",
    "            # out: 2048 x 4 x 4\n",
    "\n",
    "            nn.ConvTranspose2d(2048, 1024, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(1024),\n",
    "            nn.ReLU(True),\n",
    "            # out: 1024 x 8 x 8\n",
    "\n",
    "            nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(512),\n",
    "            nn.ReLU(True),\n",
    "            # out: 512 x 16 x 16\n",
    "\n",
    "            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(256),\n",
    "            nn.ReLU(True),\n",
    "            # out: 256 x 32 x 32\n",
    "\n",
    "            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.BatchNorm2d(128),\n",
    "            nn.ReLU(True),\n",
    "            # out: 128 x 64 x 64\n",
    "\n",
    "            nn.ConvTranspose2d(128, 3, kernel_size=4, stride=2, padding=1, bias=False),\n",
    "            nn.Tanh()\n",
    "            # out: 3 x 128 x 128\n",
    "        )\n",
    "\n",
    "        if \"Generator.pth\" in os.listdir(weights_dir):\n",
    "            self.load()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.net(x)\n",
    "        return out\n",
    "\n",
    "    def save(self):\n",
    "        torch.save(self.state_dict(), join(weights_dir, \"Generator.pth\"))\n",
    "\n",
    "    def load(self):\n",
    "        assert \"Generator.pth\" in os.listdir(weights_dir), \"No generator weights found\"\n",
    "        self.load_state_dict(torch.load(join(weights_dir, \"Generator.pth\")))\n",
    "        print(\"Generator weights loaded successfully!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_size = 128\n",
    "batch_size = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_discriminator(D: nn.Module, G: nn.Module, images: Tensor, opt_d: torch.nn.functional, device: torch.device):\n",
    "    opt_d.zero_grad()\n",
    "\n",
    "    real_preds = D(images)\n",
    "    real_targets = torch.ones(images.size(0), 1, device=device)\n",
    "    real_noisy_targets = (0.7 - 1.2) * torch.rand(images.size(0), 1, device=device) + 1.2\n",
    "    real_loss = F.binary_cross_entropy(real_preds, real_noisy_targets)\n",
    "    real_score = torch.mean(real_preds).item()\n",
    "\n",
    "    x = torch.randn(images.size(0), latent_size, 1, 1, device=device)\n",
    "    fake_images = G(x)\n",
    "\n",
    "    fake_preds = D(fake_images)\n",
    "    fake_targets = torch.zeros(fake_images.size(0), 1, device=device)\n",
    "    fake_noisy_targets = (0.0 - 0.3) * torch.rand(fake_images.size(0), 1, device=device) + 0.3\n",
    "    fake_loss = F.binary_cross_entropy(fake_preds, fake_noisy_targets)\n",
    "    fake_score = torch.mean(fake_preds).item()\n",
    "\n",
    "    loss = real_loss + fake_loss\n",
    "    loss.backward()\n",
    "    opt_d.step()\n",
    "\n",
    "    return loss.item(), real_score, fake_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_generator(D: nn.Module, G: nn.Module, batch_size: int, opt_g: torch.nn.functional, device: torch.device):\n",
    "    opt_g.zero_grad()\n",
    "\n",
    "    x = torch.randn(batch_size, latent_size, 1, 1, device=device)\n",
    "    fake_images = G(x)\n",
    "\n",
    "    preds = D(fake_images)\n",
    "    targets = torch.ones(batch_size, 1, device=device)\n",
    "    loss = F.binary_cross_entropy(preds, targets)\n",
    "\n",
    "    loss.backward()\n",
    "    opt_g.step()\n",
    "    \n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def fit(D: nn.Module, G: nn.Module, train_dl: DataLoader, epochs: int, lr: float, device: torch.device, start_idx=1):\n",
    "    # Create backups\n",
    "    shutil.copy(join(history_dir, \"history.json\"), join(history_dir, \"history_backup.json\"))\n",
    "\n",
    "    if \"Discriminator.pth\" in os.listdir(weights_dir):\n",
    "        shutil.copy(join(weights_dir, \"Discriminator.pth\"), join(weights_dir, \"Discriminator-backup.pth\"))\n",
    "\n",
    "    if \"Generator.pth\" in os.listdir(weights_dir):\n",
    "        shutil.copy(join(weights_dir, \"Generator.pth\"), join(weights_dir, \"Generator-backup.pth\"))\n",
    "\n",
    "    with open(join(history_dir, \"history.json\"), 'r') as f:\n",
    "        history = json.load(f)[\"history\"]\n",
    "\n",
    "    losses_d = []\n",
    "    losses_g = []\n",
    "    real_scores = []\n",
    "    fake_scores = []\n",
    "\n",
    "    opt_d = torch.optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "    opt_g = torch.optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        for real_images in tqdm(train_dl):\n",
    "            loss_d, real_score, fake_score = train_discriminator(D, G, real_images, opt_d, device)\n",
    "            loss_g = train_generator(D, G, real_images.size(0), opt_g, device)\n",
    "\n",
    "        losses_d.append(loss_d)\n",
    "        losses_g.append(loss_g)\n",
    "        real_scores.append(real_score)\n",
    "        fake_scores.append(fake_score)\n",
    "\n",
    "        epoch_results = \"Epoch [{}/{}], loss_g: {:.4f}, loss_d: {:.4f}, real_score: {:.4f}, fake_score: {:.4f}\".format(\n",
    "                        epoch+1, epochs, loss_g, loss_d, real_score, fake_score)\n",
    "        \n",
    "        print(epoch_results)\n",
    "        history.append(epoch_results[epoch_results.find(',')+2:])\n",
    "\n",
    "        x = torch.randn(batch_size, latent_size, 1, 1, device=device)\n",
    "        save_samples(G, epoch+start_idx, x)\n",
    "        D.save()\n",
    "        G.save()\n",
    "\n",
    "    print(\"Saving results...\")\n",
    "    with open(join(history_dir, \"history.json\"), 'w') as f:\n",
    "        json.dump({\"history\": history}, f)\n",
    "\n",
    "    make_video()\n",
    "\n",
    "    return history"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.nn.functional as F\n",
    "from random import shuffle\n",
    "import torchvision.transforms as T\n",
    "from torch.utils.data import DataLoader\n",
    "import shutil\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run():\n",
    "    try:\n",
    "        epochs = int(input(\"Enter the number of epochs to train for (1-5000): \"))\n",
    "        assert epochs in range(1, 5001), \"Enter a number between 1 and 5000\"\n",
    "        lr = float(input(\"Enter the learning rate (5e-6 - 5e-4): \"))\n",
    "        assert lr >= 5e-6 and lr <= 5e-4, \"Enter a number between 5e-6 and 5e-4\"\n",
    "\n",
    "    except AssertionError:\n",
    "        print(\"Restarting program...\")\n",
    "        run()\n",
    "\n",
    "    device = get_default_device()\n",
    "    train_ds = PokemonDataset(train_dir, transform=train_transform)\n",
    "    train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True)\n",
    "    train_dl = DeviceDataLoader(train_dl, device)\n",
    "    D = to_device(Discriminator(), device)\n",
    "    G = to_device(Generator(), device)\n",
    "\n",
    "    return fit(D, G, train_dl, epochs, lr, device, start_idx=len(os.listdir(fake_dir))+1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "img_size = 128\n",
    "batch_size = 256\n",
    "train_stats = [0.1874, 0.1779, 0.1681], [1.0, 1.0, 1.0]\n",
    "\n",
    "train_transform = T.Compose([\n",
    "    T.Lambda(lambda img: transform_image(img)),\n",
    "    T.Resize(img_size),\n",
    "    T.CenterCrop(img_size),\n",
    "    T.RandomHorizontalFlip(0.2),\n",
    "    T.RandomRotation(3, fill=0),\n",
    "    T.ToTensor(),\n",
    "    T.Normalize(*train_stats)\n",
    "])\n",
    "\n",
    "run()\n",
    "input(\"Training finished, press enter to exit...\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Visualize results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_results();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = Generator()\n",
    "\n",
    "x = torch.randn(batch_size, latent_size, 1, 1)\n",
    "fake_images = G(x)\n",
    "num_to_show = 64\n",
    "show_fakes(fake_images, num_to_show=num_to_show);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# My results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an example batch of Pokemon this model created after training for 10,000 epochs on my PC. The results are not perfect, but there are definitely some interesting samples!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=example.png>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.2 ('Jovian')",
   "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.2"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "63e5a82fdc887cbdaa189a2800f151fe693184309a9dd3cdcdd49c9afd53da1a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}