CS269.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CS269.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "mTG4skThPhgG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "SCALE =  1.0                   # scale for k e.g. 1.0\n",
        "FOLDER = \"269_svd_d_optimal_1\" # folder name on drive e.g. \"269_svd_d_optimal_5\""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "_baD78ErF6JY",
        "colab_type": "code",
        "outputId": "156ab598-c4ea-4946-b331-4ebf4cabeca2",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 136
        }
      },
      "cell_type": "code",
      "source": [
        "!pip install torch torchvision\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data import DataLoader\n",
        "from torchvision import datasets, transforms\n",
        "import torchvision.utils as vutils\n",
        "\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "dtype = torch.FloatTensor\n",
        "\n",
        "if torch.cuda.is_available():\n",
        "  dtype = torch.cuda.FloatTensor\n",
        "  \n",
        "print(dtype)"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: torch in /usr/local/lib/python3.6/dist-packages (1.0.1.post2)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.6/dist-packages (0.2.2.post3)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torchvision) (1.14.6)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision) (4.1.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from torchvision) (1.11.0)\n",
            "Requirement already satisfied: olefile in /usr/local/lib/python3.6/dist-packages (from pillow>=4.1.1->torchvision) (0.46)\n",
            "<class 'torch.cuda.FloatTensor'>\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "ommkPL67GBDQ",
        "colab_type": "code",
        "outputId": "e1df0128-e5b8-43a0-f873-ed283120f1d7",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 51
        }
      },
      "cell_type": "code",
      "source": [
        "image_size = 32\n",
        "batch_size = 125\n",
        "transform = transform=transforms.Compose([\n",
        "                               transforms.Resize(image_size),\n",
        "                               transforms.ToTensor(),\n",
        "                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))   #Normalizes to [-1, 1]\n",
        "                           ])\n",
        "\n",
        "trainset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
        "                                        download=True, transform=transform)\n",
        "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,\n",
        "                                          shuffle=True, num_workers=2)\n",
        "\n",
        "# testset = torchvision.datasets.CIFAR10(root='./data', train=False,\n",
        "#                                        download=True, transform=transform)\n",
        "# testloader = torch.utils.data.DataLoader(testset, batch_size=1000,\n",
        "#                                          shuffle=False, num_workers=2)\n",
        "\n",
        "classes = ('plane', 'car', 'bird', 'cat',\n",
        "           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\r0it [00:00, ?it/s]"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|█████████▉| 170287104/170498071 [01:10<00:00, 3643247.74it/s]"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "metadata": {
        "id": "GktlWXfwG_OZ",
        "colab_type": "code",
        "outputId": "8e3bf738-b393-4d0e-ed86-1c9e1dd5bbcc",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 187
        }
      },
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "%cd ..\n",
        "!rm -rf GAN-Stabilization/\n",
        "%cd -\n",
        "\n",
        "!git clone https://github.com/thearnavgarg/GAN-Stabilization.git\n",
        "%cd GAN-Stabilization"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/\n",
            "/content\n",
            "Cloning into 'GAN-Stabilization'...\n",
            "remote: Enumerating objects: 116, done.\u001b[K\n",
            "remote: Counting objects: 100% (116/116), done.\u001b[K\n",
            "remote: Compressing objects: 100% (87/87), done.\u001b[K\n",
            "remote: Total 116 (delta 60), reused 78 (delta 28), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (116/116), 41.05 MiB | 31.35 MiB/s, done.\n",
            "Resolving deltas: 100% (60/60), done.\n",
            "/content/GAN-Stabilization\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "BjLiuCo5NPCR",
        "colab_type": "code",
        "outputId": "6bd0f507-1336-4174-ef04-97586109ea8c",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 122
        }
      },
      "cell_type": "code",
      "source": [
        "from layers import SVDConv2d\n",
        "from gan import Generator, Discriminator, weights_init\n",
        "\n",
        "#Connect google drive for loading and saving weights\n",
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "\r170500096it [01:30, 3643247.74it/s]                               "
          ],
          "name": "stderr"
        },
        {
          "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&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&response_type=code\n",
            "\n",
            "Enter your authorization code:\n",
            "··········\n",
            "Mounted at /content/drive\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "metadata": {
        "id": "2XEcAZ4ZMtd_",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "nz = 100             # Number of latent z variables\n",
        "nc = 3               # Number of channels\n",
        "ngpu = 1\n",
        "# Generator\n",
        "ngf = 64\n",
        "netG = Generator(ngpu, nz, ngf, nc).type(dtype)\n",
        "netG.apply(weights_init)\n",
        "print(netG)\n",
        "\n",
        "# Discrimiator\n",
        "ndf = 64\n",
        "netD = Discriminator(ngpu, nz, ndf, nc, SCALE).type(dtype)\n",
        "print(netD)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "OMabxxGlRlYK",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "criterion = nn.BCELoss().type(dtype)\n",
        "beta1 = 0.5\n",
        "lr = 0.0002\n",
        "niter = 250\n",
        "\n",
        "fixed_noise = torch.randn(64, nz, 1, 1).type(dtype)\n",
        "real_label = 1\n",
        "fake_label = 0\n",
        "\n",
        "# setup optimizer\n",
        "optimizerD = optim.Adam(netD.parameters(), lr=lr, betas=(beta1, 0.999))\n",
        "optimizerG = optim.Adam(netG.parameters(), lr=lr, betas=(beta1, 0.999))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "NJMSLvi9Svy0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import time\n",
        "lbda = 0.1\n",
        "gamma = 10\n",
        "noise = torch.randn(batch_size, nz, 1, 1).type(dtype)\n",
        "img_list = []\n",
        "start = time.time()\n",
        "print(\"starting training at: \" + str(start))\n",
        "prev = time.time()\n",
        "with open('/content/drive/My Drive/' + FOLDER +'/log.txt', 'w') as f:\n",
        "    for epoch in range(niter):\n",
        "        curr = time.time()\n",
        "        print(\"*\"*10)\n",
        "        print(\"Time for last epoch: \" + str(curr-prev))\n",
        "        f.write(\"*\"*10 + \"\\n\")\n",
        "        f.write(\"Time for last epoch: \" + str(curr-prev) + \"\\n\")\n",
        "        prev = curr\n",
        "        for i, data in enumerate(trainloader, 0):\n",
        "            ############################\n",
        "            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n",
        "            ###########################\n",
        "            # train with real\n",
        "            netD.zero_grad()\n",
        "            real_cpu = data[0].type(dtype)\n",
        "            label = torch.full((batch_size,), real_label).type(dtype)\n",
        "\n",
        "            output = netD(real_cpu).view(-1)\n",
        "            errD_real = criterion(output, label)\n",
        "            errD_real += lbda*netD.orth_reg()\n",
        "            errD_real += gamma*netD.D_optimal_reg()\n",
        "            errD_real.backward()\n",
        "            D_x = output.mean().item()\n",
        "\n",
        "            # train with fake\n",
        "            noise = torch.randn(batch_size, nz, 1, 1).type(dtype)\n",
        "            fake = netG(noise)\n",
        "            label.fill_(fake_label)\n",
        "            output = netD(fake.detach()).view(-1)\n",
        "            errD_fake = criterion(output, label)\n",
        "            errD_fake.backward()\n",
        "\n",
        "            D_G_z1 = output.mean().item()\n",
        "            errD = errD_real + errD_fake\n",
        "            optimizerD.step()\n",
        "\n",
        "            ############################\n",
        "            # (2) Update G network: maximize log(D(G(z)))\n",
        "            ###########################\n",
        "            netG.zero_grad()\n",
        "            label.fill_(real_label)  # fake labels are real for generator cost\n",
        "            output = netD.forward(fake).view(-1)\n",
        "            errG = criterion(output, label)\n",
        "            errG.backward()\n",
        "            D_G_z2 = output.mean().item()\n",
        "            optimizerG.step()\n",
        "\n",
        "            if i % 50 == 0:\n",
        "                print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'\n",
        "                      % (epoch, niter, i, len(trainloader),\n",
        "                         errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n",
        "                f.write('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f\\n'\n",
        "                      % (epoch, niter, i, len(trainloader),\n",
        "                         errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n",
        "\n",
        "        torch.save(netG.state_dict(), '/content/drive/My Drive/' + FOLDER +'/G.pt')\n",
        "        torch.save(netD.state_dict(), '/content/drive/My Drive/' + FOLDER +'/D.pt')\n",
        "        torch.save(optimizerD.state_dict(), '/content/drive/My Drive/' + FOLDER +'/G_optim.pt')\n",
        "        torch.save(optimizerG.state_dict(), '/content/drive/My Drive/' + FOLDER +'/D_optim.pt')\n",
        "\n",
        "        if epoch % 5 == 0:\n",
        "            #Visualization of images\n",
        "            with torch.no_grad():\n",
        "                fake = netG(fixed_noise).detach().cpu()\n",
        "            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))\n",
        "            fake_images = np.transpose(img_list[-1],(1,2,0))\n",
        "            plt.imsave('/content/drive/My Drive/' + FOLDER +'/fake_image{}.png'.format(epoch), fake_images)\n",
        "    f.close()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "1qZ_RERNfjrv",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "with torch.no_grad():\n",
        "    fake = netG(fixed_noise).detach().cpu()\n",
        "img_list.append(vutils.make_grid(fake, padding=2, normalize=True))\n",
        "fake_images = np.transpose(img_list[-1],(1,2,0))\n",
        "plt.imsave('/content/drive/My Drive/' + FOLDER +'/fake_image{}.png'.format(epoch), fake_images)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "QZtQ54Nw_XCk",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}