t81_558_class_07_2_train_gan.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4mB_sFildiDh"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jeffheaton/t81_558_deep_learning/blob/master/t81_558_class_07_2_train_gan.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "303iHntmdiDj"
      },
      "source": [
        "# T81-558: Applications of Deep Neural Networks\n",
        "**Module 7: Generative Adversarial Networks**\n",
        "* Instructor: [Jeff Heaton](https://sites.wustl.edu/jeffheaton/), McKelvey School of Engineering, [Washington University in St. Louis](https://engineering.wustl.edu/Programs/Pages/default.aspx)\n",
        "* For more information visit the [class website](https://sites.wustl.edu/jeffheaton/t81-558/)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d5n2iv9udiDk"
      },
      "source": [
        "# Module 7 Material\n",
        "\n",
        "* Part 7.1: Introduction to GANs for Image and Data Generation [[Video]](https://www.youtube.com/watch?v=hZw-AjbdN5k&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/t81_558_deep_learning/blob/master/t81_558_class_07_1_gan_intro.ipynb)\n",
        "* **Part 7.2: Train StyleGAN3 with your Own Images** [[Video]](https://www.youtube.com/watch?v=R546LYsQk5M&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/t81_558_deep_learning/blob/master/t81_558_class_07_2_train_gan.ipynb)\n",
        "* Part 7.3: Exploring the StyleGAN Latent Vector [[Video]](https://www.youtube.com/watch?v=goQzp8QSb2s&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/t81_558_deep_learning/blob/master/t81_558_class_07_3_latent_vector.ipynb)\n",
        "* Part 7.4: GANs to Enhance Old Photographs Deoldify [[Video]](https://www.youtube.com/watch?v=0OTd5GlHRx4&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/t81_558_deep_learning/blob/master/t81_558_class_07_4_deoldify.ipynb)\n",
        "* Part 7.5: GANs for Tabular Synthetic Data Generation [[Video]](https://www.youtube.com/watch?v=yujdA46HKwA&list=PLjy4p-07OYzulelvJ5KVaT2pDlxivl_BN) [[Notebook]](https://github.com/jeffheaton/t81_558_deep_learning/blob/master/t81_558_class_07_5_tabular_synthetic.ipynb)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Ohl-hxAAaNp"
      },
      "source": [
        "# Google CoLab Instructions\n",
        "\n",
        "The following code ensures that Google CoLab is running the correct version of TensorFlow."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y8_-1h5ddiDp"
      },
      "outputs": [],
      "source": [
        "\n",
        "try:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive', force_remount=True)\n",
        "    COLAB = True\n",
        "    print(\"Note: using Google CoLab\")\n",
        "    %tensorflow_version 2.x\n",
        "except:\n",
        "    print(\"Note: not using Google CoLab\")\n",
        "    COLAB = False"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1yqlUD4sdiDk"
      },
      "source": [
        "# Part 7.2: Train StyleGAN3 with your Images\n",
        "\n",
        "Training GANs with StyleGAN is resource-intensive. The NVIDA StyleGAN researchers used computers with eight high-end GPUs for the high-resolution face GANs trained by NVIDIA. The GPU used by NVIDIA is an A100, which has more memory and cores than the P100 or V100 offered by even Colab Pro+. In this part, we will use StyleGAN2 to train rather than StyleGAN3. You can use networks trained with StyleGAN2 from StyleGAN3; however, StyleGAN3 usually is more effective at training than StyleGAN2.\n",
        "\n",
        "Unfortunately, StyleGAN3 is compute-intensive and will perform slowly on any GPU that is not the latest Ampere technology. Because Colab does not provide such technology, I am keeping the training guide at the StyleGAN2 level. Switching to StyleGAN3 is relatively easy, as will be pointed out later. \n",
        "\n",
        "Make sure that you are running this notebook with a GPU runtime. You can train GANs with either Google Colab Free or Pro. I recommend at least the Pro version due to better GPU instances, longer runtimes, and timeouts. Additionally, the capability of Google Colab Pro to run in the background is valuable when training GANs, as you can close your browser or reboot your laptop while training continues.\n",
        "\n",
        "\n",
        "\n",
        "You will store your training data and trained neural networks to GDRIVE. For GANs, I lay out my GDRIVE like this:\n",
        "\n",
        "* ./data/gan/images - RAW images I wish to train on.\n",
        "* ./data/gan/datasets - Actual training datasets that I convert from the raw images.\n",
        "* ./data/gan/experiments - The output from StyleGAN2, my image previews, and saved network snapshots.\n",
        "\n",
        "You will mount the drive at the following location.\n",
        "\n",
        "```\n",
        "/content/drive/MyDrive/data\n",
        "```\n",
        "\n",
        "## What Sort of GPU do you Have?\n",
        "\n",
        "The type of GPU assigned to you by Colab will significantly affect your training time. Some sample times that I achieved with Colab are given here. I've found that Colab Pro generally starts you with a V100, however, if you run scripts non-stop for 24hrs straight for a few days in a row, you will generally be throttled back to a P100.\n",
        "\n",
        "* 1024x1024 - V100 - 566 sec/tick (CoLab Pro)\n",
        "* 1024x1024 - P100 - 1819 sec/tick (CoLab Pro)\n",
        "* 1024x1024 - T4 - 2188 sec/tick (CoLab Free)\n",
        "\n",
        "By comparison, a 1024x1024 GAN trained with StyleGAN3 on a V100 is 3087 sec/tick.\n",
        "\n",
        "If you use Google CoLab Pro, generally, it will not disconnect before 24 hours, even if you (but not your script) are inactive. Free CoLab WILL disconnect a perfectly good running script if you do not interact for a few hours. The following describes how to circumvent this issue.\n",
        "\n",
        "* [How to prevent Google Colab from disconnecting?](https://stackoverflow.com/questions/57113226/how-to-prevent-google-colab-from-disconnecting)\n",
        "\n",
        "## Set Up New Environment\n",
        "\n",
        "You will likely need to train for >24 hours. Colab will disconnect you. You must be prepared to restart training when this eventually happens. Training is divided into ticks, every so many ticks (50 by default), your neural network is evaluated, and a snapshot is saved. When CoLab shuts down, all training after the last snapshot is lost. It might seem desirable to snapshot after each tick; however, this snapshotting process itself takes nearly an hour. Learning an optimal snapshot size for your resolution and training data is important.\n",
        "\n",
        "We will mount GDRIVE so that you will save your snapshots there. You must also place your training images in GDRIVE.\n",
        "\n",
        "You must also install NVIDIA StyleGAN2 ADA PyTorch. We also need to downgrade PyTorch to a version that supports StyleGAN."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9LVr5sV8CKz0"
      },
      "outputs": [],
      "source": [
        "!pip uninstall jax jaxlib -y\n",
        "!pip install \"jax[cuda11_cudnn805]==0.3.10\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
        "!pip install torch==1.8.1 torchvision==0.9.1\n",
        "!git clone https://github.com/NVlabs/stylegan2-ada-pytorch.git\n",
        "!pip install ninja"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oBwmbfKJCORs"
      },
      "source": [
        "## Find Your Files\n",
        "\n",
        "The drive is mounted to the following location.\n",
        "\n",
        "```\n",
        "/content/drive/MyDrive/data\n",
        "```\n",
        "\n",
        "It might be helpful to use an ```ls``` command to establish the exact path for your images."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wwpXtlrCCRHe"
      },
      "outputs": [],
      "source": [
        "!ls /content/drive/MyDrive/data/gan/images"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rl6ShUWrCUrV"
      },
      "source": [
        "## Convert Your Images\n",
        "\n",
        "You must convert your images into a data set form that PyTorch can directly utilize. The following command converts your images and writes the resulting data set to another directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qOK8m0N-Caqa"
      },
      "outputs": [],
      "source": [
        "CMD = \"python /content/stylegan2-ada-pytorch/dataset_tool.py \"\\\n",
        "  \"--source /content/drive/MyDrive/data/gan/images/circuit \"\\\n",
        "  \"--dest /content/drive/MyDrive/data/gan/dataset/circuit\"\n",
        "\n",
        "!{CMD}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LDbsiP35CdgK"
      },
      "source": [
        "You can use the following command to clear out the newly created dataset.  If something goes wrong and you need to clean up your images and rerun the above command, you should delete your partially completed dataset directory."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bpF072wiCgys"
      },
      "outputs": [],
      "source": [
        "#!rm -R /content/drive/MyDrive/data/gan/dataset/circuit/*"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L7sUkkb-CkrS"
      },
      "source": [
        "## Clean Up your Images\n",
        "\n",
        "All images must have the same dimensions and color depth.  This code can identify images that have issues."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QuWfEfEoColz"
      },
      "outputs": [],
      "source": [
        "from os import listdir\n",
        "from os.path import isfile, join\n",
        "import os\n",
        "from PIL import Image\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "IMAGE_PATH = '/content/drive/MyDrive/data/gan/images/fish'\n",
        "files = [f for f in listdir(IMAGE_PATH) if isfile(join(IMAGE_PATH, f))]\n",
        "\n",
        "base_size = None\n",
        "for file in tqdm(files):\n",
        "  file2 = os.path.join(IMAGE_PATH,file)\n",
        "  img = Image.open(file2)\n",
        "  sz = img.size\n",
        "  if base_size and sz!=base_size:\n",
        "    print(f\"Inconsistant size: {file2}\")\n",
        "  elif img.mode!='RGB':\n",
        "    print(f\"Inconsistant color format: {file2}\")\n",
        "  else:\n",
        "    base_size = sz\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_HyFZgnjCsPb"
      },
      "source": [
        "## Perform Initial Training\n",
        "\n",
        "This code performs the initial training.  Set SNAP low enough to get a snapshot before Colab forces you to quit."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4-ZBcql7Cv14"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Modify these to suit your needs\n",
        "EXPERIMENTS = \"/content/drive/MyDrive/data/gan/experiments\"\n",
        "DATA = \"/content/drive/MyDrive/data/gan/dataset/circuit\"\n",
        "SNAP = 10\n",
        "\n",
        "# Build the command and run it\n",
        "cmd = f\"/usr/bin/python3 /content/stylegan2-ada-pytorch/train.py \"\\\n",
        "  f\"--snap {SNAP} --outdir {EXPERIMENTS} --data {DATA}\"\n",
        "!{cmd}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JhOQRZrJDCvA"
      },
      "source": [
        "## Resume Training\n",
        "\n",
        "You can now resume training after you are interrupted by something in the pervious step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P1TFgr9MDJlS"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "# Modify these to suit your needs\n",
        "EXPERIMENTS = \"/content/drive/MyDrive/data/gan/experiments\"\n",
        "NETWORK = \"network-snapshot-000100.pkl\"\n",
        "RESUME = os.path.join(EXPERIMENTS, \\\n",
        "                \"00008-circuit-auto1-resumecustom\", NETWORK)\n",
        "DATA = \"/content/drive/MyDrive/data/gan/dataset/circuit\"\n",
        "SNAP = 10\n",
        "\n",
        "# Build the command and run it\n",
        "cmd = f\"/usr/bin/python3 /content/stylegan2-ada-pytorch/train.py \"\\\n",
        "  f\"--snap {SNAP} --resume {RESUME} --outdir {EXPERIMENTS} --data {DATA}\"\n",
        "!{cmd}"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "anaconda-cloud": {},
    "colab": {
      "collapsed_sections": [],
      "name": "fix of t81_558_class_07_2_train_gan.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.9 (tensorflow)",
      "language": "python",
      "name": "tensorflow"
    },
    "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.9.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}