SmallGan.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Copy of SmallGan.ipynb",
      "version": "0.3.2",
      "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/pytorch-small-dataset-image-generation/blob/master/SmallGan.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "metadata": {
        "id": "CW5drKMaJn7Z",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/apple2373/PyTorch-SmallGAN.git"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "CdA4HKbfKapJ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "cd PyTorch-SmallGAN/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "JFPejR4aKgon",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!gdown https://drive.google.com/uc?id=1nAle7FCVFZdix2--ks0r5JBkFnKw8ctW"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "QP5nbDOop5vA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!unzip BigGAN_ch96_bs256x8_138k.zip"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "SVXA-z7e0Z3Z",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!mv 138k/G_ema.pth data/"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Sdw84oTRzydN",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Training"
      ]
    },
    {
      "metadata": {
        "id": "qnprDx1XHtwi",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!wget https://www.dropbox.com/s/kzv6zo9jcmjcanj/checkpoint_iter20500.pth.tar"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "fFnT9bcv0sfP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "!python train.py --dataset anime --gpu 0 --pretrained ./data/G_ema.pth --resume \"/content/PyTorch-SmallGAN/checkpoint_iter20500.pth.tar\" --iters 50000"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "v9Dek_FYz0MT",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Testing"
      ]
    },
    {
      "metadata": {
        "id": "gOTSDxih-3Kd",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import glob\n",
        "import os\n",
        "import matplotlib\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import json\n",
        "\n",
        "%matplotlib inline\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "import torchvision\n",
        "from models.setup_model import setup_model\n",
        "from dataloaders.setup_dataloader_smallgan import setup_dataloader\n",
        "\n",
        "\n",
        "def reconstruct(model,out_path,indices):\n",
        "    model.eval()\n",
        "    device = next(model.parameters()).device\n",
        "    dataset_size = model.embeddings.weight.size()[0]\n",
        "    assert type(indices)==torch.Tensor\n",
        "    indices = indices.to(device)        \n",
        "    embeddings = model.embeddings(indices)\n",
        "    batch_size = embeddings.size()[0]\n",
        "    image_tensors = model(embeddings)\n",
        "    with torch.no_grad():\n",
        "        torchvision.utils.save_image(\n",
        "            image_tensors,\n",
        "            out_path,\n",
        "            nrow=int(batch_size ** 0.5),\n",
        "            normalize=True,\n",
        "        )\n",
        "        \n",
        "#see https://github.com/nogu-atsu/SmallGAN/blob/2293700dce1e2cd97e25148543532814659516bd/gen_models/ada_generator.py#L37-L53\n",
        "def interpolate(model,out_path,source,dist,trncate=0.4,num=5):\n",
        "    model.eval()\n",
        "    device = next(model.parameters()).device\n",
        "    dataset_size = model.embeddings.weight.size()[0]\n",
        "    indices = torch.tensor([source,dist],device=device)\n",
        "    indices = indices.to(device) \n",
        "    embeddings = model.embeddings(indices)\n",
        "    embeddings = embeddings[[0]] * torch.linspace(1, 0, num,device=device)[:, None] + embeddings[[1]]* torch.linspace(0, 1, num,device=device)[:, None]\n",
        "    batch_size = embeddings.size()[0]\n",
        "    image_tensors = model(embeddings)\n",
        "    with torch.no_grad():\n",
        "        torchvision.utils.save_image(\n",
        "            image_tensors,\n",
        "            out_path,\n",
        "            nrow=batch_size,\n",
        "            normalize=True,\n",
        "        )\n",
        "\n",
        "#from https://github.com/nogu-atsu/SmallGAN/blob/2293700dce1e2cd97e25148543532814659516bd/gen_models/ada_generator.py#L37-L53        \n",
        "def random(model,out_path,tmp=0.4, n=9, truncate=False):\n",
        "    from scipy.stats import truncnorm\n",
        "    model.eval()\n",
        "    device = next(model.parameters()).device\n",
        "    dataset_size = model.embeddings.weight.size()[0]\n",
        "    dim_z = model.embeddings.weight.size(1)\n",
        "    if truncate:\n",
        "        embeddings = truncnorm(-tmp, tmp).rvs(n * dim_z).astype(\"float32\").reshape(n, dim_z)\n",
        "    else:\n",
        "        embeddings = np.random.normal(0, tmp, size=(n, dim_z)).astype(\"float32\")\n",
        "    embeddings = torch.tensor(embeddings,device=device)\n",
        "    batch_size = embeddings.size()[0]\n",
        "    image_tensors = model(embeddings)\n",
        "    with torch.no_grad():\n",
        "        torchvision.utils.save_image(\n",
        "            image_tensors,\n",
        "            out_path,\n",
        "            nrow=int(batch_size ** 0.5),\n",
        "            normalize=True,\n",
        "        )"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "MiF96ABIAeEh",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "dataloader = setup_dataloader(\"anime\",batch_size=2)\n",
        "dataset_size = len(dataloader.dataset)\n",
        "exp_dir = \"./experiments/train_dataset-anime_model-biggan128-ada_2019-04-26-19-11-39/\"\n",
        "print(json.load(open(exp_dir+\"args.json\")))\n",
        "model = setup_model(\"biggan128-ada\",dataset_size=50,resume=exp_dir+\"checkpoint_iter500.pth.tar\")\n",
        "model = model.cuda()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "rs-4AHmcAvFt",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "reconstruct(model,out_path=\"./samples/anime_reconstruct.jpg\",indices= torch.arange(9))\n",
        "interpolate(model,out_path=\"./samples/anime_interpolate.jpg\",source=1,dist=2)\n",
        "random(model,out_path=\"./samples/anime_random.jpg\",tmp=0.2, n=9, truncate=True)\n",
        "\n",
        "plt.figure(figsize=(10,10))\n",
        "im = Image.open(\"./samples/anime_reconstruct.jpg\")\n",
        "plt.imshow(im)\n",
        "plt.show()\n",
        "\n",
        "plt.figure(figsize=(10,10))\n",
        "im = Image.open(\"./samples/anime_interpolate.jpg\")\n",
        "plt.imshow(im)\n",
        "plt.show()\n",
        "\n",
        "plt.figure(figsize=(10,10))\n",
        "im = Image.open(\"./samples/anime_random.jpg\")\n",
        "plt.imshow(im)\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "oAgxl015GsDz",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}