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": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 156
    },
    "id": "Y8_-1h5ddiDp",
    "outputId": "c623348e-940c-48b5-be88-c5998b0e8314"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Go to this URL in a browser: https://accounts.google.com/o/oauth2/auth?client_id=947318989803-6bn6qk8qdgf4n4g3pfee6491hc0brc4i.apps.googleusercontent.com&redirect_uri=urn%3aietf%3awg%3aoauth%3a2.0%3aoob&response_type=code&scope=email%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdocs.test%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive%20https%3a%2f%2fwww.googleapis.com%2fauth%2fdrive.photos.readonly%20https%3a%2f%2fwww.googleapis.com%2fauth%2fpeopleapi.readonly\n",
      "\n",
      "Enter your authorization code:\n",
      "··········\n",
      "Mounted at /content/drive\n",
      "Note: using Google CoLab\n",
      "TensorFlow 2.x selected.\n"
     ]
    }
   ],
   "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 Own Images\n",
    "\n",
    "Training GANs with StyleGAN is resource-intensive. For the high-resolution face GANs trained by NVIDIA, the researchers used computers with 8 high-end GPUs. 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, they will not be trained as well as StyleGAN3.\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. GANs can be trained 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",
    "Your training data and trained neural networks will be stored 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",
    "The drive is mounted to 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 greatly 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. It is important to learn an optimal snapshot-size for your resolution and training data.\n",
    "\n",
    "We will mount GDRIVE so that your snapshots are saved 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": 1,
   "metadata": {
    "id": "9LVr5sV8CKz0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[31mERROR: Could not find a version that satisfies the requirement torch==1.8.1 (from versions: 1.9.0, 1.10.0, 1.10.1)\u001b[0m\n",
      "\u001b[31mERROR: No matching distribution found for torch==1.8.1\u001b[0m\n",
      "Cloning into 'stylegan2-ada-pytorch'...\n",
      "remote: Enumerating objects: 128, done.\u001b[K\n",
      "remote: Total 128 (delta 0), reused 0 (delta 0), pack-reused 128\u001b[K\n",
      "Receiving objects: 100% (128/128), 1.12 MiB | 2.33 MiB/s, done.\n",
      "Resolving deltas: 100% (58/58), done.\n",
      "Requirement already satisfied: ninja in /Users/jeff/miniforge3/lib/python3.9/site-packages (1.10.2.2)\n"
     ]
    }
   ],
   "source": [
    "!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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qOK8m0N-Caqa"
   },
   "outputs": [],
   "source": [
    "!python /content/stylegan2-ada-pytorch/dataset_tool.py --source /content/drive/MyDrive/data/gan/images/circuit --dest /content/drive/MyDrive/data/gan/dataset/circuit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LDbsiP35CdgK"
   },
   "source": [
    "The following command can be used 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 created 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",
    "It is important that all images 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 that you 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 --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, \"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 --snap {SNAP} --resume {RESUME} --outdir {EXPERIMENTS} --data {DATA}\"\n",
    "!{cmd}"
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "colab": {
   "collapsed_sections": [],
   "name": "Copy 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": 1
}