Steam_Lightweight_GAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Steam-Lightweight-GAN.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Q7TuTbDXFlZ"
      },
      "source": [
        "# Steam-Lightweight-GAN\r\n",
        "\r\n",
        "Reference: https://github.com/lucidrains/lightweight-gan"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1XUm7SVvUMrA"
      },
      "source": [
        "## Check GPU specifications"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "f792_Dm-nfd5",
        "outputId": "961d862e-1e8e-4335-b98c-49cd615e6d8f"
      },
      "source": [
        "!nvidia-smi"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Tue Dec 29 09:10:46 2020       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.27.04    Driver Version: 418.67       CUDA Version: 10.1     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla T4            Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   55C    P8    10W /  70W |      0MiB / 15079MiB |      0%      Default |\n",
            "|                               |                      |                 ERR! |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CyoO8-BVVd3y"
      },
      "source": [
        "## Install requirements"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SFQfIcgRVQT8"
      },
      "source": [
        "%pip install google-colab-transfer lightweight-gan moviepy > /dev/null"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VmTyOV0NaBZk",
        "outputId": "d273a4b9-8b82-4710-d231-52addc34ffd7"
      },
      "source": [
        "import colab_transfer as ct\r\n",
        "\r\n",
        "ct.mount_google_drive()"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Mounted at /content/drive/\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DIOUESF2VOX3"
      },
      "source": [
        "## Load image data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oCkgPlhH_uCQ"
      },
      "source": [
        "archive_name = 'steam-oneface-small-lr.tar.gz'"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7V2Zf-5LAx4V"
      },
      "source": [
        "# Reference: https://github.com/woctezuma/steam-filtered-image-data#steam-oneface-dataset\r\n",
        "archive_dict = {\r\n",
        "    'steam-oneface-lr.tar.gz': '1QptHrW9vloTtP--YJsxMY8PZWI2D8NJt',\r\n",
        "    'steam-oneface-lr_with_retinaface.tar.gz': '1-0Nk7H6Cn3Nt60EdHG_NWSA8ohi2oBqr',\r\n",
        "    'steam-oneface-lr_with_dlib.tar.gz': '1-4RIn9G9Bee2JZ1bK1gkkgkLocHuWJJ4',\r\n",
        "    'steam-oneface-small-lr.tar.gz': '1-1V5fDhPo75iDtAbrD18rppV-lf51bPW',\r\n",
        "    'steam-oneface-tiny-lr.tar.gz': '1-2sCVgBUmu6LFug1pzBfmL8zNFFBq27F',\r\n",
        "}"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H4LjwtJQVFYt",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f54de04e-edca-48b1-b223-40b8180e6f70"
      },
      "source": [
        "from pathlib import Path\r\n",
        "\r\n",
        "gdrive_id = archive_dict[archive_name]\r\n",
        "print('Processing file {} with id = {}'.format(archive_name, gdrive_id))\r\n",
        "\r\n",
        "if not Path(archive_name).exists():\r\n",
        "  !gdown --id {gdrive_id}\r\n",
        "  !tar xf {archive_name}"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Processing file steam-oneface-small-lr.tar.gz with id = 1-1V5fDhPo75iDtAbrD18rppV-lf51bPW\n",
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1-1V5fDhPo75iDtAbrD18rppV-lf51bPW\n",
            "To: /content/steam-oneface-small-lr.tar.gz\n",
            "30.3MB [00:00, 115MB/s] \n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DmuQFPAFVnlL"
      },
      "source": [
        "## Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2_a6ubirBfDB"
      },
      "source": [
        "### Parameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fHuGlq85A6UG"
      },
      "source": [
        "base_name = archive_name.replace('.tar.gz', '')"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N8yo8yF4WImb",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "f4591927-7765-457f-901a-3f05eb412c48"
      },
      "source": [
        "image_dir = '/content/{}/'.format(base_name)\r\n",
        "\r\n",
        "print('Image directory: {}'.format(image_dir))"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Image directory: /content/steam-oneface-small-lr/\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-LBO5KeqBBjl",
        "outputId": "1312d66f-5133-460f-b56a-7cbd5e5937b7"
      },
      "source": [
        "run_name = 'run-{}'.format(base_name)\r\n",
        "\r\n",
        "suffixe = '-v0.14.4'\r\n",
        "run_name += suffixe\r\n",
        "\r\n",
        "print('Run name: {}'.format(run_name))"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Run name: run-steam-oneface-small-lr\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RQxVbjtsnGxi"
      },
      "source": [
        "model_dir = '/content/drive/MyDrive/lightweight-gan/'\r\n",
        "%mkdir -p {model_dir}"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qnAdjbMSBiAM"
      },
      "source": [
        "### Run"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sfcz610tBrM9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "0dbe53c5-aee2-464a-9168-035ffac1c374"
      },
      "source": [
        "# An arbitrary augmentation plan:\r\n",
        "# -   iter   0k -  27k   [translation] with p = 0.4\r\n",
        "# -   iter  27k -  54k   [translation,color] with p = 0.4\r\n",
        "# -   iter  54k -  71k   [translation,color,cutout] with p = 0.4\r\n",
        "# -   iter  71k -  78k   [translation,color,cutout] with p = 0.6\r\n",
        "\r\n",
        "mini_batch_size = 64\r\n",
        "\r\n",
        "# mini-batch:\r\n",
        "# -  32 for Tesla  K80 GPU,\r\n",
        "# -  64 for Tesla   T4 GPU,\r\n",
        "# - 128 for Tesla P100 GPU.\r\n",
        "\r\n",
        "!lightweight_gan \\\r\n",
        " --data {image_dir} \\\r\n",
        " --name {run_name} \\\r\n",
        " --models-dir {model_dir} \\\r\n",
        " --amp \\\r\n",
        " --aug-prob 0.4 \\\r\n",
        " --aug-types [translation,color,cutout] \\\r\n",
        " --batch-size {mini_batch_size} \\\r\n",
        " --gradient-accumulate-every 1 \\\r\n",
        " --save-every 1000 \\ \r\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "continuing from previous epoch - 62\n",
            "loading from version 0.14.4\n",
            "run-steam-oneface-small-lr</content/steam-oneface-small-lr/>:  41% 62000/150000 [00:00<?, ?it/s]G: 3.35 | D: 0.18 | GP: 0.54 | SS: 0.02\n",
            "run-steam-oneface-small-lr</content/steam-oneface-small-lr/>:  41% 62048/150000 [01:15<38:54:30,  1.59s/it]G: 1.16 | D: 0.28 | GP: 0.41 | SS: 0.02\n",
            "run-steam-oneface-small-lr</content/steam-oneface-small-lr/>:  41% 62097/150000 [02:34<39:12:56,  1.61s/it]G: 3.52 | D: 0.03 | GP: 0.44 | SS: 0.01\n",
            "run-steam-oneface-small-lr</content/steam-oneface-small-lr/>:  41% 62104/150000 [02:44<38:11:29,  1.56s/it]"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FAB6OkoHXe-D"
      },
      "source": [
        "## Sample"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OpBW8uM-Wlir",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "902ef7f4-f591-48f8-ac74-9cab081a7c28"
      },
      "source": [
        "!lightweight_gan --generate \\\r\n",
        " --num-image-tiles 3 \\\r\n",
        " --name {run_name} \\\r\n",
        " --models-dir {model_dir} \\\r\n"
      ],
      "execution_count": 16,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "continuing from previous epoch - 62\n",
            "loading from version 0.14.4\n",
            "sample images generated at ./results/run-steam-oneface-small-lr/generated-12-29-2020_11-31-56\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CBcrUISYqlGM"
      },
      "source": [
        "## Interpolate"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-kgq4yijWmsr",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "a4c5b6fc-400e-40b5-dd2c-8426ad5e3da2"
      },
      "source": [
        "!lightweight_gan --generate-interpolation \\\r\n",
        " --num-image-tiles 3 \\\r\n",
        " --name {run_name} \\\r\n",
        " --models-dir {model_dir} \\\r\n"
      ],
      "execution_count": 17,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "continuing from previous epoch - 62\n",
            "loading from version 0.14.4\n",
            "100% 100/100 [00:46<00:00,  2.14it/s]\n",
            "interpolation generated at ./results/run-steam-oneface-small-lr/generated-12-29-2020_11-32-13\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HXl4xixFVSXu"
      },
      "source": [
        "import glob\r\n",
        "\r\n",
        "input_path = 'results/{}/'.format(run_name)\r\n",
        "gif_files = glob.glob(input_path + '*.gif')"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9kpPQfHipWU9",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8aee8d85-d5f8-4818-e59b-d8faab9bf326"
      },
      "source": [
        "import moviepy.editor as mp\r\n",
        "\r\n",
        "input_name = sorted(gif_files)[-1]\r\n",
        "output_name = input_name.replace('.gif', '.mp4')\r\n",
        "\r\n",
        "clip = mp.VideoFileClip(input_name)\r\n",
        "clip.write_videofile(output_name)"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[MoviePy] >>>> Building video results/run-steam-oneface-small-lr/generated-12-29-2020_11-32-13.mp4\n",
            "[MoviePy] Writing video results/run-steam-oneface-small-lr/generated-12-29-2020_11-32-13.mp4\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|█████████▉| 793/794 [02:40<00:00,  4.94it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "[MoviePy] Done.\n",
            "[MoviePy] >>>> Video ready: results/run-steam-oneface-small-lr/generated-12-29-2020_11-32-13.mp4 \n",
            "\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "crIGnD_xRSPX"
      },
      "source": [
        "Done."
      ]
    }
  ]
}