styleganportraits.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "styleganportraits.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ak9250/stylegan-art/blob/master/styleganportraits.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eZh_nbVi0xiT",
        "colab_type": "text"
      },
      "source": [
        "Setup:\n",
        "\n",
        "1) Make sure GPU is enabled, go to edit->notebook settings->Hardware Accelerator GPU\n",
        "\n",
        "2) make a copy to your google drive, click on copy to drive in panel\n",
        "\n",
        "Note: colab will reset after 12 hours make sure to save your model checkpoints to google drive around 6 hours mark (one tick), the second tick will show up before the reset but I have lost work by not being able to save it to drive in time, then go to runtime->reset all runtimes. Now copy your train model back into colab and start training again from the previous checkpoint.\n",
        "\n",
        "It will take about 6 hours per tick on a k80 which colab currently uses to train one tick with transfer learning for stylegan, I usually save my checkpoint at that point and reset all runtimes to start training again."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ehlPXfTzK-Dm",
        "colab_type": "text"
      },
      "source": [
        "Clone the repo from github"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VpuHnDC29LyO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/NVlabs/stylegan.git"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NLHoBrULLC-2",
        "colab_type": "text"
      },
      "source": [
        "cd into the directory"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5-myoXv_9yqA",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cd stylegan/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dm8PIdV4LH2F",
        "colab_type": "text"
      },
      "source": [
        "download dataset from google drive, (dataset from giacomo nanni dcgan-art-tensorflow)(https://github.com/sinanatra/DCGAN-Art-Tensorflow) (optional if you dont plan to scrape your own dataset and just testing on portraits)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bt5EjdTPXoTf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!gdown https://drive.google.com/uc?id=1DMGmyr-SuB8ugKMafxPEYBe-4ta30aUf"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AEC9SEFRzu-_",
        "colab_type": "text"
      },
      "source": [
        "another wikiart portrait dataset I scrapped recently, about 5k images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UJJH2dpWzuFW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!wget https://www.dropbox.com/s/yk02zgrzddztazx/portrait5k.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Di90KT8sT1uH",
        "colab_type": "text"
      },
      "source": [
        "if you need images from google, use a google images scrapper "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jsfXa3rZT5FG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!apt-get update\n",
        "!apt install chromium-chromedriver"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HsLLlLr2T9pp",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install selenium"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ETbE_e96T_3S",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import sys\n",
        "sys.path.insert(0,'/usr/lib/chromium-browser/chromedriver')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zt0kJV_EUEIN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from selenium import webdriver\n",
        "\n",
        "chrome_options = webdriver.ChromeOptions()\n",
        "chrome_options.add_argument('--headless')\n",
        "chrome_options.add_argument('--no-sandbox')\n",
        "chrome_options.add_argument('--disable-dev-shm-usage')\n",
        "wd = webdriver.Chrome('chromedriver',options=chrome_options)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Rs3OtWw2UGZ_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install google_images_download"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J-kjqBtbUIvs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!googleimagesdownload --keywords \"portrait art\" --limit 2000 -cd /usr/bin/chromedriver"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UlaYVdq6h0Kg",
        "colab_type": "text"
      },
      "source": [
        "or use a wikiart scraper"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KdSLK2wuh25K",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!git clone https://github.com/lucasdavid/wikiart.git"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QKWcJSMZh56d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cd wikiart"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6XETR4l4iGNV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir dataset"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QUtQPNFUh7UR",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python3 wikiart.py --datadir /content/stylegan/wikiart/dataset/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ioSyCs_p1XAb",
        "colab_type": "text"
      },
      "source": [
        "or use gallery-dl to scrape flickr, tumblr etc..\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RCTSaENU1TV7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install gallery-dl"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5BtDo04R2Sbh",
        "colab_type": "text"
      },
      "source": [
        "add range for how many images you want to scrape example uses flickr (note when scrapping I noticed colab will freeze around 4k to 5k images so either scrape 1k multiple times or you might have to restart colab if you set range to 1-5000)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v-_YTlw8186-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!gallery-dl --range 1-1000 https://www.flickr.com/search/?text=portrait%20art"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GApMHjxgLNEN",
        "colab_type": "text"
      },
      "source": [
        "make directory for portraits (if you are using the default portraits datasets provided)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2QH-nX8r1a2_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir portraitbig"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ws5BRkaLcIr",
        "colab_type": "text"
      },
      "source": [
        "move portrait zip to directory"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XrkO6LXBYldQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mv /content/portrait.zip /content/stylegan/portraitbig/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PJr6L8Fu8_h3",
        "colab_type": "text"
      },
      "source": [
        "unzip portrait zip file"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xQMVux4j8m7d",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cd portraitbig/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8iBCLma_9sDi",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!unzip portrait.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8cJnnozq-8id",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm portrait.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TwVtYSdsw15l",
        "colab_type": "text"
      },
      "source": [
        "if not doing faces and need resize and/or augment images use Augmentor"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zQVQzBnww-7T",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install Augmentor"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SscWjN1vw_34",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import Augmentor\n",
        "p = Augmentor.Pipeline(\"/content/stylegan/portraitbig/\", \"/path/to/output/folder/\", save_format=\"JPEG\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WgP1joV3xF6N",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "p.resize(probability=1.0, width=512, height=512)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uVJBzTzh2Njx",
        "colab_type": "text"
      },
      "source": [
        "run this cell only for augmentation, see https://augmentor.readthedocs.io/en/master/ for different  features, this is for changes in color"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1iDQvhCw2QG2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "p.random_color(probability=1.0,min_factor=0.5,max_factor=0.9)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jhoiBr1tCS4N",
        "colab_type": "text"
      },
      "source": [
        "use p.process to just resize p.sample and a higher number than original dataset to augment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H9fLUcb71VKr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "p.process()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hhJGqB3BxIE-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "p.sample(5000)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zoqqMYmR2iMp",
        "colab_type": "text"
      },
      "source": [
        "remove any images that cant be opened"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ytdL_Qqw2kkj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from fastai.vision import *\n",
        "verify_images('/path/to/dataset/', delete=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dRwv1do1Lf2J",
        "colab_type": "text"
      },
      "source": [
        "install autocrop to crop faces if training faces"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2HSnh8kd5Z7D",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!pip install autocrop"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z8ajb_v-LkqH",
        "colab_type": "text"
      },
      "source": [
        "Make directory for faces and rejected faces and then run autocrop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "POf7T93_24dl",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cd .."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "axagObqQ2Dgn",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir Faces"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "me3ZuslxNXBY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!mkdir Rejects"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ohZ-YK59MZDh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!autocrop -i /path/to/dataset/ -o /content/stylegan/Faces/ -w 512 -H 512 --facePercent 50 -r /content/stylegan/Rejects/ "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kV1o_J_0LsAP",
        "colab_type": "text"
      },
      "source": [
        "download latest model for transfer learning"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uoBBmmWptRYo",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!gdown https://drive.google.com/uc?id=1cJQtMeTy_QldOP7n64F8stCDXY6Esup9"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DieqF2rbLvxq",
        "colab_type": "text"
      },
      "source": [
        "generate tfrecords"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C1ME1QLCtme5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!python dataset_tool.py create_from_images datasets/smalls/ /content/stylegan/Faces/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rC5RLqIGLyqS",
        "colab_type": "text"
      },
      "source": [
        "change any settings if needed"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tLsyyZkXMVbj",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cd training/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yJIqRrpfiwHQ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%%writefile training_loop.py \n",
        "# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.\n",
        "#\n",
        "# This work is licensed under the Creative Commons Attribution-NonCommercial\n",
        "# 4.0 International License. To view a copy of this license, visit\n",
        "# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to\n",
        "# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.\n",
        "\n",
        "\"\"\"Main training script.\"\"\"\n",
        "\n",
        "import os\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "from dnnlib.tflib.autosummary import autosummary\n",
        "\n",
        "import config\n",
        "import train\n",
        "from training import dataset\n",
        "from training import misc\n",
        "from metrics import metric_base\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "# Just-in-time processing of training images before feeding them to the networks.\n",
        "\n",
        "def process_reals(x, lod, mirror_augment, drange_data, drange_net):\n",
        "    with tf.name_scope('ProcessReals'):\n",
        "        with tf.name_scope('DynamicRange'):\n",
        "            x = tf.cast(x, tf.float32)\n",
        "            x = misc.adjust_dynamic_range(x, drange_data, drange_net)\n",
        "        if mirror_augment:\n",
        "            with tf.name_scope('MirrorAugment'):\n",
        "                s = tf.shape(x)\n",
        "                mask = tf.random_uniform([s[0], 1, 1, 1], 0.0, 1.0)\n",
        "                mask = tf.tile(mask, [1, s[1], s[2], s[3]])\n",
        "                x = tf.where(mask < 0.5, x, tf.reverse(x, axis=[3]))\n",
        "        with tf.name_scope('FadeLOD'): # Smooth crossfade between consecutive levels-of-detail.\n",
        "            s = tf.shape(x)\n",
        "            y = tf.reshape(x, [-1, s[1], s[2]//2, 2, s[3]//2, 2])\n",
        "            y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)\n",
        "            y = tf.tile(y, [1, 1, 1, 2, 1, 2])\n",
        "            y = tf.reshape(y, [-1, s[1], s[2], s[3]])\n",
        "            x = tflib.lerp(x, y, lod - tf.floor(lod))\n",
        "        with tf.name_scope('UpscaleLOD'): # Upscale to match the expected input/output size of the networks.\n",
        "            s = tf.shape(x)\n",
        "            factor = tf.cast(2 ** tf.floor(lod), tf.int32)\n",
        "            x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])\n",
        "            x = tf.tile(x, [1, 1, 1, factor, 1, factor])\n",
        "            x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])\n",
        "        return x\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "# Evaluate time-varying training parameters.\n",
        "\n",
        "def training_schedule(\n",
        "    cur_nimg,\n",
        "    training_set,\n",
        "    num_gpus,\n",
        "    lod_initial_resolution  = 4,        # Image resolution used at the beginning.\n",
        "    lod_training_kimg       = 600,      # Thousands of real images to show before doubling the resolution.\n",
        "    lod_transition_kimg     = 600,      # Thousands of real images to show when fading in new layers.\n",
        "    minibatch_base          = 16,       # Maximum minibatch size, divided evenly among GPUs.\n",
        "    minibatch_dict          = {},       # Resolution-specific overrides.\n",
        "    max_minibatch_per_gpu   = {},       # Resolution-specific maximum minibatch size per GPU.\n",
        "    G_lrate_base            = 0.001,    # Learning rate for the generator.\n",
        "    G_lrate_dict            = {},       # Resolution-specific overrides.\n",
        "    D_lrate_base            = 0.001,    # Learning rate for the discriminator.\n",
        "    D_lrate_dict            = {},       # Resolution-specific overrides.\n",
        "    lrate_rampup_kimg       = 0,        # Duration of learning rate ramp-up.\n",
        "    tick_kimg_base          = 160,      # Default interval of progress snapshots.\n",
        "    tick_kimg_dict          = {4: 160, 8:140, 16:120, 32:100, 64:80, 128:60, 256:40, 512:30, 1024:20}): # Resolution-specific overrides.\n",
        "\n",
        "    # Initialize result dict.\n",
        "    s = dnnlib.EasyDict()\n",
        "    s.kimg = cur_nimg / 1000.0\n",
        "\n",
        "    # Training phase.\n",
        "    phase_dur = lod_training_kimg + lod_transition_kimg\n",
        "    phase_idx = int(np.floor(s.kimg / phase_dur)) if phase_dur > 0 else 0\n",
        "    phase_kimg = s.kimg - phase_idx * phase_dur\n",
        "\n",
        "    # Level-of-detail and resolution.\n",
        "    s.lod = training_set.resolution_log2\n",
        "    s.lod -= np.floor(np.log2(lod_initial_resolution))\n",
        "    s.lod -= phase_idx\n",
        "    if lod_transition_kimg > 0:\n",
        "        s.lod -= max(phase_kimg - lod_training_kimg, 0.0) / lod_transition_kimg\n",
        "    s.lod = max(s.lod, 0.0)\n",
        "    s.resolution = 2 ** (training_set.resolution_log2 - int(np.floor(s.lod)))\n",
        "\n",
        "    # Minibatch size.\n",
        "    s.minibatch = minibatch_dict.get(s.resolution, minibatch_base)\n",
        "    s.minibatch -= s.minibatch % num_gpus\n",
        "    if s.resolution in max_minibatch_per_gpu:\n",
        "        s.minibatch = min(s.minibatch, max_minibatch_per_gpu[s.resolution] * num_gpus)\n",
        "\n",
        "    # Learning rate.\n",
        "    s.G_lrate = G_lrate_dict.get(s.resolution, G_lrate_base)\n",
        "    s.D_lrate = D_lrate_dict.get(s.resolution, D_lrate_base)\n",
        "    if lrate_rampup_kimg > 0:\n",
        "        rampup = min(s.kimg / lrate_rampup_kimg, 1.0)\n",
        "        s.G_lrate *= rampup\n",
        "        s.D_lrate *= rampup\n",
        "\n",
        "    # Other parameters.\n",
        "    s.tick_kimg = tick_kimg_dict.get(s.resolution, tick_kimg_base)\n",
        "    return s\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "# Main training script.\n",
        "\n",
        "def training_loop(\n",
        "    submit_config,\n",
        "    G_args                  = {},       # Options for generator network.\n",
        "    D_args                  = {},       # Options for discriminator network.\n",
        "    G_opt_args              = {},       # Options for generator optimizer.\n",
        "    D_opt_args              = {},       # Options for discriminator optimizer.\n",
        "    G_loss_args             = {},       # Options for generator loss.\n",
        "    D_loss_args             = {},       # Options for discriminator loss.\n",
        "    dataset_args            = {},       # Options for dataset.load_dataset().\n",
        "    sched_args              = {},       # Options for train.TrainingSchedule.\n",
        "    grid_args               = {},       # Options for train.setup_snapshot_image_grid().\n",
        "    metric_arg_list         = [],       # Options for MetricGroup.\n",
        "    tf_config               = {},       # Options for tflib.init_tf().\n",
        "    G_smoothing_kimg        = 10.0,     # Half-life of the running average of generator weights.\n",
        "    D_repeats               = 1,        # How many times the discriminator is trained per G iteration.\n",
        "    minibatch_repeats       = 4,        # Number of minibatches to run before adjusting training parameters.\n",
        "    reset_opt_for_new_lod   = True,     # Reset optimizer internal state (e.g. Adam moments) when new layers are introduced?\n",
        "    total_kimg              = 15000,    # Total length of the training, measured in thousands of real images.\n",
        "    mirror_augment          = False,    # Enable mirror augment?\n",
        "    drange_net              = [-1,1],   # Dynamic range used when feeding image data to the networks.\n",
        "    image_snapshot_ticks    = 1,        # How often to export image snapshots?\n",
        "    network_snapshot_ticks  = 1,       # How often to export network snapshots?\n",
        "    save_tf_graph           = False,    # Include full TensorFlow computation graph in the tfevents file?\n",
        "    save_weight_histograms  = False,    # Include weight histograms in the tfevents file?\n",
        "    resume_run_id           = \"/content/stylegan/network-snapshot-011125.pkl\",     # Run ID or network pkl to resume training from, None = start from scratch.\n",
        "    resume_snapshot         = None,     # Snapshot index to resume training from, None = autodetect.\n",
        "    resume_kimg             = 11125,      # Assumed training progress at the beginning. Affects reporting and training schedule.\n",
        "    resume_time             = 0.0):     # Assumed wallclock time at the beginning. Affects reporting.\n",
        "\n",
        "    # Initialize dnnlib and TensorFlow.\n",
        "    ctx = dnnlib.RunContext(submit_config, train)\n",
        "    tflib.init_tf(tf_config)\n",
        "\n",
        "    # Load training set.\n",
        "    training_set = dataset.load_dataset(data_dir=config.data_dir, verbose=True, **dataset_args)\n",
        "\n",
        "    # Construct networks.\n",
        "    with tf.device('/gpu:0'):\n",
        "        if resume_run_id is not None:\n",
        "            network_pkl = misc.locate_network_pkl(resume_run_id, resume_snapshot)\n",
        "            print('Loading networks from \"%s\"...' % network_pkl)\n",
        "            G, D, Gs = misc.load_pkl(network_pkl)\n",
        "        else:\n",
        "            print('Constructing networks...')\n",
        "            G = tflib.Network('G', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **G_args)\n",
        "            D = tflib.Network('D', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **D_args)\n",
        "            Gs = G.clone('Gs')\n",
        "    G.print_layers(); D.print_layers()\n",
        "\n",
        "    print('Building TensorFlow graph...')\n",
        "    with tf.name_scope('Inputs'), tf.device('/cpu:0'):\n",
        "        lod_in          = tf.placeholder(tf.float32, name='lod_in', shape=[])\n",
        "        lrate_in        = tf.placeholder(tf.float32, name='lrate_in', shape=[])\n",
        "        minibatch_in    = tf.placeholder(tf.int32, name='minibatch_in', shape=[])\n",
        "        minibatch_split = minibatch_in // submit_config.num_gpus\n",
        "        Gs_beta         = 0.5 ** tf.div(tf.cast(minibatch_in, tf.float32), G_smoothing_kimg * 1000.0) if G_smoothing_kimg > 0.0 else 0.0\n",
        "\n",
        "    G_opt = tflib.Optimizer(name='TrainG', learning_rate=lrate_in, **G_opt_args)\n",
        "    D_opt = tflib.Optimizer(name='TrainD', learning_rate=lrate_in, **D_opt_args)\n",
        "    for gpu in range(submit_config.num_gpus):\n",
        "        with tf.name_scope('GPU%d' % gpu), tf.device('/gpu:%d' % gpu):\n",
        "            G_gpu = G if gpu == 0 else G.clone(G.name + '_shadow')\n",
        "            D_gpu = D if gpu == 0 else D.clone(D.name + '_shadow')\n",
        "            lod_assign_ops = [tf.assign(G_gpu.find_var('lod'), lod_in), tf.assign(D_gpu.find_var('lod'), lod_in)]\n",
        "            reals, labels = training_set.get_minibatch_tf()\n",
        "            reals = process_reals(reals, lod_in, mirror_augment, training_set.dynamic_range, drange_net)\n",
        "            with tf.name_scope('G_loss'), tf.control_dependencies(lod_assign_ops):\n",
        "                G_loss = dnnlib.util.call_func_by_name(G=G_gpu, D=D_gpu, opt=G_opt, training_set=training_set, minibatch_size=minibatch_split, **G_loss_args)\n",
        "            with tf.name_scope('D_loss'), tf.control_dependencies(lod_assign_ops):\n",
        "                D_loss = dnnlib.util.call_func_by_name(G=G_gpu, D=D_gpu, opt=D_opt, training_set=training_set, minibatch_size=minibatch_split, reals=reals, labels=labels, **D_loss_args)\n",
        "            G_opt.register_gradients(tf.reduce_mean(G_loss), G_gpu.trainables)\n",
        "            D_opt.register_gradients(tf.reduce_mean(D_loss), D_gpu.trainables)\n",
        "    G_train_op = G_opt.apply_updates()\n",
        "    D_train_op = D_opt.apply_updates()\n",
        "\n",
        "    Gs_update_op = Gs.setup_as_moving_average_of(G, beta=Gs_beta)\n",
        "    with tf.device('/gpu:0'):\n",
        "        try:\n",
        "            peak_gpu_mem_op = tf.contrib.memory_stats.MaxBytesInUse()\n",
        "        except tf.errors.NotFoundError:\n",
        "            peak_gpu_mem_op = tf.constant(0)\n",
        "\n",
        "    print('Setting up snapshot image grid...')\n",
        "    grid_size, grid_reals, grid_labels, grid_latents = misc.setup_snapshot_image_grid(G, training_set, **grid_args)\n",
        "    sched = training_schedule(cur_nimg=total_kimg*1000, training_set=training_set, num_gpus=submit_config.num_gpus, **sched_args)\n",
        "    grid_fakes = Gs.run(grid_latents, grid_labels, is_validation=True, minibatch_size=sched.minibatch//submit_config.num_gpus)\n",
        "\n",
        "    print('Setting up run dir...')\n",
        "    misc.save_image_grid(grid_reals, os.path.join(submit_config.run_dir, 'reals.png'), drange=training_set.dynamic_range, grid_size=grid_size)\n",
        "    misc.save_image_grid(grid_fakes, os.path.join(submit_config.run_dir, 'fakes%06d.png' % resume_kimg), drange=drange_net, grid_size=grid_size)\n",
        "    summary_log = tf.summary.FileWriter(submit_config.run_dir)\n",
        "    if save_tf_graph:\n",
        "        summary_log.add_graph(tf.get_default_graph())\n",
        "    if save_weight_histograms:\n",
        "        G.setup_weight_histograms(); D.setup_weight_histograms()\n",
        "    metrics = metric_base.MetricGroup(metric_arg_list)\n",
        "\n",
        "    print('Training...\\n')\n",
        "    ctx.update('', cur_epoch=resume_kimg, max_epoch=total_kimg)\n",
        "    maintenance_time = ctx.get_last_update_interval()\n",
        "    cur_nimg = int(resume_kimg * 1000)\n",
        "    cur_tick = 0\n",
        "    tick_start_nimg = cur_nimg\n",
        "    prev_lod = -1.0\n",
        "    while cur_nimg < total_kimg * 1000:\n",
        "        if ctx.should_stop(): break\n",
        "\n",
        "        # Choose training parameters and configure training ops.\n",
        "        sched = training_schedule(cur_nimg=cur_nimg, training_set=training_set, num_gpus=submit_config.num_gpus, **sched_args)\n",
        "        training_set.configure(sched.minibatch // submit_config.num_gpus, sched.lod)\n",
        "        if reset_opt_for_new_lod:\n",
        "            if np.floor(sched.lod) != np.floor(prev_lod) or np.ceil(sched.lod) != np.ceil(prev_lod):\n",
        "                G_opt.reset_optimizer_state(); D_opt.reset_optimizer_state()\n",
        "        prev_lod = sched.lod\n",
        "\n",
        "        # Run training ops.\n",
        "        for _mb_repeat in range(minibatch_repeats):\n",
        "            for _D_repeat in range(D_repeats):\n",
        "                tflib.run([D_train_op, Gs_update_op], {lod_in: sched.lod, lrate_in: sched.D_lrate, minibatch_in: sched.minibatch})\n",
        "                cur_nimg += sched.minibatch\n",
        "            tflib.run([G_train_op], {lod_in: sched.lod, lrate_in: sched.G_lrate, minibatch_in: sched.minibatch})\n",
        "\n",
        "        # Perform maintenance tasks once per tick.\n",
        "        done = (cur_nimg >= total_kimg * 1000)\n",
        "        if cur_nimg >= tick_start_nimg + sched.tick_kimg * 1000 or done:\n",
        "            cur_tick += 1\n",
        "            tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0\n",
        "            tick_start_nimg = cur_nimg\n",
        "            tick_time = ctx.get_time_since_last_update()\n",
        "            total_time = ctx.get_time_since_start() + resume_time\n",
        "\n",
        "            # Report progress.\n",
        "            print('tick %-5d kimg %-8.1f lod %-5.2f minibatch %-4d time %-12s sec/tick %-7.1f sec/kimg %-7.2f maintenance %-6.1f gpumem %-4.1f' % (\n",
        "                autosummary('Progress/tick', cur_tick),\n",
        "                autosummary('Progress/kimg', cur_nimg / 1000.0),\n",
        "                autosummary('Progress/lod', sched.lod),\n",
        "                autosummary('Progress/minibatch', sched.minibatch),\n",
        "                dnnlib.util.format_time(autosummary('Timing/total_sec', total_time)),\n",
        "                autosummary('Timing/sec_per_tick', tick_time),\n",
        "                autosummary('Timing/sec_per_kimg', tick_time / tick_kimg),\n",
        "                autosummary('Timing/maintenance_sec', maintenance_time),\n",
        "                autosummary('Resources/peak_gpu_mem_gb', peak_gpu_mem_op.eval() / 2**30)))\n",
        "            autosummary('Timing/total_hours', total_time / (60.0 * 60.0))\n",
        "            autosummary('Timing/total_days', total_time / (24.0 * 60.0 * 60.0))\n",
        "\n",
        "            # Save snapshots.\n",
        "            if cur_tick % image_snapshot_ticks == 0 or done:\n",
        "                grid_fakes = Gs.run(grid_latents, grid_labels, is_validation=True, minibatch_size=sched.minibatch//submit_config.num_gpus)\n",
        "                misc.save_image_grid(grid_fakes, os.path.join(submit_config.run_dir, 'fakes%06d.png' % (cur_nimg // 1000)), drange=drange_net, grid_size=grid_size)\n",
        "            if cur_tick % network_snapshot_ticks == 0 or done or cur_tick == 1:\n",
        "                pkl = os.path.join(submit_config.run_dir, 'network-snapshot-%06d.pkl' % (cur_nimg // 1000))\n",
        "                misc.save_pkl((G, D, Gs), pkl)\n",
        "                metrics.run(pkl, run_dir=submit_config.run_dir, num_gpus=submit_config.num_gpus, tf_config=tf_config)\n",
        "\n",
        "            # Update summaries and RunContext.\n",
        "            metrics.update_autosummaries()\n",
        "            tflib.autosummary.save_summaries(summary_log, cur_nimg)\n",
        "            ctx.update('%.2f' % sched.lod, cur_epoch=cur_nimg // 1000, max_epoch=total_kimg)\n",
        "            maintenance_time = ctx.get_last_update_interval() - tick_time\n",
        "\n",
        "    # Write final results.\n",
        "    misc.save_pkl((G, D, Gs), os.path.join(submit_config.run_dir, 'network-final.pkl'))\n",
        "    summary_log.close()\n",
        "\n",
        "    ctx.close()\n",
        "\n",
        "#----------------------------------------------------------------------------\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j7Byxw69MkfU",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "cd .."
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lQ7KoklS1I33",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.\n",
        "#\n",
        "# This work is licensed under the Creative Commons Attribution-NonCommercial\n",
        "# 4.0 International License. To view a copy of this license, visit\n",
        "# http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to\n",
        "# Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.\n",
        "\n",
        "\"\"\"Main entry point for training StyleGAN and ProGAN networks.\"\"\"\n",
        "\n",
        "import copy\n",
        "import dnnlib\n",
        "from dnnlib import EasyDict\n",
        "\n",
        "import config\n",
        "from metrics import metric_base\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "# Official training configs for StyleGAN, targeted mainly for FFHQ.\n",
        "\n",
        "if 1:\n",
        "    desc          = 'sgan'                                                                 # Description string included in result subdir name.\n",
        "    train         = EasyDict(run_func_name='training.training_loop.training_loop')         # Options for training loop.\n",
        "    G             = EasyDict(func_name='training.networks_stylegan.G_style')               # Options for generator network.\n",
        "    D             = EasyDict(func_name='training.networks_stylegan.D_basic')               # Options for discriminator network.\n",
        "    G_opt         = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)                          # Options for generator optimizer.\n",
        "    D_opt         = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)                          # Options for discriminator optimizer.\n",
        "    G_loss        = EasyDict(func_name='training.loss.G_logistic_nonsaturating')           # Options for generator loss.\n",
        "    D_loss        = EasyDict(func_name='training.loss.D_logistic_simplegp', r1_gamma=10.0) # Options for discriminator loss.\n",
        "    dataset       = EasyDict()                                                             # Options for load_dataset().\n",
        "    sched         = EasyDict()                                                             # Options for TrainingSchedule.\n",
        "    grid          = EasyDict(size='4k', layout='random')                                   # Options for setup_snapshot_image_grid().\n",
        "    #metrics       = [metric_base.fid50k]                                                   # Options for MetricGroup.\n",
        "    submit_config = dnnlib.SubmitConfig()                                                  # Options for dnnlib.submit_run().\n",
        "    tf_config     = {'rnd.np_random_seed': 1000}                                           # Options for tflib.init_tf().\n",
        "\n",
        "    # Dataset.\n",
        "    desc += '-custom';     dataset = EasyDict(tfrecord_dir='smalls');              train.mirror_augment = True\n",
        "    #desc += '-celebahq'; dataset = EasyDict(tfrecord_dir='celebahq');          train.mirror_augment = True\n",
        "    #desc += '-bedroom';  dataset = EasyDict(tfrecord_dir='lsun-bedroom-full'); train.mirror_augment = False\n",
        "    #desc += '-car';      dataset = EasyDict(tfrecord_dir='lsun-car-512x384');  train.mirror_augment = False\n",
        "    #desc += '-cat';      dataset = EasyDict(tfrecord_dir='lsun-cat-full');     train.mirror_augment = False\n",
        "\n",
        "    # Number of GPUs.\n",
        "    desc += '-1gpu'; submit_config.num_gpus = 1; sched.minibatch_base = 4; sched.minibatch_dict = {4: 128, 8: 128, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8, 512: 4}\n",
        "    #desc += '-2gpu'; submit_config.num_gpus = 2; sched.minibatch_base = 8; sched.minibatch_dict = {4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8}\n",
        "    #desc += '-4gpu'; submit_config.num_gpus = 4; sched.minibatch_base = 16; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16}\n",
        "    #desc += '-8gpu'; submit_config.num_gpus = 8; sched.minibatch_base = 32; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32}\n",
        "\n",
        "    # Default options.\n",
        "    train.total_kimg = 25000\n",
        "    sched.lod_initial_resolution = 8\n",
        "    sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}\n",
        "    sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)\n",
        "\n",
        "    # WGAN-GP loss for CelebA-HQ.\n",
        "    #desc += '-wgangp'; G_loss = EasyDict(func_name='training.loss.G_wgan'); D_loss = EasyDict(func_name='training.loss.D_wgan_gp'); sched.G_lrate_dict = {k: min(v, 0.002) for k, v in sched.G_lrate_dict.items()}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict)\n",
        "\n",
        "    # Table 1.\n",
        "    #desc += '-tuned-baseline'; G.use_styles = False; G.use_pixel_norm = True; G.use_instance_norm = False; G.mapping_layers = 0; G.truncation_psi = None; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False\n",
        "    #desc += '-add-mapping-and-styles'; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False\n",
        "    #desc += '-remove-traditional-input'; G.style_mixing_prob = 0.0; G.use_noise = False\n",
        "    #desc += '-add-noise-inputs'; G.style_mixing_prob = 0.0\n",
        "    #desc += '-mixing-regularization' # default\n",
        "\n",
        "    # Table 2.\n",
        "    #desc += '-mix0'; G.style_mixing_prob = 0.0\n",
        "    #desc += '-mix50'; G.style_mixing_prob = 0.5\n",
        "    #desc += '-mix90'; G.style_mixing_prob = 0.9 # default\n",
        "    #desc += '-mix100'; G.style_mixing_prob = 1.0\n",
        "\n",
        "    # Table 4.\n",
        "    #desc += '-traditional-0'; G.use_styles = False; G.use_pixel_norm = True; G.use_instance_norm = False; G.mapping_layers = 0; G.truncation_psi = None; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False\n",
        "    #desc += '-traditional-8'; G.use_styles = False; G.use_pixel_norm = True; G.use_instance_norm = False; G.mapping_layers = 8; G.truncation_psi = None; G.const_input_layer = False; G.style_mixing_prob = 0.0; G.use_noise = False\n",
        "    #desc += '-stylebased-0'; G.mapping_layers = 0\n",
        "    #desc += '-stylebased-1'; G.mapping_layers = 1\n",
        "    #desc += '-stylebased-2'; G.mapping_layers = 2\n",
        "    #desc += '-stylebased-8'; G.mapping_layers = 8 # default\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "# Official training configs for Progressive GAN, targeted mainly for CelebA-HQ.\n",
        "\n",
        "if 0:\n",
        "    desc          = 'pgan'                                                         # Description string included in result subdir name.\n",
        "    train         = EasyDict(run_func_name='training.training_loop.training_loop') # Options for training loop.\n",
        "    G             = EasyDict(func_name='training.networks_progan.G_paper')         # Options for generator network.\n",
        "    D             = EasyDict(func_name='training.networks_progan.D_paper')         # Options for discriminator network.\n",
        "    G_opt         = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)                  # Options for generator optimizer.\n",
        "    D_opt         = EasyDict(beta1=0.0, beta2=0.99, epsilon=1e-8)                  # Options for discriminator optimizer.\n",
        "    G_loss        = EasyDict(func_name='training.loss.G_wgan')                     # Options for generator loss.\n",
        "    D_loss        = EasyDict(func_name='training.loss.D_wgan_gp')                  # Options for discriminator loss.\n",
        "    dataset       = EasyDict()                                                     # Options for load_dataset().\n",
        "    sched         = EasyDict()                                                     # Options for TrainingSchedule.\n",
        "    grid          = EasyDict(size='1080p', layout='random')                        # Options for setup_snapshot_image_grid().\n",
        "    #metrics       = [metric_base.fid50k]                                           # Options for MetricGroup.\n",
        "    submit_config = dnnlib.SubmitConfig()                                          # Options for dnnlib.submit_run().\n",
        "    tf_config     = {'rnd.np_random_seed': 1000}                                   # Options for tflib.init_tf().\n",
        "\n",
        "    # Dataset (choose one).\n",
        "    #desc += '-celebahq';            dataset = EasyDict(tfrecord_dir='celebahq'); train.mirror_augment = True\n",
        "    #desc += '-celeba';              dataset = EasyDict(tfrecord_dir='celeba'); train.mirror_augment = True\n",
        "    #desc += '-cifar10';             dataset = EasyDict(tfrecord_dir='cifar10')\n",
        "    #desc += '-cifar100';            dataset = EasyDict(tfrecord_dir='cifar100')\n",
        "    #desc += '-svhn';                dataset = EasyDict(tfrecord_dir='svhn')\n",
        "    #desc += '-mnist';               dataset = EasyDict(tfrecord_dir='mnist')\n",
        "    #desc += '-mnistrgb';            dataset = EasyDict(tfrecord_dir='mnistrgb')\n",
        "    #desc += '-syn1024rgb';          dataset = EasyDict(class_name='training.dataset.SyntheticDataset', resolution=1024, num_channels=3)\n",
        "    #desc += '-lsun-airplane';       dataset = EasyDict(tfrecord_dir='lsun-airplane-100k');       train.mirror_augment = True\n",
        "    #desc += '-lsun-bedroom';        dataset = EasyDict(tfrecord_dir='lsun-bedroom-100k');        train.mirror_augment = True\n",
        "    #desc += '-lsun-bicycle';        dataset = EasyDict(tfrecord_dir='lsun-bicycle-100k');        train.mirror_augment = True\n",
        "    #desc += '-lsun-bird';           dataset = EasyDict(tfrecord_dir='lsun-bird-100k');           train.mirror_augment = True\n",
        "    #desc += '-lsun-boat';           dataset = EasyDict(tfrecord_dir='lsun-boat-100k');           train.mirror_augment = True\n",
        "    #desc += '-lsun-bottle';         dataset = EasyDict(tfrecord_dir='lsun-bottle-100k');         train.mirror_augment = True\n",
        "    #desc += '-lsun-bridge';         dataset = EasyDict(tfrecord_dir='lsun-bridge-100k');         train.mirror_augment = True\n",
        "    #desc += '-lsun-bus';            dataset = EasyDict(tfrecord_dir='lsun-bus-100k');            train.mirror_augment = True\n",
        "    #desc += '-lsun-car';            dataset = EasyDict(tfrecord_dir='lsun-car-100k');            train.mirror_augment = True\n",
        "    #desc += '-lsun-cat';            dataset = EasyDict(tfrecord_dir='lsun-cat-100k');            train.mirror_augment = True\n",
        "    #desc += '-lsun-chair';          dataset = EasyDict(tfrecord_dir='lsun-chair-100k');          train.mirror_augment = True\n",
        "    #desc += '-lsun-churchoutdoor';  dataset = EasyDict(tfrecord_dir='lsun-churchoutdoor-100k');  train.mirror_augment = True\n",
        "    #desc += '-lsun-classroom';      dataset = EasyDict(tfrecord_dir='lsun-classroom-100k');      train.mirror_augment = True\n",
        "    #desc += '-lsun-conferenceroom'; dataset = EasyDict(tfrecord_dir='lsun-conferenceroom-100k'); train.mirror_augment = True\n",
        "    #desc += '-lsun-cow';            dataset = EasyDict(tfrecord_dir='lsun-cow-100k');            train.mirror_augment = True\n",
        "    #desc += '-lsun-diningroom';     dataset = EasyDict(tfrecord_dir='lsun-diningroom-100k');     train.mirror_augment = True\n",
        "    #desc += '-lsun-diningtable';    dataset = EasyDict(tfrecord_dir='lsun-diningtable-100k');    train.mirror_augment = True\n",
        "    #desc += '-lsun-dog';            dataset = EasyDict(tfrecord_dir='lsun-dog-100k');            train.mirror_augment = True\n",
        "    #desc += '-lsun-horse';          dataset = EasyDict(tfrecord_dir='lsun-horse-100k');          train.mirror_augment = True\n",
        "    #desc += '-lsun-kitchen';        dataset = EasyDict(tfrecord_dir='lsun-kitchen-100k');        train.mirror_augment = True\n",
        "    #desc += '-lsun-livingroom';     dataset = EasyDict(tfrecord_dir='lsun-livingroom-100k');     train.mirror_augment = True\n",
        "    #desc += '-lsun-motorbike';      dataset = EasyDict(tfrecord_dir='lsun-motorbike-100k');      train.mirror_augment = True\n",
        "    #desc += '-lsun-person';         dataset = EasyDict(tfrecord_dir='lsun-person-100k');         train.mirror_augment = True\n",
        "    #desc += '-lsun-pottedplant';    dataset = EasyDict(tfrecord_dir='lsun-pottedplant-100k');    train.mirror_augment = True\n",
        "    #desc += '-lsun-restaurant';     dataset = EasyDict(tfrecord_dir='lsun-restaurant-100k');     train.mirror_augment = True\n",
        "    #desc += '-lsun-sheep';          dataset = EasyDict(tfrecord_dir='lsun-sheep-100k');          train.mirror_augment = True\n",
        "    #desc += '-lsun-sofa';           dataset = EasyDict(tfrecord_dir='lsun-sofa-100k');           train.mirror_augment = True\n",
        "    #desc += '-lsun-tower';          dataset = EasyDict(tfrecord_dir='lsun-tower-100k');          train.mirror_augment = True\n",
        "    #desc += '-lsun-train';          dataset = EasyDict(tfrecord_dir='lsun-train-100k');          train.mirror_augment = True\n",
        "    #desc += '-lsun-tvmonitor';      dataset = EasyDict(tfrecord_dir='lsun-tvmonitor-100k');      train.mirror_augment = True\n",
        "\n",
        "    # Conditioning & snapshot options.\n",
        "    #desc += '-cond'; dataset.max_label_size = 'full' # conditioned on full label\n",
        "    #desc += '-cond1'; dataset.max_label_size = 1 # conditioned on first component of the label\n",
        "    #desc += '-g4k'; grid.size = '4k'\n",
        "    #desc += '-grpc'; grid.layout = 'row_per_class'\n",
        "\n",
        "    # Config presets (choose one).\n",
        "    #desc += '-preset-v1-1gpu'; submit_config.num_gpus = 1; D.mbstd_group_size = 16; sched.minibatch_base = 16; sched.minibatch_dict = {256: 14, 512: 6, 1024: 3}; sched.lod_training_kimg = 800; sched.lod_transition_kimg = 800; train.total_kimg = 19000\n",
        "    #desc += '-preset-v2-1gpu'; submit_config.num_gpus = 1; sched.minibatch_base = 4; sched.minibatch_dict = {4: 128, 8: 128, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8, 512: 4}; sched.G_lrate_dict = {1024: 0.0015}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000\n",
        "    #desc += '-preset-v2-2gpus'; submit_config.num_gpus = 2; sched.minibatch_base = 8; sched.minibatch_dict = {4: 256, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16, 256: 8}; sched.G_lrate_dict = {512: 0.0015, 1024: 0.002}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000\n",
        "    #desc += '-preset-v2-4gpus'; submit_config.num_gpus = 4; sched.minibatch_base = 16; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32, 128: 16}; sched.G_lrate_dict = {256: 0.0015, 512: 0.002, 1024: 0.003}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000\n",
        "    #desc += '-preset-v2-8gpus'; submit_config.num_gpus = 8; sched.minibatch_base = 32; sched.minibatch_dict = {4: 512, 8: 256, 16: 128, 32: 64, 64: 32}; sched.G_lrate_dict = {128: 0.0015, 256: 0.002, 512: 0.003, 1024: 0.003}; sched.D_lrate_dict = EasyDict(sched.G_lrate_dict); train.total_kimg = 12000\n",
        "\n",
        "    # Numerical precision (choose one).\n",
        "    #desc += '-fp32'; sched.max_minibatch_per_gpu = {256: 16, 512: 8, 1024: 4}\n",
        "    #desc += '-fp16'; G.dtype = 'float16'; D.dtype = 'float16'; G.pixelnorm_epsilon=1e-4; G_opt.use_loss_scaling = True; D_opt.use_loss_scaling = True; sched.max_minibatch_per_gpu = {512: 16, 1024: 8}\n",
        "\n",
        "    # Disable individual features.\n",
        "    #desc += '-nogrowing'; sched.lod_initial_resolution = 1024; sched.lod_training_kimg = 0; sched.lod_transition_kimg = 0; train.total_kimg = 10000\n",
        "    #desc += '-nopixelnorm'; G.use_pixelnorm = False\n",
        "    #desc += '-nowscale'; G.use_wscale = False; D.use_wscale = False\n",
        "    #desc += '-noleakyrelu'; G.use_leakyrelu = False\n",
        "    #desc += '-nosmoothing'; train.G_smoothing_kimg = 0.0\n",
        "    #desc += '-norepeat'; train.minibatch_repeats = 1\n",
        "    #desc += '-noreset'; train.reset_opt_for_new_lod = False\n",
        "\n",
        "    # Special modes.\n",
        "    #desc += '-BENCHMARK'; sched.lod_initial_resolution = 4; sched.lod_training_kimg = 3; sched.lod_transition_kimg = 3; train.total_kimg = (8*2+1)*3; sched.tick_kimg_base = 1; sched.tick_kimg_dict = {}; train.image_snapshot_ticks = 1000; train.network_snapshot_ticks = 1000\n",
        "    #desc += '-BENCHMARK0'; sched.lod_initial_resolution = 1024; train.total_kimg = 10; sched.tick_kimg_base = 1; sched.tick_kimg_dict = {}; train.image_snapshot_ticks = 1000; train.network_snapshot_ticks = 1000\n",
        "    #desc += '-VERBOSE'; sched.tick_kimg_base = 1; sched.tick_kimg_dict = {}; train.image_snapshot_ticks = 1; train.network_snapshot_ticks = 100\n",
        "    #desc += '-GRAPH'; train.save_tf_graph = True\n",
        "    #desc += '-HIST'; train.save_weight_histograms = True\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "# Main entry point for training.\n",
        "# Calls the function indicated by 'train' using the selected options.\n",
        "\n",
        "def main():\n",
        "    kwargs = EasyDict(train)\n",
        "    kwargs.update(G_args=G, D_args=D, G_opt_args=G_opt, D_opt_args=D_opt, G_loss_args=G_loss, D_loss_args=D_loss)\n",
        "    kwargs.update(dataset_args=dataset, sched_args=sched, grid_args=grid, tf_config=tf_config)\n",
        "    kwargs.submit_config = copy.deepcopy(submit_config)\n",
        "    kwargs.submit_config.run_dir_root = dnnlib.submission.submit.get_template_from_path(config.result_dir)\n",
        "    kwargs.submit_config.run_dir_ignore += config.run_dir_ignore\n",
        "    kwargs.submit_config.run_desc = desc\n",
        "    dnnlib.submit_run(**kwargs)\n",
        "\n",
        "#----------------------------------------------------------------------------\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n",
        "\n",
        "#----------------------------------------------------------------------------\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R7HSZ3Y95zKb",
        "colab_type": "text"
      },
      "source": [
        "using trained model to generate images and interpolations video (slighty modfied from gene kogans fork here https://github.com/genekogan/stylegan/blob/master/generate.ipynb)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LuNalnk36QP7",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import pickle\n",
        "import numpy as np\n",
        "import math\n",
        "import random\n",
        "import matplotlib.pyplot as plt\n",
        "import PIL.Image\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "import config\n",
        "import moviepy.editor as mpy\n",
        "from tqdm import tqdm_notebook as tqdm\n",
        "\n",
        "tflib.init_tf()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0__wXmZg6XOa",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "model = '/content/stylegan/results/00000-sgan-custom-1gpu/network-snapshot-011185.pkl'\n",
        "\n",
        "with open(model, 'rb') as f:\n",
        "    _G, _D, Gs = pickle.load(f)\n",
        "\n",
        "fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n",
        "synthesis_kwargs = dict(output_transform=dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True), minibatch_size=8)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CAlhVax26eqM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "truncation = 0.7\n",
        "\n",
        "\n",
        "def bookmark(latents, new_faves):\n",
        "    for f in new_faves:\n",
        "        faves.append(latents[f])\n",
        "\n",
        "def show_faves(faves):\n",
        "    latents = np.array(faves)\n",
        "    labels = np.zeros([latents.shape[0]] + Gs.input_shapes[1][1:])\n",
        "    n = len(faves)\n",
        "    nr, nc = math.ceil(n / 6), 6\n",
        "    for r in range(nr):\n",
        "        images = Gs.run(latents[6*r:min(n-1, 6*(r+1))], None, truncation_psi=truncation, randomize_noise=False, output_transform=fmt)\n",
        "        img1 = np.concatenate([img for img in images], axis=1)\n",
        "        plt.figure(figsize=(24,4))\n",
        "        plt.imshow(img1)\n",
        "        \n",
        "def random_sample(num_images, scale):\n",
        "    latents = np.random.RandomState(int(1000*random.random())).randn(num_images, *Gs.input_shapes[0][1:])\n",
        "    labels = np.zeros([latents.shape[0]] + Gs.input_shapes[1][1:])\n",
        "    images = Gs.run(latents, None, truncation_psi=truncation, randomize_noise=False, output_transform=fmt)\n",
        "    images_ct = np.concatenate([img for img in images], axis=1)\n",
        "    plt.figure(figsize=(scale*num_images, scale))\n",
        "    plt.imshow(images_ct)\n",
        "    return images, latents\n",
        "\n",
        "def get_latent_interpolation(endpoints, num_frames_per, mode, shuffle):\n",
        "    if shuffle:\n",
        "        random.shuffle(endpoints)\n",
        "    num_endpoints, dim = len(endpoints), len(endpoints[0])\n",
        "    num_frames = num_frames_per * num_endpoints\n",
        "    endpoints = np.array(endpoints)\n",
        "    latents = np.zeros((num_frames, dim))\n",
        "    for e in range(num_endpoints):\n",
        "        e1, e2 = e, (e+1)%num_endpoints\n",
        "        for t in range(num_frames_per):\n",
        "            frame = e * num_frames_per + t\n",
        "            r = 0.5 - 0.5 * np.cos(np.pi*t/(num_frames_per-1)) if mode == 'ease' else float(t) / num_frames_per\n",
        "            latents[frame, :] = (1.0-r) * endpoints[e1,:] + r * endpoints[e2,:]\n",
        "    return latents\n",
        "\n",
        "def get_latent_interpolation_bspline(endpoints, nf, k, s, shuffle):\n",
        "    if shuffle:\n",
        "        random.shuffle(endpoints)\n",
        "    x = np.array(endpoints)\n",
        "    x = np.append(x, x[0,:].reshape(1, x.shape[1]), axis=0)\n",
        "    nd = x.shape[1]\n",
        "    latents = np.zeros((nd, nf))\n",
        "    nss = list(range(1, 10)) + [10]*(nd-19) + list(range(10,0,-1))\n",
        "    for i in tqdm(range(nd-9)):\n",
        "        idx = list(range(i,i+10))\n",
        "        tck, u = interpolate.splprep([x[:,j] for j in range(i,i+10)], k=k, s=s)\n",
        "        out = interpolate.splev(np.linspace(0, 1, num=nf, endpoint=True), tck)\n",
        "        latents[i:i+10,:] += np.array(out)\n",
        "    latents = latents / np.array(nss).reshape((512,1))\n",
        "    return latents.T\n",
        "\n",
        "\n",
        "def generate_images(latents, labels):\n",
        "    batch_size = 8\n",
        "    num_frames = latents.shape[0]\n",
        "    num_batches = int(np.ceil(num_frames/batch_size))\n",
        "    images = []\n",
        "    for b in tqdm(range(num_batches)):\n",
        "        new_images = Gs.run(latents[b*batch_size:min((b+1)*batch_size, num_frames-1), :], None, truncation_psi=truncation, randomize_noise=False, output_transform=fmt)\n",
        "        for img in new_images:\n",
        "            images.append(img)\n",
        "    return images\n",
        "\n",
        "def make_movie(images, out_dir, out_name):\n",
        "    temp_dir = 'frames%06d'%int(1000000*random.random())\n",
        "    os.system('mkdir %s'%temp_dir)\n",
        "    for idx in tqdm(range(len(images))):\n",
        "        PIL.Image.fromarray(images[idx], 'RGB').save('%s/frame%05d.png' % (temp_dir, idx))\n",
        "    cmd = 'ffmpeg -i %s/frame%05d.png -c:v libx264 -pix_fmt yuv420p %s/%s.mp4' % (temp_dir, out_dir, out_name)\n",
        "    print(cmd)\n",
        "    os.system(cmd)\n",
        "    os.system('rm -rf %s'%temp_dir)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7ajgYVtV6hns",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#from google.colab import files\n",
        "\n",
        "def random_sample(num_images, scale):\n",
        "    latents = np.random.RandomState(int(1000*random.random())).randn(num_images, *Gs.input_shapes[0][1:])\n",
        "    labels = np.zeros([latents.shape[0]] + Gs.input_shapes[1][1:])\n",
        "    images = Gs.run(latents, None, truncation_psi=truncation, randomize_noise=False, output_transform=fmt)\n",
        "    images_ct = np.concatenate([img for img in images], axis=1)\n",
        "    plt.figure(figsize=(scale*num_images, scale))\n",
        "    plt.imshow(images_ct)\n",
        "    plt.axis('off')\n",
        "    #plt.savefig('download.png')\n",
        "    return images, latents\n",
        "\n",
        "images, latents = random_sample(40, scale=8)\n",
        "\n",
        "\n",
        "#files.download('download.png')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ViRzB7aG6mn2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "latents = get_latent_interpolation(latents, 30, 'linear', False)\n",
        "labels = np.zeros([latents.shape[0]] + Gs.input_shapes[1][1:])\n",
        "\n",
        "images = generate_images(latents, labels)    \n",
        "\n",
        "make_movie(images, '.', 'faves13')"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}