BiGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "BiGAN.ipynb",
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "metadata": {
        "id": "kcQdGgP4WBhA",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "# Task 1 - Generative Models"
      ]
    },
    {
      "metadata": {
        "id": "s6KOURPRVwAE",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Introduction"
      ]
    },
    {
      "metadata": {
        "id": "B139n3lXfbqb",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "This is a research environment for Task 1 - Generative Models. The implementations of the networks are adapted from William Beng's implementation of the BiGAN network.\n",
        "[[Paper]](https://arxiv.org/abs/1605.09782) [[Original implementation]](https://github.com/WilliBee/bigan_SRL)\n",
        "\n",
        "Additionally, this relies on Cyrille Rossant's post 'An illustrated introduction to the t-SNE algorithm' for t-SNE visualisation utilities. [[Post]](https://www.oreilly.com/learning/an-illustrated-introduction-to-the-t-sne-algorithm)\n",
        "\n",
        "Triplet data processing and triplet loss calculations based on Adam Bielski's implementation of triplet networks. [[Implementation]](https://github.com/adambielski/siamese-triplet/blob/master/losses.py)\n",
        "\n",
        "The BiCoGAN architecture comes from a paper by Ayush Jaiswal, Wael AbdAlmageed, Yue Wu, Premkumar Natarajan. [[Paper]](https://arxiv.org/abs/1711.07461)\n",
        "\n",
        "The overriding principle has been to prototype rapidly and get research results as quickly as possible. This means that there are some obvious inefficiencies in the code, e.g. network error printing could be rewritten as a seperate function and reused as needed.\n",
        "\n",
        "Keeping different architectures in separate classes is a conscious design decision.\n",
        "\n",
        "Thanks to the Colab environment, further architectures and variations of the mentioned ones were tested. In many cases, the results were either similar to the ones achieved on the BiCoGAN architecture or were not substantially different from the baseline method (BiGAN) and were not reported.\n",
        "\n",
        "Let's get the ball rolling."
      ]
    },
    {
      "metadata": {
        "id": "LeJ8AjerVjCf",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Preliminaries"
      ]
    },
    {
      "metadata": {
        "id": "xGQu1sKYftH9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "from torch.autograd import Variable\n",
        "import torch.optim as optim\n",
        "from torch.utils.data import Dataset\n",
        "from torchvision import datasets, transforms\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.gridspec as gridspec\n",
        "import matplotlib.patheffects as PathEffects\n",
        "import seaborn as sns\n",
        "import numpy as np\n",
        "import os\n",
        "from itertools import *\n",
        "import math\n",
        "from PIL import Image\n",
        "from sklearn.manifold import TSNE\n",
        "# Faster implementation of t-SNE:\n",
        "# from MulticoreTSNE import MulticoreTSNE as TSNE\n",
        "from sklearn.neighbors import KNeighborsClassifier\n",
        "from sklearn import metrics\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Fr4yfdkyvr0_",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#from google.colab import drive\n",
        "#drive.mount('/content/gdrive')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "VVsWHW6yVrs6",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Shared Utilities"
      ]
    },
    {
      "metadata": {
        "id": "7XKlQ4IJlcy9",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# print out net and number of parameters\n",
        "\n",
        "def print_network(net):\n",
        "    num_params = 0\n",
        "    for param in net.parameters():\n",
        "        num_params += param.numel()\n",
        "    print(net)\n",
        "    print('Total number of parameters: %d' % num_params)\n",
        "    \n",
        "            "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "l3Nsn7OjWUiG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# weight initialization\n",
        "\n",
        "def initialize_weights(net):\n",
        "    for m in net.modules():\n",
        "        if isinstance(m, nn.Linear):\n",
        "            m.weight.data.normal_(0, 0.02)\n",
        "            m.bias.data.zero_()\n",
        "\n",
        "            "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ROPpDlEGXwNt",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# handles log cases with x = 0\n",
        "\n",
        "def log(x):\n",
        "      return torch.log(x + 1e-8)\n",
        "\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "EdNHdpAncX85",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# scatter plot for t-SNE\n",
        "\n",
        "def scatter(x, colors):\n",
        "    # We choose a color palette with seaborn.\n",
        "    palette = np.array(sns.color_palette(\"hls\", 10))\n",
        "\n",
        "    # We create a scatter plot.\n",
        "    f = plt.figure(figsize=(8, 8))\n",
        "    ax = plt.subplot(aspect='equal')\n",
        "    sc = ax.scatter(x[:,0], x[:,1], lw=0, s=40,\n",
        "                    c=palette[colors.astype(np.int)])\n",
        "    plt.xlim(-25, 25)\n",
        "    plt.ylim(-25, 25)\n",
        "    ax.axis('off')\n",
        "    ax.axis('tight')\n",
        "\n",
        "    # We add the labels for each digit.\n",
        "    txts = []\n",
        "    for i in range(10):\n",
        "        # Position of each label.\n",
        "        xtext, ytext = np.median(x[colors == i, :], axis=0)\n",
        "        txt = ax.text(xtext, ytext, str(i), fontsize=24)\n",
        "        txt.set_path_effects([\n",
        "            PathEffects.Stroke(linewidth=5, foreground=\"w\"),\n",
        "            PathEffects.Normal()])\n",
        "        txts.append(txt)\n",
        "\n",
        "    return f, ax, sc, txts"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "3MdTVNdXaMW1",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#!rm -rf sample_data\n",
        "#!zip -r /content/file.zip /content\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Wps2cNnRYD7Z",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Datasets"
      ]
    },
    {
      "metadata": {
        "id": "vra803YZXwio",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# classic MNIST\n",
        "\n",
        "class Mnist:\n",
        "    def __init__(self, batch_size):\n",
        "        MNIST_MEAN = 0.1307\n",
        "        MNIST_STD = 0.3081\n",
        "\n",
        "        dataset_transform = transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       # transforms.Normalize((MNIST_MEAN,), (MNIST_STD,))\n",
        "                   ])\n",
        "\n",
        "        train_dataset = datasets.MNIST('/home/zelazny/Downloads/tooploox/data/', train=True, download=True, transform=dataset_transform)\n",
        "        test_dataset = datasets.MNIST('/home/zelazny/Downloads/tooploox/data/', train=False, download=True, transform=dataset_transform)\n",
        "\n",
        "        self.train_loader  = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
        "        self.test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)\n",
        "\n",
        "        "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "SWGKw7fqiL0F",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# Dataset to transform MNIST into triplets\n",
        "\n",
        "class TripletMNIST(Dataset):\n",
        "    \"\"\"\n",
        "    Train: For each sample (anchor) randomly chooses a positive and negative samples\n",
        "    Test: Creates fixed triplets for testing\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, mnist_dataset):\n",
        "        self.mnist_dataset = mnist_dataset\n",
        "        self.train = self.mnist_dataset.train\n",
        "        self.transform = self.mnist_dataset.transform\n",
        "\n",
        "        if self.train:\n",
        "            self.targets= self.mnist_dataset.targets\n",
        "            self.data = self.mnist_dataset.data\n",
        "            self.labels_set = set(self.targets.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.targets.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "        else:\n",
        "            self.targets = self.mnist_dataset.targets\n",
        "            self.data = self.mnist_dataset.data\n",
        "            # generate fixed triplets for testing\n",
        "            self.labels_set = set(self.targets.numpy())\n",
        "            self.label_to_indices = {label: np.where(self.targets.numpy() == label)[0]\n",
        "                                     for label in self.labels_set}\n",
        "\n",
        "            random_state = np.random.RandomState(29)\n",
        "\n",
        "            triplets = [[i,\n",
        "                         random_state.choice(self.label_to_indices[self.targets[i].item()]),\n",
        "                         random_state.choice(self.label_to_indices[\n",
        "                                                 np.random.choice(\n",
        "                                                     list(self.labels_set - set([self.targets[i].item()]))\n",
        "                                                 )\n",
        "                                             ])\n",
        "                         ]\n",
        "                        for i in range(len(self.data))]\n",
        "            self.test_triplets = triplets\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        if self.train:\n",
        "            img1, label1 = self.data[index], self.targets[index].item()\n",
        "            positive_index = index\n",
        "            while positive_index == index:\n",
        "                positive_index = np.random.choice(self.label_to_indices[label1])\n",
        "            negative_label = np.random.choice(list(self.labels_set - set([label1])))\n",
        "            negative_index = np.random.choice(self.label_to_indices[negative_label])\n",
        "            img2 = self.data[positive_index]\n",
        "            img3 = self.data[negative_index]\n",
        "        else:\n",
        "            img1 = self.data[self.test_triplets[index][0]]\n",
        "            img2 = self.data[self.test_triplets[index][1]]\n",
        "            img3 = self.data[self.test_triplets[index][2]]\n",
        "            label1 = self.targets[self.test_triplets[index][0]]\n",
        "\n",
        "        img1 = Image.fromarray(img1.numpy(), mode='L')\n",
        "        img2 = Image.fromarray(img2.numpy(), mode='L')\n",
        "        img3 = Image.fromarray(img3.numpy(), mode='L')\n",
        "        if self.transform is not None:\n",
        "            img1 = self.transform(img1)\n",
        "            img2 = self.transform(img2)\n",
        "            img3 = self.transform(img3)\n",
        "        #return (img1, img2, img3), []\n",
        "        return (img1, img2, img3), label1\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.mnist_dataset)\n",
        "\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "VSgc6_9OicF-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# MNIST transformed into triplets\n",
        "\n",
        "class TripletMnist:\n",
        "    def __init__(self, batch_size):\n",
        "        MNIST_MEAN = 0.1307\n",
        "        MNIST_STD = 0.3081\n",
        "\n",
        "        dataset_transform = transforms.Compose([\n",
        "                       transforms.ToTensor(),\n",
        "                       # transforms.Normalize((MNIST_MEAN,), (MNIST_STD,))\n",
        "                   ])\n",
        "\n",
        "        train_dataset = datasets.MNIST('/home/zelazny/Downloads/tooploox/data/', train=True, download=True, transform=dataset_transform)\n",
        "        test_dataset = datasets.MNIST('/home/zelazny/Downloads/tooploox/data/', train=False, download=True, transform=dataset_transform)\n",
        "        \n",
        "        triplet_train_dataset = TripletMNIST(train_dataset)\n",
        "        triplet_test_dataset = TripletMNIST(test_dataset)\n",
        "\n",
        "        self.train_loader  = torch.utils.data.DataLoader(triplet_train_dataset, batch_size=batch_size, shuffle=True)\n",
        "        self.test_loader = torch.utils.data.DataLoader(triplet_test_dataset, batch_size=batch_size, shuffle=False)\n",
        "\n",
        "        "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "5mQpw8v8WwB7",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## BiGAN Architecture & Utilities"
      ]
    },
    {
      "metadata": {
        "id": "8KsDPbxpWWPC",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# fully-connected generator\n",
        "\n",
        "\n",
        "class Generator_FC(nn.Module):\n",
        "    \"\"\"\n",
        "    MLP with two hidden layers of dimension h_dim.\n",
        "    Input is a vector from representation space of dimension z_dim.\n",
        "    Output is a vector from image space of dimension X_dim.\n",
        "    \"\"\"\n",
        "    def __init__(self, z_dim, h_dim, X_dim):\n",
        "        super(Generator_FC, self).__init__()\n",
        "\n",
        "        self.z_dim = z_dim\n",
        "        self.h_dim = h_dim\n",
        "        self.X_dim = X_dim\n",
        "\n",
        "        self.fc = torch.nn.Sequential(\n",
        "            torch.nn.Linear(z_dim, h_dim),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.Linear(h_dim, h_dim),\n",
        "            torch.nn.BatchNorm1d(h_dim),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.Linear(h_dim, X_dim),\n",
        "            torch.nn.Sigmoid()\n",
        "            )\n",
        "\n",
        "        initialize_weights(self)\n",
        "\n",
        "    def forward(self, input):\n",
        "        x = self.fc(input)\n",
        "        return x\n",
        "\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "sNcxLxQkI0D4",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "The Discriminator is implemented as outlined in the BiGan paper:\n",
        "\n",
        "*The BiGAN\n",
        "discriminator D(x, z) takes data x as its initial input, and at each linear layer thereafter, the latent\n",
        "representation z is transformed using a learned linear transformation to the hidden layer dimension\n",
        "and added to the non-linearity input.*\n",
        "\n",
        "Such a formulation is not standard. Alternative implementations based on concatenations of X and z are possible and perform similarly.\n",
        "\n",
        "This Discriminator is also used in the basic variants of the BiCoGAN and TriBiGAN architectures."
      ]
    },
    {
      "metadata": {
        "id": "kj3HXc7NWfgG",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# fully-connected Discriminator\n",
        "\n",
        "\n",
        "class Discriminator_FC(nn.Module):\n",
        "    \"\"\"\n",
        "    MLP with two hidden layers of dimension h_dim.\n",
        "    Initial input is vector X from data space (or sampled from representation space).\n",
        "    Vector z from latent space (representation space) is transformed linearly\n",
        "    and added to activation input in hidden layers.\n",
        "    For example, if X comes from the dataset, corresponding\n",
        "    z is Encoder(X), and if z is sampled from representation space, X is Generator(z).\n",
        "    \"\"\"\n",
        "    def __init__(self, z_dim, h_dim, X_dim):\n",
        "        super(Discriminator_FC, self).__init__()\n",
        "\n",
        "        self.z_dim = z_dim\n",
        "        self.h_dim = h_dim\n",
        "        self.X_dim = X_dim\n",
        "\n",
        "        self.fc1 = torch.nn.Linear(X_dim, h_dim)\n",
        "        \n",
        "        self.z1 = torch.nn.Linear(z_dim, h_dim)   \n",
        "        \n",
        "        self.leaky1 = nn.LeakyReLU(0.2)\n",
        "        \n",
        "        self.fc2 = torch.nn.Linear(h_dim, h_dim)\n",
        "        \n",
        "        self.z2 = torch.nn.Linear(z_dim, h_dim)\n",
        "        \n",
        "        self.leaky2 = nn.LeakyReLU(0.2)\n",
        "        \n",
        "        self.fc3 = torch.nn.Sequential(\n",
        "            torch.nn.Linear(h_dim, 1),\n",
        "            torch.nn.Sigmoid()\n",
        "            )\n",
        "\n",
        "        initialize_weights(self)\n",
        "\n",
        "    def forward(self, input_x, input_z):\n",
        "        x = self.fc1(input_x)\n",
        "        x = torch.add(x, self.z1(input_z))\n",
        "        x = self.leaky1(x)\n",
        "        x = self.fc2(x)\n",
        "        x = torch.add(x, self.z2(input_z))\n",
        "        x = self.leaky2(x)\n",
        "        return self.fc3(x)\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "oRTF1CHqWy7u",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# fully-connected Encoder\n",
        "\n",
        "class Encoder_FC(nn.Module):\n",
        "    \"\"\"\n",
        "    MLP with with two hidden layers of dimension h_dim.\n",
        "    Input is vector X from image space if dimension X_dim.\n",
        "    Output is vector z from representation space of dimension z_dim.\n",
        "    \"\"\"\n",
        "    def __init__(self, z_dim, h_dim, X_dim):\n",
        "        super(Encoder_FC, self).__init__()\n",
        "\n",
        "        self.z_dim = z_dim\n",
        "        self.h_dim = h_dim\n",
        "        self.X_dim = X_dim\n",
        "        \n",
        "        self.fc = torch.nn.Sequential(\n",
        "            torch.nn.Linear(X_dim, h_dim),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            torch.nn.Linear(h_dim, h_dim),\n",
        "            torch.nn.BatchNorm1d(h_dim),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            torch.nn.Linear(h_dim, z_dim),\n",
        "            )\n",
        "        \n",
        "        initialize_weights(self)\n",
        "\n",
        "    def forward(self, input):\n",
        "        x = self.fc(input)\n",
        "        return x\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "tZ0KzsjMM_CT",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# save plots of loss functions after training\n",
        "\n",
        "\n",
        "def save_plot_losses(train_D_loss, train_G_loss, eval_D_loss, eval_G_loss, model_used, z_dim, epochs, lr, batch_size):\n",
        "\n",
        "    x = np.arange(1, len(train_D_loss) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_G_loss, label=\"Train G loss\", linewidth=2)\n",
        "    plt.plot(x, eval_G_loss, label=\"Eval G loss\", linewidth=2)\n",
        "\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the Train and Eval losses of G\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_G_losses.eps\", format='eps', dpi=1000)\n",
        "\n",
        "\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_D_loss, label=\"Train D loss\", linewidth=2)\n",
        "    plt.plot(x, eval_D_loss, label=\"Eval D loss\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "\n",
        "    plt.suptitle(\"Evolution of the Train and Eval losses of D\")\n",
        "\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_D_losses.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n",
        "\n",
        "\n",
        "def save_plot_pixel_norm(mean_pixel_norm, model_used, z_dim, epochs, lr, batch_size):\n",
        "    x = np.arange(1, len(mean_pixel_norm) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, mean_pixel_norm, label=\"Reconstruction error\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Norm')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the reconstruction error between X and G(E(X))\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"pix2pix_norm.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n",
        "\n",
        "def save_plot_z_norm(mean_z_norm, model_used, z_dim, epochs, lr, batch_size):\n",
        "    x = np.arange(1, len(mean_z_norm) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, mean_z_norm, label=\"Reconstruction error\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Norm')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the reconstruction error between z and E(G(z))\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"z_norm.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "aOSi6z9Bbg1D",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# BiGAN net\n",
        "\n",
        "class BIGAN(object):\n",
        "    \"\"\"\n",
        "    Class implementing a BiGAN network that trains from an observations dataset.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, kwargs):\n",
        "        \n",
        "        self.epoch = kwargs['epoch']\n",
        "        self.batch_size = kwargs['batch_size']\n",
        "        self.save_dir = kwargs['save_dir']\n",
        "        self.result_dir = kwargs['result_dir']\n",
        "        self.log_dir = kwargs['log_dir']\n",
        "        self.gpu_mode = kwargs['gpu_mode']\n",
        "        self.learning_rate = kwargs['lr']\n",
        "        self.lr_decay = kwargs['lr_decay']\n",
        "        self.beta1 = kwargs['beta1']\n",
        "        self.beta2 = kwargs['beta2']\n",
        "        self.decay = kwargs['decay']\n",
        "        self.network_type = kwargs['network_type']\n",
        "        self.dataset = kwargs['dataset']\n",
        "        self.dataset_path = kwargs['dataset_path']\n",
        "\n",
        "        # BIGAN parameters\n",
        "        self.z_dim = kwargs['z_dim']    #dimension of feature space\n",
        "        self.h_dim = kwargs['h_dim']    #dimension of the hidden layer\n",
        "\n",
        "        if kwargs['dataset'] == 'mnist':\n",
        "            self.X_dim = 28*28    #dimension of data\n",
        "            self.num_channels = 1\n",
        "\n",
        "        if kwargs['network_type'] == 'FC':\n",
        "            # networks init\n",
        "            self.G = Generator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.D = Discriminator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.E = Encoder_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "        else:\n",
        "            raise Exception(\"[!] There is no option for \" + kwargs['network_type'])\n",
        "\n",
        "        if self.gpu_mode:\n",
        "            self.G.cuda()\n",
        "            self.D.cuda()\n",
        "            self.E.cuda()\n",
        "\n",
        "        \n",
        "        self.G_solver = optim.Adam(chain(self.E.parameters(), self.G.parameters()), lr=self.learning_rate, betas=[self.beta1,self.beta2], weight_decay=self.decay)\n",
        "        self.D_solver = optim.Adam(self.D.parameters(), lr=self.learning_rate, betas=[self.beta1,self.beta2], weight_decay=self.decay)\n",
        "        \n",
        "        # exponential learning rate decay starting halfway through learning\n",
        "        self.G_scheduler = optim.lr_scheduler.MultiStepLR(self.G_solver, milestones=list(range(self.epoch//2, self.epoch)), gamma=self.lr_decay)\n",
        "        self.D_scheduler = optim.lr_scheduler.MultiStepLR(self.D_solver, milestones=list(range(self.epoch//2, self.epoch)), gamma=self.lr_decay)\n",
        "\n",
        "\n",
        "\n",
        "        print('---------- Networks architecture -------------')\n",
        "        print_network(self.G)\n",
        "        print_network(self.E)\n",
        "        print_network(self.D)\n",
        "        print('-----------------------------------------------')\n",
        "\n",
        "\n",
        "\n",
        "    def D_(self, X, z):\n",
        "        return self.D(X, z)\n",
        "\n",
        "    def reset_grad(self):\n",
        "        self.E.zero_grad()\n",
        "        self.G.zero_grad()\n",
        "        self.D.zero_grad()\n",
        "\n",
        "\n",
        "    def train(self):\n",
        "        if self.dataset == 'mnist':\n",
        "            dataset = Mnist(self.batch_size)\n",
        "\n",
        "\n",
        "        self.train_hist = {}\n",
        "        self.train_hist['D_loss'] = []\n",
        "        self.train_hist['G_loss'] = []\n",
        "\n",
        "        self.eval_hist = {}\n",
        "        self.eval_hist['D_loss'] = []\n",
        "        self.eval_hist['G_loss'] = []\n",
        "        self.eval_hist['pixel_norm'] = []\n",
        "        self.eval_hist['z_norm'] = []\n",
        "\n",
        "\n",
        "        for epoch in range(self.epoch):\n",
        "            print(\"epoch \",str(epoch))\n",
        "\n",
        "            self.D.train()\n",
        "            self.E.train()\n",
        "            self.G.train()\n",
        "\n",
        "            train_loss_G = 0\n",
        "            train_loss_D = 0\n",
        "            \n",
        "            # learning rate schedule\n",
        "            self.G_scheduler.step()\n",
        "            self.D_scheduler.step()\n",
        "\n",
        "\n",
        "            for batch_id, (data, target) in enumerate(dataset.train_loader):\n",
        "\n",
        "                if self.gpu_mode:\n",
        "                    # sample z\n",
        "                    z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1)).cuda() # uniform on [-1,1]\n",
        "                    # X is a real image from the dataset\n",
        "                    X = data\n",
        "                    X = Variable(X).cuda()\n",
        "                else:\n",
        "                    z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1)) # uniform on [-1,1]\n",
        "                    X = data\n",
        "                    X = Variable(X)\n",
        "\n",
        "                # sometimes batchsize of X is not equal to actual batch_size\n",
        "                if X.size(0) == self.batch_size:\n",
        "\n",
        "                    if self.network_type == 'FC':\n",
        "                        X = X.view(self.batch_size, -1)\n",
        "                        z_hat = self.E(X)\n",
        "                        X_hat = self.G(z)\n",
        "\n",
        "                        D_enc = self.D_(X, z_hat)\n",
        "                        D_gen = self.D_(X_hat, z)\n",
        "\n",
        "\n",
        "                    D_loss = -torch.mean(log(D_enc) + log(1 - D_gen))\n",
        "                    G_loss = -torch.mean(log(D_gen) + log(1 - D_enc))\n",
        "\n",
        "                    D_loss.backward(retain_graph=True)\n",
        "                    self.D_solver.step()\n",
        "                    self.reset_grad()\n",
        "\n",
        "                    G_loss.backward()\n",
        "                    self.G_solver.step()\n",
        "                    self.reset_grad()\n",
        "\n",
        "                    train_loss_G += G_loss.data\n",
        "                    train_loss_D += D_loss.data\n",
        "\n",
        "                    if batch_id % 1000 == 0:\n",
        "                        # Print and plot every now and then\n",
        "                        samples = X_hat.data.cpu().numpy()\n",
        "\n",
        "                        fig = plt.figure(figsize=(8, 4))\n",
        "                        gs = gridspec.GridSpec(4, 8)\n",
        "                        gs.update(wspace=0.05, hspace=0.05)\n",
        "\n",
        "                        for i, sample in enumerate(samples):\n",
        "                            if i<32:\n",
        "                                ax = plt.subplot(gs[i])\n",
        "                                plt.axis('off')\n",
        "                                ax.set_xticklabels([])\n",
        "                                ax.set_yticklabels([])\n",
        "                                ax.set_aspect('equal')\n",
        "\n",
        "                                if self.network_type == 'FC':\n",
        "                                    if self.dataset == 'mnist':\n",
        "                                        sample = sample.reshape(28, 28)\n",
        "                                        plt.imshow(sample, cmap='Greys_r')\n",
        "\n",
        "\n",
        "                        if not os.path.exists(self.result_dir + '/train/'):\n",
        "                            os.makedirs(self.result_dir + '/train/')\n",
        "\n",
        "                        filename = \"epoch_\" + str(epoch) + \"_batchid_\" + str(batch_id)\n",
        "                        plt.savefig(self.result_dir + '/train/{}.png'.format(filename, bbox_inches='tight'))\n",
        "                        plt.close()\n",
        "\n",
        "            print(\"Train loss G:\", train_loss_G / len(dataset.train_loader))\n",
        "            print(\"Train loss D:\", train_loss_D / len(dataset.train_loader))\n",
        "\n",
        "            self.train_hist['D_loss'].append(train_loss_D / len(dataset.train_loader))\n",
        "            self.train_hist['G_loss'].append(train_loss_G / len(dataset.train_loader))\n",
        "\n",
        "\n",
        "            self.D.eval()\n",
        "            self.E.eval()\n",
        "            self.G.eval()\n",
        "            test_loss_G = 0\n",
        "            test_loss_D = 0\n",
        "\n",
        "            mean_pixel_norm = 0\n",
        "            mean_z_norm = 0\n",
        "            norm_counter = 1\n",
        "\n",
        "            for batch_id, (data, target) in enumerate(dataset.test_loader):\n",
        "                # Sample data\n",
        "                z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1)) # uniform on [-1,1]\n",
        "                X_data = Variable(data)\n",
        "\n",
        "                if self.gpu_mode:\n",
        "                    z = z.cuda()\n",
        "                    X_data = X_data.cuda()\n",
        "\n",
        "                if X_data.size(0) == self.batch_size:\n",
        "                    X = X_data\n",
        "                    if self.network_type == 'FC':\n",
        "                        X = X.view(self.batch_size, -1)\n",
        "                        z_hat = self.E(X)\n",
        "                        X_hat = self.G(z)\n",
        "\n",
        "                        D_enc = self.D_(X, z_hat)\n",
        "                        D_gen = self.D_(X_hat, z)\n",
        "\n",
        "                    D_loss = -torch.mean(log(D_enc) + log(1 - D_gen))\n",
        "                    G_loss = -torch.mean(log(D_gen) + log(1 - D_enc))\n",
        "\n",
        "                    test_loss_G += G_loss.data\n",
        "                    test_loss_D += D_loss.data\n",
        "\n",
        "                    pixel_norm = X -  self.G(z_hat)\n",
        "                    pixel_norm = pixel_norm.norm().data / float(self.X_dim)\n",
        "                    mean_pixel_norm += pixel_norm\n",
        "\n",
        "\n",
        "                    z_norm = z - self.E(X_hat)\n",
        "                    z_norm = z_norm.norm().data / float(self.z_dim)\n",
        "                    mean_z_norm += z_norm\n",
        "\n",
        "                    norm_counter += 1\n",
        "\n",
        "\n",
        "            print(\"Eval loss G:\", test_loss_G / norm_counter)\n",
        "            print(\"Eval loss D:\", test_loss_D / norm_counter)\n",
        "\n",
        "            self.eval_hist['D_loss'].append(test_loss_D / norm_counter)\n",
        "            self.eval_hist['G_loss'].append(test_loss_G / norm_counter)\n",
        "\n",
        "            print(\"Pixel norm:\", mean_pixel_norm / norm_counter)\n",
        "            self.eval_hist['pixel_norm'].append( mean_pixel_norm / norm_counter )\n",
        "\n",
        "            with open('pixel_error_BIGAN.txt', 'a') as f:\n",
        "                f.writelines(str(mean_pixel_norm / norm_counter) + '\\n')\n",
        "\n",
        "            print(\"z norm:\", mean_z_norm / norm_counter)\n",
        "            self.eval_hist['z_norm'].append( mean_z_norm / norm_counter )\n",
        "\n",
        "            with open('z_error_BIGAN.txt', 'a') as f:\n",
        "                f.writelines(str(mean_z_norm / norm_counter) + '\\n')\n",
        "\n",
        "            ##### At the end of the epoch, save X and its reconstruction G(E(X))\n",
        "            samples = X.data.cpu().numpy()\n",
        "\n",
        "            fig = plt.figure(figsize=(10, 2))\n",
        "            gs = gridspec.GridSpec(2, 10)\n",
        "            gs.update(wspace=0.05, hspace=0.05)\n",
        "\n",
        "            for i, sample in enumerate(samples):\n",
        "                if i<10:\n",
        "                    ax = plt.subplot(gs[i])\n",
        "                    plt.axis('off')\n",
        "                    ax.set_xticklabels([])\n",
        "                    ax.set_yticklabels([])\n",
        "                    ax.set_aspect('equal')\n",
        "                    if self.network_type == 'FC':\n",
        "                        if self.dataset == 'mnist':\n",
        "                            sample = sample.reshape(28, 28)\n",
        "                            plt.imshow(sample, cmap='Greys_r')\n",
        "                        \n",
        "                        \n",
        "            X_hat = self.G(self.E(X).view(self.batch_size, self.z_dim))\n",
        "            samples = X_hat.data.cpu().numpy()\n",
        "\n",
        "\n",
        "            for i, sample in enumerate(samples):\n",
        "                if i<10:\n",
        "                    ax = plt.subplot(gs[10+i])\n",
        "                    plt.axis('off')\n",
        "                    ax.set_xticklabels([])\n",
        "                    ax.set_yticklabels([])\n",
        "                    ax.set_aspect('equal')\n",
        "                    if self.network_type == 'FC':\n",
        "                        if self.dataset == 'mnist':\n",
        "                            sample = sample.reshape(28, 28)\n",
        "                            plt.imshow(sample, cmap='Greys_r')\n",
        "                        \n",
        "\n",
        "            if not os.path.exists(self.result_dir + '/recons/'):\n",
        "                os.makedirs(self.result_dir + '/recons/')\n",
        "\n",
        "            filename = \"epoch_\" + str(epoch)\n",
        "            plt.savefig(self.result_dir + '/recons/{}.png'.format(filename), bbox_inches='tight')\n",
        "            plt.close()\n",
        "\n",
        "\n",
        "        save_plot_losses(self.train_hist['D_loss'], self.train_hist['G_loss'], self.eval_hist['D_loss'], self.eval_hist['G_loss'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "        save_plot_pixel_norm(self.eval_hist['pixel_norm'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "        save_plot_z_norm(self.eval_hist['z_norm'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "\n",
        "    def save_model(self):\n",
        "        if not os.path.exists(self.save_dir):\n",
        "            os.makedirs(self.save_dir)\n",
        "        \n",
        "        torch.save(self.G.state_dict(), self.save_dir + \"/G.pt\")\n",
        "        torch.save(self.E.state_dict(), self.save_dir + \"/E.pt\")\n",
        "        torch.save(self.D.state_dict(), self.save_dir + \"/D.pt\")\n",
        "\n",
        "    def load_model(self, kwargs):\n",
        "        if kwargs['network_type'] == 'FC':\n",
        "            # networks init\n",
        "            self.G = Generator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.D = Discriminator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.E = Encoder_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "\n",
        "        self.G.load_state_dict(torch.load(\"/models/G.pt\"))\n",
        "        self.E.load_state_dict(torch.load(\"/models/E.pt\"))\n",
        "        self.D.load_state_dict(torch.load(\"/models/D.pt\"))\n",
        "\n",
        "        if self.gpu_mode:\n",
        "            self.G.cuda()\n",
        "            self.D.cuda()\n",
        "            self.E.cuda()\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "avzTxKXxazeO",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## BiGAN Training"
      ]
    },
    {
      "metadata": {
        "id": "wTl6VjBQdVH4",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\"\"\"config\"\"\"\n",
        "\n",
        "kwargs = {\n",
        "    'dataset': 'mnist',\n",
        "    'dataset_path': '/home/zelazny/Downloads/tooploox/data/',\n",
        "    'gpu_mode': True,\n",
        "    'save_dir': 'models',\n",
        "    'result_dir': 'results',\n",
        "    'log_dir': 'logs',\n",
        "    'epoch': 400,\n",
        "    'batch_size': 128,\n",
        "    'lr': 1e-4, # initial learning rate\n",
        "    'lr_decay': (0.01)**(1/200), # exponential decay to 1e-6 over 200 epochs\n",
        "    'beta1': 0.5, # Adam1\n",
        "    'beta2': 0.999, # Adam2\n",
        "    'decay': 2.5*1e-5, # weight decay\n",
        "    'network_type': 'FC',\n",
        "    'z_dim': 50, # latent space dimension\n",
        "    'h_dim': 1024 # number of hidden units\n",
        "}\n",
        "\n",
        "\"\"\"check arguments\"\"\"\n",
        "\n",
        "def check_kwargs(kwargs):\n",
        "    # save_dir\n",
        "    if not os.path.exists(kwargs['save_dir']):\n",
        "        os.makedirs(kwargs['save_dir'])\n",
        "\n",
        "    # result_dir\n",
        "    if not os.path.exists(kwargs['result_dir']):\n",
        "        os.makedirs(kwargs['result_dir'])\n",
        "\n",
        "    # log_dir\n",
        "    if not os.path.exists(kwargs['log_dir']):\n",
        "        os.makedirs(kwargs['log_dir'])\n",
        "\n",
        "    # epoch\n",
        "    try:\n",
        "        assert kwargs['epoch'] >= 1\n",
        "    except:\n",
        "        print('number of epochs must be larger than or equal to one')\n",
        "\n",
        "    # batch_size\n",
        "    try:\n",
        "        assert kwargs['batch_size'] >= 1\n",
        "    except:\n",
        "        print('batch size must be larger than or equal to one')\n",
        "\n",
        "    return kwargs\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "xkbQw3S6vdoL",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\"\"\"main\"\"\"\n",
        "def main():\n",
        "    # check arguments\n",
        "    if kwargs is None:\n",
        "        exit()\n",
        "    else:\n",
        "        check_kwargs(kwargs)\n",
        "\n",
        "    bigan = BIGAN(kwargs)\n",
        "\n",
        "    # wipe old files\n",
        "    with open('pixel_error_BIGAN.txt', 'w') as f:\n",
        "        f.writelines('')\n",
        "    with open('z_error_BIGAN.txt', 'w') as f:\n",
        "        f.writelines('')\n",
        "\n",
        "    bigan.train()\n",
        "    print(\" [*] Training finished!\")\n",
        "\n",
        "    bigan.save_model()\n",
        "    \n",
        "\n",
        "if __name__ == '__main__':\n",
        "    main()\n",
        "\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "ZSgx8h_ma6uC",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## BiGAN Evaluation"
      ]
    },
    {
      "metadata": {
        "id": "nKpp0Dz-1Qb-",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# load trained net\n",
        "\n",
        "bigan = BIGAN(kwargs)\n",
        "bigan.load_model(kwargs)\n",
        "\n",
        "dataset = Mnist(128)\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "pAsAdJyfb4pR",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate encodings for dataset\n",
        "\n",
        "def generate_encodings(dataset, train_test):\n",
        "    if train_test == 'train':\n",
        "        loader = dataset.train_loader\n",
        "    elif train_test == 'test':\n",
        "        loader = dataset.test_loader\n",
        "    \n",
        "    encodings = []\n",
        "    labels = []\n",
        "    \n",
        "    for batch_id, (data, target) in enumerate(loader):\n",
        "    \n",
        "        X_data = Variable(data)\n",
        "    \n",
        "        if bigan.gpu_mode:\n",
        "            X_data = X_data.cuda()\n",
        "        \n",
        "        if X_data.size(0) == bigan.batch_size:\n",
        "            X = X_data\n",
        "            X = X.view(bigan.batch_size, -1)\n",
        "            z_hat = bigan.E(X)\n",
        "            encodings.append(z_hat)\n",
        "            labels.append(target)\n",
        "\n",
        "\n",
        "    encodings = torch.cat(encodings).data.cpu().numpy()\n",
        "    labels = torch.cat(labels).data.cpu().numpy()\n",
        "\n",
        "    return encodings, labels\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "kQkNsnhecOea",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate encodings for test set\n",
        "\n",
        "encodings_train, labels_train = generate_encodings(dataset, 'train')\n",
        "encodings_test, labels_test = generate_encodings(dataset, 'test')\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GNS4BaIk-oq5",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate and visualize t-SNE encodings for test set\n",
        "\n",
        "tsne_encodings_test = TSNE().fit_transform(encodings_test)\n",
        "scatter(tsne_encodings_test, labels_test)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "Fvt7xBP3BJVY",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# kNN classification accuracy on test set\n",
        "\n",
        "def knn_results(n):\n",
        "    knn = KNeighborsClassifier(n_neighbors=n)\n",
        "    knn.fit(encodings_train, labels_train)\n",
        "    labels_hat = knn.predict(encodings_test)\n",
        "\n",
        "    print('%sNN classification accuracy (%%)' %n, round(metrics.accuracy_score(labels_test, labels_hat)*100, 2))\n",
        "\n",
        "    \n",
        "for i in range(1, 11):\n",
        "    knn_results(i)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "dfPyciFBfKZk",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## BiCoGan Architecture and Utilities"
      ]
    },
    {
      "metadata": {
        "id": "KKhnpNyk1ElP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# BiCoGAN Generator\n",
        "\n",
        "class Generator_BiCoGAN(nn.Module):\n",
        "    \"\"\"\n",
        "    MLP incorporating label data. Input: two vectors, one from representation\n",
        "    space of dimension z_dim, the second one a one-hot encoding of the label,\n",
        "    of dimension y_dim. Intermediate layers are of dimension h_dim. Output is\n",
        "    a vector from image space of dimension X_dim.\n",
        "    \"\"\"\n",
        "    def __init__(self, z_dim, h_dim, X_dim, y_dim):\n",
        "        super(Generator_BiCoGAN, self).__init__()\n",
        "\n",
        "        self.z_dim = z_dim\n",
        "        self.h_dim = h_dim\n",
        "        self.X_dim = X_dim\n",
        "        self.y_dim = y_dim\n",
        "        \n",
        "        self.fc_y = torch.nn.Linear(y_dim, h_dim)\n",
        "        self.fc_z = torch.nn.Linear(z_dim, h_dim)\n",
        "\n",
        "        self.fc = torch.nn.Sequential(\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.Linear(h_dim, h_dim),\n",
        "            torch.nn.BatchNorm1d(h_dim),\n",
        "            torch.nn.ReLU(),\n",
        "            torch.nn.Linear(h_dim, X_dim),\n",
        "            torch.nn.Sigmoid()\n",
        "            )\n",
        "        \n",
        "        initialize_weights(self)\n",
        "\n",
        "    def forward(self, input_z, input_y):\n",
        "        y = self.fc_y(input_y)\n",
        "        z = self.fc_z(input_z)\n",
        "        x = y*z # element-wise multiplication\n",
        "        x = self.fc(x)\n",
        "        return x\n",
        "\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "v5Ci1z732yN2",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# BiCoGAN Encoder\n",
        "\n",
        "class Encoder_BiCoGAN(nn.Module):\n",
        "    \"\"\"\n",
        "    MLP incorporating label data. Input: vector from data space\n",
        "    of dimension X_dim. Output: a vector from representation space of dimension\n",
        "    z_dim and a second one which is a prediction of the label,of dimension y_dim.\n",
        "    Intermediate layers are of dimension h_dim.\n",
        "    \"\"\"\n",
        "    def __init__(self, z_dim, h_dim, X_dim, y_dim):\n",
        "        super(Encoder_BiCoGAN, self).__init__()\n",
        "\n",
        "        self.z_dim = z_dim\n",
        "        self.h_dim = h_dim\n",
        "        self.X_dim = X_dim\n",
        "        self.y_dim = y_dim\n",
        "\n",
        "        \n",
        "        self.fc = torch.nn.Sequential(\n",
        "            torch.nn.Linear(X_dim, h_dim),\n",
        "            nn.LeakyReLU(0.2),\n",
        "            torch.nn.Linear(h_dim, h_dim),\n",
        "            torch.nn.BatchNorm1d(h_dim),\n",
        "            nn.LeakyReLU(0.2)\n",
        "            )\n",
        "        \n",
        "        self.z_out = torch.nn.Linear(h_dim, z_dim)\n",
        "        \n",
        "        self.y_out = torch.nn.Sequential(\n",
        "            torch.nn.Linear(h_dim, y_dim),\n",
        "            torch.nn.Softmax()\n",
        "            )        \n",
        "\n",
        "        initialize_weights(self)\n",
        "\n",
        "    def forward(self, input_x):\n",
        "        x = self.fc(input_x)\n",
        "        y = self.y_out(x)\n",
        "        x = self.z_out(x)\n",
        "        return [x, y]\n",
        "\n",
        "    "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "65DFbE1hM21v",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# save plots of loss functions after training\n",
        "\n",
        "def save_plot_losses_BiCoGAN(train_D_loss, train_G_loss, train_E_loss, eval_D_loss, eval_G_loss, eval_E_loss, model_used, z_dim, epochs, lr, batch_size):\n",
        "\n",
        "    x = np.arange(1, len(train_D_loss) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_G_loss, label=\"Train G loss\", linewidth=2)\n",
        "    plt.plot(x, eval_G_loss, label=\"Eval G loss\", linewidth=2)\n",
        "\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the Train and Eval losses of G\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_G_losses.eps\", format='eps', dpi=1000)\n",
        "\n",
        "\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_D_loss, label=\"Train D loss\", linewidth=2)\n",
        "    plt.plot(x, eval_D_loss, label=\"Eval D loss\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "\n",
        "    plt.suptitle(\"Evolution of the Train and Eval losses of D\")\n",
        "\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_D_losses.eps\", format='eps', dpi=1000)\n",
        "    \n",
        "    \n",
        "    \n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_E_loss, label=\"Train E loss\", linewidth=2)\n",
        "    plt.plot(x, eval_E_loss, label=\"Eval E loss\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "\n",
        "    plt.suptitle(\"Evolution of the Train and Eval losses of E\")\n",
        "\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_E_losses.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n",
        "\n",
        "\n",
        "def save_plot_pixel_norm(mean_pixel_norm, model_used, z_dim, epochs, lr, batch_size):\n",
        "    x = np.arange(1, len(mean_pixel_norm) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, mean_pixel_norm, label=\"Reconstruction error\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Norm')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the reconstruction error between X and G(E(X))\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"pix2pix_norm.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n",
        "\n",
        "def save_plot_z_norm(mean_z_norm, model_used, z_dim, epochs, lr, batch_size):\n",
        "    x = np.arange(1, len(mean_z_norm) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, mean_z_norm, label=\"Reconstruction error\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Norm')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the reconstruction error between z and E(G(z))\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"z_norm.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "qefw49yNwgZw",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# BiCoGAN net\n",
        "\n",
        "class BiCoGAN(object):\n",
        "    \"\"\"\n",
        "    Class implementing a BiCoGAN network that trains from an observations dataset.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, kwargs):\n",
        "        \n",
        "        self.epoch = kwargs['epoch']\n",
        "        self.batch_size = kwargs['batch_size']\n",
        "        self.save_dir = kwargs['save_dir']\n",
        "        self.result_dir = kwargs['result_dir']\n",
        "        self.log_dir = kwargs['log_dir']\n",
        "        self.gpu_mode = kwargs['gpu_mode']\n",
        "        self.alpha = kwargs['alpha']\n",
        "        self.phi = kwargs['phi']\n",
        "        self.rho = kwargs['rho']\n",
        "        self.learning_rate = kwargs['lr']\n",
        "        self.lr_decay = kwargs['lr_decay']\n",
        "        self.beta1 = kwargs['beta1']\n",
        "        self.beta2 = kwargs['beta2']\n",
        "        self.decay = kwargs['decay']\n",
        "        self.network_type = kwargs['network_type']\n",
        "        self.dataset = kwargs['dataset']\n",
        "        self.dataset_path = kwargs['dataset_path']\n",
        "\n",
        "        # BIGAN parameters\n",
        "        self.z_dim = kwargs['z_dim']    #dimension of feature space\n",
        "        self.h_dim = kwargs['h_dim']    #dimension of the hidden layer\n",
        "        self.y_dim = kwargs['y_dim']    #number of label classes\n",
        "\n",
        "        if kwargs['dataset'] == 'mnist':\n",
        "            self.X_dim = 28*28    #dimension of data\n",
        "            self.num_channels = 1\n",
        "\n",
        "        if kwargs['network_type'] == 'FC':\n",
        "            # networks init\n",
        "            self.G = Generator_BiCoGAN(self.z_dim, self.h_dim, self.X_dim, self.y_dim)\n",
        "            self.D = Discriminator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.E = Encoder_BiCoGAN(self.z_dim, self.h_dim, self.X_dim, self.y_dim)\n",
        "        else:\n",
        "            raise Exception(\"[!] There is no option for \" + kwargs['network_type'])\n",
        "\n",
        "        if self.gpu_mode:\n",
        "            self.G.cuda()\n",
        "            self.D.cuda()\n",
        "            self.E.cuda()\n",
        "\n",
        "        \n",
        "        self.G_solver = optim.Adam(chain(self.E.parameters(), self.G.parameters()), lr=self.learning_rate, betas=[self.beta1,self.beta2], weight_decay=self.decay)\n",
        "        self.D_solver = optim.Adam(self.D.parameters(), lr=self.learning_rate, betas=[self.beta1,self.beta2], weight_decay=self.decay)\n",
        "        \n",
        "        # exponential learning rate decay starting halfway through learning\n",
        "        self.G_scheduler = optim.lr_scheduler.MultiStepLR(self.G_solver, milestones=list(range(self.epoch//2, self.epoch)), gamma=self.lr_decay)\n",
        "        self.D_scheduler = optim.lr_scheduler.MultiStepLR(self.D_solver, milestones=list(range(self.epoch//2, self.epoch)), gamma=self.lr_decay)\n",
        "\n",
        "\n",
        "\n",
        "        print('---------- Networks architecture -------------')\n",
        "        print_network(self.G)\n",
        "        print_network(self.E)\n",
        "        print_network(self.D)\n",
        "        print('-----------------------------------------------')\n",
        "\n",
        "\n",
        "\n",
        "    def D_(self, X, z):\n",
        "        return self.D(X, z)\n",
        "\n",
        "    def reset_grad(self):\n",
        "        self.E.zero_grad()\n",
        "        self.G.zero_grad()\n",
        "        self.D.zero_grad()\n",
        "\n",
        "\n",
        "    def train(self):\n",
        "        if self.dataset == 'mnist':\n",
        "            dataset = Mnist(self.batch_size)\n",
        "\n",
        "\n",
        "        self.train_hist = {}\n",
        "        self.train_hist['D_loss'] = []\n",
        "        self.train_hist['G_loss'] = []\n",
        "        self.train_hist['E_loss'] = []\n",
        "\n",
        "        self.eval_hist = {}\n",
        "        self.eval_hist['D_loss'] = []\n",
        "        self.eval_hist['G_loss'] = []\n",
        "        self.eval_hist['E_loss'] = []\n",
        "        self.eval_hist['pixel_norm'] = []\n",
        "        self.eval_hist['z_norm'] = []\n",
        "\n",
        "                    \n",
        "        for epoch in range(self.epoch):\n",
        "            print(\"epoch \", str(epoch))\n",
        "\n",
        "            self.D.train()\n",
        "            self.E.train()\n",
        "            self.G.train()\n",
        "\n",
        "            train_loss_G = 0\n",
        "            train_loss_D = 0\n",
        "            train_loss_E = 0\n",
        "            \n",
        "            # learning rate schedule\n",
        "            self.G_scheduler.step()\n",
        "            self.D_scheduler.step()\n",
        "            \n",
        "            self.efl_gamma = min(self.alpha*math.exp(self.rho*epoch), self.phi)\n",
        "\n",
        "\n",
        "            for batch_id, (data, target) in enumerate(dataset.train_loader):\n",
        "                # sometimes batchsize of X is not equal to actual batch_size\n",
        "                if data.size(0) == self.batch_size:\n",
        "\n",
        "                    if self.gpu_mode:\n",
        "                        # sample z\n",
        "                        z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1)).cuda() # uniform on [-1,1]\n",
        "                        # X is a real image from the dataset\n",
        "                        X = data\n",
        "                        X = Variable(X).cuda()\n",
        "                        y_onehot = torch.FloatTensor(self.batch_size, self.y_dim).cuda() # dummy for one-hot encodings\n",
        "                        y = Variable(target).cuda()                    \n",
        "                        target = target.cuda()\n",
        "                    else:\n",
        "                        z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1)) # uniform on [-1,1]\n",
        "                        X = data\n",
        "                        X = Variable(X)\n",
        "                        y_onehot = torch.FloatTensor(self.batch_size, self.y_dim) # dummy for one-hot encodings\n",
        "                        y = Variable(target)\n",
        "\n",
        "\n",
        "                    if self.network_type == 'FC':\n",
        "                        X = X.view(self.batch_size, -1)\n",
        "                        y = y.view(self.batch_size, -1)                        \n",
        "                        y_onehot.zero_()\n",
        "                        y_onehot.scatter_(1, y, 1)                        \n",
        "                        z_hat = self.E(X)[0]\n",
        "                        y_hat = self.E(X)[1]\n",
        "                        X_hat = self.G(z, y_onehot)\n",
        "\n",
        "                        D_enc = self.D_(X, z_hat)\n",
        "                        D_gen = self.D_(X_hat, z)\n",
        "\n",
        "\n",
        "                    D_loss = -torch.mean(log(D_enc) + log(1 - D_gen))\n",
        "                    G_loss = -torch.mean(log(D_gen) + log(1 - D_enc))\n",
        "                    \n",
        "                    EFL_loss = torch.nn.CrossEntropyLoss()\n",
        "                    E_loss = torch.mean(self.efl_gamma*EFL_loss(y_hat, target))\n",
        "                    \n",
        "                    G_EFL_loss = G_loss + E_loss\n",
        "\n",
        "                    D_loss.backward(retain_graph=True)\n",
        "                    self.D_solver.step()\n",
        "                    self.reset_grad()\n",
        "\n",
        "                    G_EFL_loss.backward()\n",
        "                    self.G_solver.step()\n",
        "                    self.reset_grad()\n",
        "                    \n",
        "\n",
        "                    train_loss_G += G_loss.data\n",
        "                    train_loss_D += D_loss.data\n",
        "                    train_loss_E += E_loss.data\n",
        "\n",
        "                    if batch_id % 1000 == 0:\n",
        "                        # Print and plot every now and then\n",
        "                        samples = X_hat.data.cpu().numpy()\n",
        "\n",
        "                        fig = plt.figure(figsize=(8, 4))\n",
        "                        gs = gridspec.GridSpec(4, 8)\n",
        "                        gs.update(wspace=0.05, hspace=0.05)\n",
        "\n",
        "                        for i, sample in enumerate(samples):\n",
        "                            if i<32:\n",
        "                                ax = plt.subplot(gs[i])\n",
        "                                plt.axis('off')\n",
        "                                ax.set_xticklabels([])\n",
        "                                ax.set_yticklabels([])\n",
        "                                ax.set_aspect('equal')\n",
        "\n",
        "                                if self.network_type == 'FC':\n",
        "                                    if self.dataset == 'mnist':\n",
        "                                        sample = sample.reshape(28, 28)\n",
        "                                        plt.imshow(sample, cmap='Greys_r')\n",
        "\n",
        "\n",
        "                        if not os.path.exists(self.result_dir + '/train/'):\n",
        "                            os.makedirs(self.result_dir + '/train/')\n",
        "\n",
        "                        filename = \"epoch_\" + str(epoch) + \"_batchid_\" + str(batch_id)\n",
        "                        plt.savefig(self.result_dir + '/train/{}.png'.format(filename, bbox_inches='tight'))\n",
        "                        plt.close()\n",
        "\n",
        "            print(\"Train loss G:\", train_loss_G / len(dataset.train_loader))\n",
        "            print(\"Train loss D:\", train_loss_D / len(dataset.train_loader))\n",
        "            print(\"Train loss E:\", train_loss_E / len(dataset.train_loader))\n",
        "\n",
        "            self.train_hist['D_loss'].append(train_loss_D / len(dataset.train_loader))\n",
        "            self.train_hist['G_loss'].append(train_loss_G / len(dataset.train_loader))\n",
        "            self.train_hist['E_loss'].append(train_loss_E / len(dataset.train_loader))\n",
        "\n",
        "\n",
        "            self.D.eval()\n",
        "            self.E.eval()\n",
        "            self.G.eval()\n",
        "            test_loss_G = 0\n",
        "            test_loss_D = 0\n",
        "            test_loss_E = 0\n",
        "\n",
        "            mean_pixel_norm = 0\n",
        "            mean_z_norm = 0\n",
        "            norm_counter = 1\n",
        "\n",
        "            for batch_id, (data, target) in enumerate(dataset.test_loader):\n",
        "                if data.size(0) == self.batch_size:\n",
        "                    # Sample data\n",
        "                    z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1)) # uniform on [-1,1]\n",
        "                    X_data = Variable(data)\n",
        "                    y_onehot = torch.FloatTensor(self.batch_size, self.y_dim) # dummy for one-hot encodings\n",
        "                    y = Variable(target)\n",
        "\n",
        "                    if self.gpu_mode:\n",
        "                        z = z.cuda()\n",
        "                        X_data = X_data.cuda()\n",
        "                        y_onehot = y_onehot.cuda()\n",
        "                        y = y.cuda()\n",
        "                        target = target.cuda()\n",
        "\n",
        "                    \n",
        "                    X = X_data\n",
        "                    if self.network_type == 'FC':                        \n",
        "                        X = X.view(self.batch_size, -1)\n",
        "                        y = y.view(self.batch_size, -1)\n",
        "                        y_onehot.zero_()\n",
        "                        y_onehot.scatter_(1, y, 1) # one-hot encoding of lables\n",
        "                        z_hat =  self.E(X)[0].view(self.batch_size, -1)\n",
        "                        y_hat = self.E(X)[1]\n",
        "                        X_hat = self.G(z, y_onehot)\n",
        "\n",
        "                        D_enc = self.D_(X, z_hat)\n",
        "                        D_gen = self.D_(X_hat, z)\n",
        "\n",
        "                    D_loss = -torch.mean(log(D_enc) + log(1 - D_gen))\n",
        "                    G_loss = -torch.mean(log(D_gen) + log(1 - D_enc))                    \n",
        "                    \n",
        "                    EFL_loss = torch.nn.CrossEntropyLoss()\n",
        "                    E_loss = torch.mean(self.efl_gamma*EFL_loss(y_hat, target))\n",
        "                    \n",
        "                    # not needed for test data\n",
        "                    #G_EFL_loss = G_loss + E_loss\n",
        "                    \n",
        "                    test_loss_G += G_loss.data\n",
        "                    test_loss_D += D_loss.data\n",
        "                    test_loss_E += E_loss.data\n",
        "\n",
        "                    pixel_norm = X -  self.G(z_hat, y_onehot)\n",
        "                    pixel_norm = pixel_norm.norm().data / float(self.X_dim)\n",
        "                    mean_pixel_norm += pixel_norm\n",
        "\n",
        "\n",
        "                    z_norm = z - self.E(X_hat)[0]\n",
        "                    z_norm = z_norm.norm().data / float(self.z_dim)\n",
        "                    mean_z_norm += z_norm\n",
        "\n",
        "                    norm_counter += 1\n",
        "\n",
        "\n",
        "            print(\"Eval loss G:\", test_loss_G / norm_counter)\n",
        "            print(\"Eval loss D:\", test_loss_D / norm_counter)\n",
        "            print(\"Eval loss E:\", test_loss_E / norm_counter)\n",
        "\n",
        "            self.eval_hist['D_loss'].append(test_loss_D / norm_counter)\n",
        "            self.eval_hist['G_loss'].append(test_loss_G / norm_counter)\n",
        "            self.eval_hist['E_loss'].append(test_loss_E / norm_counter)\n",
        "\n",
        "            print(\"Pixel norm:\", mean_pixel_norm / norm_counter)\n",
        "            self.eval_hist['pixel_norm'].append( mean_pixel_norm / norm_counter )\n",
        "\n",
        "            with open('pixel_error_BIGAN.txt', 'a') as f:\n",
        "                f.writelines(str(mean_pixel_norm / norm_counter) + '\\n')\n",
        "\n",
        "            print(\"z norm:\", mean_z_norm / norm_counter)\n",
        "            self.eval_hist['z_norm'].append( mean_z_norm / norm_counter )\n",
        "\n",
        "            with open('z_error_BIGAN.txt', 'a') as f:\n",
        "                f.writelines(str(mean_z_norm / norm_counter) + '\\n')\n",
        "\n",
        "            ##### At the end of the epoch, save X and its reconstruction G(E(X))\n",
        "            samples = X.data.cpu().numpy()\n",
        "\n",
        "            fig = plt.figure(figsize=(10, 2))\n",
        "            gs = gridspec.GridSpec(2, 10)\n",
        "            gs.update(wspace=0.05, hspace=0.05)\n",
        "\n",
        "            for i, sample in enumerate(samples):\n",
        "                if i<10:\n",
        "                    ax = plt.subplot(gs[i])\n",
        "                    plt.axis('off')\n",
        "                    ax.set_xticklabels([])\n",
        "                    ax.set_yticklabels([])\n",
        "                    ax.set_aspect('equal')\n",
        "                    if self.network_type == 'FC':\n",
        "                        if self.dataset == 'mnist':\n",
        "                            sample = sample.reshape(28, 28)\n",
        "                            plt.imshow(sample, cmap='Greys_r')\n",
        "                        \n",
        "            # alterative logic: keep Encoder output as is                        \n",
        "            #X_hat = self.G(self.E(X)[0].view(self.batch_size, self.z_dim), self.E(X)[1].view(self.batch_size, self.y_dim))\n",
        "            \n",
        "            # transform Encoder output to one-hot encoding\n",
        "            y_hat = torch.argmax(self.E(X)[1].view(self.batch_size, self.y_dim), dim=1)\n",
        "            y_hat = y_hat.view(self.batch_size, -1)\n",
        "            y_hat_onehot = torch.FloatTensor(self.batch_size, self.y_dim) # dummy for one-hot encodings\n",
        "            \n",
        "            if self.gpu_mode:\n",
        "                y_hat_onehot = y_hat_onehot.cuda()                           \n",
        "            \n",
        "            y_hat_onehot.zero_()\n",
        "            y_hat_onehot.scatter_(1, y_hat, 1)\n",
        "            \n",
        "            X_hat = self.G(self.E(X)[0].view(self.batch_size, self.z_dim), y_hat_onehot)\n",
        "            samples = X_hat.data.cpu().numpy()\n",
        "\n",
        "\n",
        "            for i, sample in enumerate(samples):\n",
        "                if i<10:\n",
        "                    ax = plt.subplot(gs[10+i])\n",
        "                    plt.axis('off')\n",
        "                    ax.set_xticklabels([])\n",
        "                    ax.set_yticklabels([])\n",
        "                    ax.set_aspect('equal')\n",
        "                    if self.network_type == 'FC':\n",
        "                        if self.dataset == 'mnist':\n",
        "                            sample = sample.reshape(28, 28)\n",
        "                            plt.imshow(sample, cmap='Greys_r')\n",
        "                        \n",
        "\n",
        "            if not os.path.exists(self.result_dir + '/recons/'):\n",
        "                os.makedirs(self.result_dir + '/recons/')\n",
        "\n",
        "            filename = \"epoch_\" + str(epoch)\n",
        "            plt.savefig(self.result_dir + '/recons/{}.png'.format(filename), bbox_inches='tight')\n",
        "            plt.close()\n",
        "\n",
        "\n",
        "        save_plot_losses_BiCoGAN(self.train_hist['D_loss'], self.train_hist['G_loss'], self.train_hist['E_loss'], self.eval_hist['D_loss'], self.eval_hist['G_loss'], self.eval_hist['E_loss'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "        save_plot_pixel_norm(self.eval_hist['pixel_norm'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "        save_plot_z_norm(self.eval_hist['z_norm'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "\n",
        "    def save_model(self):\n",
        "        if not os.path.exists(self.save_dir):\n",
        "            os.makedirs(self.save_dir)\n",
        "        \n",
        "        torch.save(self.G.state_dict(), self.save_dir + \"/G.pt\")\n",
        "        torch.save(self.E.state_dict(), self.save_dir + \"/E.pt\")\n",
        "        torch.save(self.D.state_dict(), self.save_dir + \"/D.pt\")\n",
        "\n",
        "    def load_model(self, kwargs):\n",
        "        if kwargs['network_type'] == 'FC':\n",
        "            # networks init\n",
        "            self.G = Generator_BiCoGAN(self.z_dim, self.h_dim, self.X_dim, self.y_dim)\n",
        "            self.D = Discriminator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.E = Encoder_BiCoGAN(self.z_dim, self.h_dim, self.X_dim, self.y_dim)\n",
        "        \n",
        "        self.G.load_state_dict(torch.load(\"/models/G.pt\"))\n",
        "        self.E.load_state_dict(torch.load(\"/models/E.pt\"))\n",
        "        self.D.load_state_dict(torch.load(\"/models/D.pt\"))\n",
        "        \n",
        "\n",
        "        if self.gpu_mode:\n",
        "            self.G.cuda()\n",
        "            self.D.cuda()\n",
        "            self.E.cuda()\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "s6W4v_nOg42v",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## BiCoGAN Training"
      ]
    },
    {
      "metadata": {
        "id": "h_hQmscHJlKA",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\"\"\"config\"\"\"\n",
        "\n",
        "kwargs_bicogan = {\n",
        "    'dataset': 'mnist',\n",
        "    'dataset_path': '/home/zelazny/Downloads/tooploox/data/',\n",
        "    'gpu_mode': True,\n",
        "    'save_dir': 'models',\n",
        "    'result_dir': 'results',\n",
        "    'log_dir': 'logs',\n",
        "    'epoch': 400,\n",
        "    'batch_size': 128,\n",
        "    'alpha' : 5,\n",
        "    'phi' : 10,\n",
        "    'rho' : 0.25,\n",
        "    'lr': 1e-4,\n",
        "    'lr_decay': (0.01)**(1/200), # exponential decay to 1e-6 over 200 epochs\n",
        "    'beta1': 0.5,\n",
        "    'beta2': 0.999,\n",
        "    # 'slope': 1e-2,\n",
        "    'decay': 2.5*1e-5,\n",
        "    # 'dropout': 0.2,\n",
        "    'network_type': 'FC',\n",
        "    'z_dim': 50,\n",
        "    'h_dim': 1024,\n",
        "    'y_dim': 10\n",
        "}\n",
        "\n",
        "\"\"\"check arguments\"\"\"\n",
        "\n",
        "def check_kwargs(kwargs_bicogan):\n",
        "    # save_dir\n",
        "    if not os.path.exists(kwargs_bicogan['save_dir']):\n",
        "        os.makedirs(kwargs_bicogan['save_dir'])\n",
        "\n",
        "    # result_dir\n",
        "    if not os.path.exists(kwargs_bicogan['result_dir']):\n",
        "        os.makedirs(kwargs_bicogan['result_dir'])\n",
        "\n",
        "    # log_dir\n",
        "    if not os.path.exists(kwargs_bicogan['log_dir']):\n",
        "        os.makedirs(kwargs_bicogan['log_dir'])\n",
        "\n",
        "    # epoch\n",
        "    try:\n",
        "        assert kwargs_bicogan['epoch'] >= 1\n",
        "    except:\n",
        "        print('number of epochs must be larger than or equal to one')\n",
        "\n",
        "    # batch_size\n",
        "    try:\n",
        "        assert kwargs_bicogan['batch_size'] >= 1\n",
        "    except:\n",
        "        print('batch size must be larger than or equal to one')\n",
        "\n",
        "    return kwargs_bicogan\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Xy0AUW6CKHMh",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#import warnings\n",
        "\n",
        "\"\"\"main\"\"\"\n",
        "def main():\n",
        "    #warnings.simplefilter('error', UserWarning)\n",
        "    # check arguments\n",
        "    if kwargs_bicogan is None:\n",
        "        exit()\n",
        "    else:\n",
        "        check_kwargs(kwargs_bicogan)\n",
        "\n",
        "    bicogan = BiCoGAN(kwargs_bicogan)\n",
        "\n",
        "    # wipe old files\n",
        "    with open('pixel_error_BIGAN.txt', 'w') as f:\n",
        "        f.writelines('')\n",
        "    with open('z_error_BIGAN.txt', 'w') as f:\n",
        "        f.writelines('')\n",
        "\n",
        "    bicogan.train()\n",
        "    print(\" [*] Training finished!\")\n",
        "\n",
        "    bicogan.save_model()\n",
        "    \n",
        "\n",
        "if __name__ == '__main__':\n",
        "    main()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "iPCD6WtrhTHf",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## BiCoGAN Evaluation"
      ]
    },
    {
      "metadata": {
        "id": "8idQmYzchDU0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# load trained net\n",
        "\n",
        "bicogan = BiCoGAN(kwargs_bicogan)\n",
        "bicogan.load_model(kwargs_bicogan)\n",
        "\n",
        "dataset = Mnist(128)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "KNIpBXSkhahB",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate encodings for dataset\n",
        "\n",
        "def generate_encodings(dataset, train_test):\n",
        "    if train_test == 'train':\n",
        "        loader = dataset.train_loader\n",
        "    elif train_test == 'test':\n",
        "        loader = dataset.test_loader\n",
        "    \n",
        "    encodings = []\n",
        "    labels = []\n",
        "    \n",
        "    for batch_id, (data, target) in enumerate(loader):\n",
        "    \n",
        "        X_data = Variable(data)\n",
        "    \n",
        "        if bicogan.gpu_mode:\n",
        "            X_data = X_data.cuda()\n",
        "        \n",
        "        if X_data.size(0) == bicogan.batch_size:\n",
        "            X = X_data\n",
        "            X = X.view(bicogan.batch_size, -1)\n",
        "            z_hat = bicogan.E(X)[0]\n",
        "            encodings.append(z_hat)\n",
        "            labels.append(target)\n",
        "\n",
        "\n",
        "    encodings = torch.cat(encodings).data.cpu().numpy()\n",
        "    labels = torch.cat(labels).data.cpu().numpy()\n",
        "\n",
        "    return encodings, labels\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "792ZyLbahbA6",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate encodings for test set\n",
        "\n",
        "encodings_train, labels_train = generate_encodings(dataset, 'train')\n",
        "encodings_test, labels_test = generate_encodings(dataset, 'test')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "RTd6wh_TkLNu",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate and visualize t-SNE encodings for test set\n",
        "\n",
        "tsne_encodings_test = TSNE().fit_transform(encodings_test)\n",
        "scatter(tsne_encodings_test, labels_test)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "DJvpGAA3nZkD",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# kNN classification accuracy on test set\n",
        "\n",
        "def knn_results(n):\n",
        "    knn = KNeighborsClassifier(n_neighbors=n)\n",
        "    knn.fit(encodings_train, labels_train)\n",
        "    labels_hat = knn.predict(encodings_test)\n",
        "\n",
        "    print('%sNN classification accuracy (%%)' %n, round(metrics.accuracy_score(labels_test, labels_hat)*100, 2))\n",
        "\n",
        "    \n",
        "for i in range(1, 11):\n",
        "    knn_results(i)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "5_UzsnikiAPX",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## TriBiGAN Architecture and Utilities"
      ]
    },
    {
      "metadata": {
        "id": "uv0J55Z40CTt",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# triplet loss\n",
        "\n",
        "class TripletLoss(nn.Module):\n",
        "    \"\"\"\n",
        "    Triplet loss\n",
        "    Takes embeddings of an anchor sample, a positive sample and a negative sample\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, margin):\n",
        "        super(TripletLoss, self).__init__()\n",
        "        self.margin = margin\n",
        "\n",
        "    def forward(self, anchor, positive, negative, size_average=True):\n",
        "        distance_positive = (anchor - positive).pow(2).sum(1)  # .pow(.5)\n",
        "        distance_negative = (anchor - negative).pow(2).sum(1)  # .pow(.5)\n",
        "        losses = F.relu(distance_positive - distance_negative + self.margin)\n",
        "        return losses.mean() if size_average else losses.sum()\n",
        "   \n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "fvdwD7CkCAqY",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# save plots of loss functions after training\n",
        "\n",
        "def save_plot_losses_triplet(train_D_loss, train_G_loss, train_triplet_loss, eval_D_loss, eval_G_loss, eval_triplet_loss, model_used, z_dim, epochs, lr, batch_size):\n",
        "\n",
        "    x = np.arange(1, len(train_D_loss) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_G_loss, label=\"Train G loss\", linewidth=2)\n",
        "    plt.plot(x, eval_G_loss, label=\"Eval G loss\", linewidth=2)\n",
        "\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the Train and Eval losses of G\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_G_losses.eps\", format='eps', dpi=1000)\n",
        "\n",
        "\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_D_loss, label=\"Train D loss\", linewidth=2)\n",
        "    plt.plot(x, eval_D_loss, label=\"Eval D loss\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "\n",
        "    plt.suptitle(\"Evolution of the Train and Eval losses of D\")\n",
        "\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_D_losses.eps\", format='eps', dpi=1000)\n",
        "    \n",
        "    \n",
        "    \n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, train_triplet_loss, label=\"Train triplet loss\", linewidth=2)\n",
        "    plt.plot(x, eval_triplet_loss, label=\"Eval triplet loss\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Loss')\n",
        "    plt.legend(loc='upper right')\n",
        "\n",
        "    plt.suptitle(\"Evolution of the Train and Eval triplet losses\")\n",
        "\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"plot_triplet_losses.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n",
        "\n",
        "\n",
        "def save_plot_pixel_norm(mean_pixel_norm, model_used, z_dim, epochs, lr, batch_size):\n",
        "    x = np.arange(1, len(mean_pixel_norm) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, mean_pixel_norm, label=\"Reconstruction error\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Norm')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the reconstruction error between X and G(E(X))\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"pix2pix_norm.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n",
        "\n",
        "def save_plot_z_norm(mean_z_norm, model_used, z_dim, epochs, lr, batch_size):\n",
        "    x = np.arange(1, len(mean_z_norm) + 1)\n",
        "\n",
        "    plt.figure(figsize=(8, 6))\n",
        "    plt.plot(x, mean_z_norm, label=\"Reconstruction error\", linewidth=2)\n",
        "\n",
        "    plt.axes().set_xlabel('Epoch')\n",
        "    plt.axes().set_ylabel('Norm')\n",
        "    plt.legend(loc='upper right')\n",
        "    plt.suptitle(\"Evolution of the reconstruction error between z and E(G(z))\")\n",
        "    params = \"Network type: \" + model_used + \", Dimension of latent space: \" + str(z_dim) + \", epochs: \" + str(epochs) + \", learning rate: \" + str(lr) + \", batch size:\" + str(batch_size)\n",
        "    plt.title(params, fontsize=8)\n",
        "    # plt.show()\n",
        "    plt.savefig(\"z_norm.eps\", format='eps', dpi=1000)\n",
        "    plt.close()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "CH5b_Ap38jMX",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# TriBiGAN net\n",
        "\n",
        "class TripletBIGAN(object):\n",
        "    \"\"\"\n",
        "    Class implementing a BIGAN network with an additional triplet loss.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, kwargs):\n",
        "        \n",
        "        self.epoch = kwargs['epoch']\n",
        "        self.batch_size = kwargs['batch_size']\n",
        "        self.save_dir = kwargs['save_dir']\n",
        "        self.result_dir = kwargs['result_dir']\n",
        "        self.log_dir = kwargs['log_dir']\n",
        "        self.gpu_mode = kwargs['gpu_mode']\n",
        "        self.triplet_margin = kwargs['triplet_margin']\n",
        "        self.learning_rate = kwargs['lr']\n",
        "        self.lr_decay = kwargs['lr_decay']\n",
        "        self.beta1 = kwargs['beta1']\n",
        "        self.beta2 = kwargs['beta2']\n",
        "        self.decay = kwargs['decay']      \n",
        "        self.network_type = kwargs['network_type']\n",
        "        self.dataset = kwargs['dataset']\n",
        "        self.dataset_path = kwargs['dataset_path']\n",
        "\n",
        "        # BIGAN parameters\n",
        "        self.z_dim = kwargs['z_dim']    #dimension of feature space\n",
        "        self.h_dim = kwargs['h_dim']    #dimension of the hidden layer\n",
        "\n",
        "        if kwargs['dataset'] == 'mnist':\n",
        "            self.X_dim = 28*28    #dimension of data\n",
        "            self.num_channels = 1\n",
        "\n",
        "        if kwargs['network_type'] == 'FC':\n",
        "            # networks init\n",
        "            self.G = Generator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.D = Discriminator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.E = Encoder_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "        else:\n",
        "            raise Exception(\"[!] There is no option for \" + kwargs['network_type'])\n",
        "\n",
        "        if self.gpu_mode:\n",
        "            self.G.cuda()\n",
        "            self.D.cuda()\n",
        "            self.E.cuda()\n",
        "\n",
        "        \n",
        "        self.G_solver = optim.Adam(chain(self.E.parameters(), self.G.parameters()), lr=self.learning_rate, betas=[self.beta1,self.beta2], weight_decay=self.decay)\n",
        "        self.D_solver = optim.Adam(self.D.parameters(), lr=self.learning_rate, betas=[self.beta1,self.beta2], weight_decay=self.decay)\n",
        "        \n",
        "        # exponential learning rate decay starting halfway through learning\n",
        "        self.G_scheduler = optim.lr_scheduler.MultiStepLR(self.G_solver, milestones=list(range(self.epoch//2, self.epoch)), gamma=self.lr_decay)\n",
        "        self.D_scheduler = optim.lr_scheduler.MultiStepLR(self.D_solver, milestones=list(range(self.epoch//2, self.epoch)), gamma=self.lr_decay)\n",
        "\n",
        "\n",
        "\n",
        "        print('---------- Networks architecture -------------')\n",
        "        print_network(self.G)\n",
        "        print_network(self.E)\n",
        "        print_network(self.D)\n",
        "        print('-----------------------------------------------')\n",
        "\n",
        "\n",
        "\n",
        "    def D_(self, X, z):\n",
        "        return self.D(X, z)\n",
        "\n",
        "    def reset_grad(self):\n",
        "        self.E.zero_grad()\n",
        "        self.G.zero_grad()\n",
        "        self.D.zero_grad()\n",
        "\n",
        "\n",
        "    def train(self):\n",
        "        if self.dataset == 'mnist':\n",
        "            dataset = TripletMnist(self.batch_size)\n",
        "\n",
        "\n",
        "        self.train_hist = {}\n",
        "        self.train_hist['D_loss'] = []\n",
        "        self.train_hist['G_loss'] = []\n",
        "        self.train_hist['triplet_loss'] = []\n",
        "\n",
        "        self.eval_hist = {}\n",
        "        self.eval_hist['D_loss'] = []\n",
        "        self.eval_hist['G_loss'] = []\n",
        "        self.eval_hist['triplet_loss'] = []\n",
        "        self.eval_hist['pixel_norm'] = []\n",
        "        self.eval_hist['z_norm'] = []\n",
        "\n",
        "\n",
        "        for epoch in range(self.epoch):\n",
        "            print(\"epoch \",str(epoch))\n",
        "\n",
        "            self.D.train()\n",
        "            self.E.train()\n",
        "            self.G.train()\n",
        "\n",
        "            train_loss_G = 0\n",
        "            train_loss_D = 0\n",
        "            train_triplet_loss = 0\n",
        "            \n",
        "            # learning rate schedule\n",
        "            self.G_scheduler.step()\n",
        "            self.D_scheduler.step()\n",
        "\n",
        "\n",
        "            for batch_id, (data, target) in enumerate(dataset.train_loader):\n",
        "                \n",
        "                # sometimes batchsize of X is not equal to actual batch_size\n",
        "                if data[0].size(0) == self.batch_size:\n",
        "                    # sample z\n",
        "                    z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1))# uniform on [-1,1]\n",
        "                    # (X = anchor) is a real image from the dataset                                                       \n",
        "                    anchor = Variable(data[0])\n",
        "                    positive = Variable(data[1])\n",
        "                    negative = Variable(data[2])\n",
        "\n",
        "                    if self.gpu_mode:\n",
        "                        z = z.cuda()\n",
        "                        anchor = anchor.cuda()\n",
        "                        positive = positive.cuda()\n",
        "                        negative = negative.cuda()\n",
        "\n",
        "                    if self.network_type == 'FC':                        \n",
        "                        anchor = anchor.view(self.batch_size, -1)\n",
        "                        positive = positive.view(self.batch_size, -1)\n",
        "                        negative = negative.view(self.batch_size, -1)\n",
        "                        z_hat = self.E(anchor)\n",
        "                        X_hat = self.G(z)\n",
        "\n",
        "                        D_enc = self.D_(anchor, z_hat)\n",
        "                        D_gen = self.D_(X_hat, z)                      \n",
        "\n",
        "\n",
        "                    D_loss = -torch.mean(log(D_enc) + log(1 - D_gen))\n",
        "                    G_loss = -torch.mean(log(D_gen) + log(1 - D_enc))                    \n",
        "                    \n",
        "                    triplet_fn = TripletLoss(self.triplet_margin)\n",
        "                    triplet_loss = triplet_fn(self.E(anchor), self.E(positive), self.E(negative))\n",
        "                    \n",
        "                    G_triplet_loss = G_loss + triplet_loss\n",
        "\n",
        "                    D_loss.backward(retain_graph=True)\n",
        "                    self.D_solver.step()\n",
        "                    self.reset_grad()\n",
        "\n",
        "                    G_triplet_loss.backward()\n",
        "                    self.G_solver.step()\n",
        "                    self.reset_grad()\n",
        "\n",
        "                    train_loss_G += G_loss.data\n",
        "                    train_loss_D += D_loss.data\n",
        "                    train_triplet_loss += triplet_loss.data\n",
        "\n",
        "                    if batch_id % 1000 == 0:\n",
        "                        # Print and plot every now and then\n",
        "                        samples = X_hat.data.cpu().numpy()\n",
        "\n",
        "                        fig = plt.figure(figsize=(8, 4))\n",
        "                        gs = gridspec.GridSpec(4, 8)\n",
        "                        gs.update(wspace=0.05, hspace=0.05)\n",
        "\n",
        "                        for i, sample in enumerate(samples):\n",
        "                            if i<32:\n",
        "                                ax = plt.subplot(gs[i])\n",
        "                                plt.axis('off')\n",
        "                                ax.set_xticklabels([])\n",
        "                                ax.set_yticklabels([])\n",
        "                                ax.set_aspect('equal')\n",
        "\n",
        "                                if self.network_type == 'FC':\n",
        "                                    if self.dataset == 'mnist':\n",
        "                                        sample = sample.reshape(28, 28)\n",
        "                                        plt.imshow(sample, cmap='Greys_r')\n",
        "\n",
        "\n",
        "                        if not os.path.exists(self.result_dir + '/train/'):\n",
        "                            os.makedirs(self.result_dir + '/train/')\n",
        "\n",
        "                        filename = \"epoch_\" + str(epoch) + \"_batchid_\" + str(batch_id)\n",
        "                        plt.savefig(self.result_dir + '/train/{}.png'.format(filename, bbox_inches='tight'))\n",
        "                        plt.close()\n",
        "\n",
        "            print(\"Train loss G:\", train_loss_G / len(dataset.train_loader))\n",
        "            print(\"Train loss D:\", train_loss_D / len(dataset.train_loader))\n",
        "            print(\"Train triplet loss:\", train_triplet_loss / len(dataset.train_loader))\n",
        "\n",
        "            self.train_hist['D_loss'].append(train_loss_D / len(dataset.train_loader))\n",
        "            self.train_hist['G_loss'].append(train_loss_G / len(dataset.train_loader))\n",
        "            self.train_hist['triplet_loss'].append(train_triplet_loss / len(dataset.train_loader))\n",
        "\n",
        "\n",
        "            self.D.eval()\n",
        "            self.E.eval()\n",
        "            self.G.eval()\n",
        "            test_loss_G = 0\n",
        "            test_loss_D = 0\n",
        "            test_triplet_loss = 0\n",
        "\n",
        "            mean_pixel_norm = 0\n",
        "            mean_z_norm = 0\n",
        "            norm_counter = 1\n",
        "\n",
        "            for batch_id, (data, target) in enumerate(dataset.test_loader):\n",
        "\n",
        "                if data[0].size(0) == self.batch_size:\n",
        "                    \n",
        "                    # Sample data\n",
        "                    z = Variable((1 - (-1))*torch.rand(self.batch_size, self.z_dim) + (-1)) # uniform on [-1,1]                    \n",
        "                    anchor = Variable(data[0])\n",
        "                    positive = Variable(data[1])\n",
        "                    negative = Variable(data[2])\n",
        "\n",
        "                    if self.gpu_mode:\n",
        "                        z = z.cuda()\n",
        "                        anchor = anchor.cuda()\n",
        "                        positive = positive.cuda()\n",
        "                        negative = negative.cuda()                \n",
        "                    \n",
        "                    if self.network_type == 'FC':\n",
        "                        anchor = anchor.view(self.batch_size, -1)\n",
        "                        positive = positive.view(self.batch_size, -1)\n",
        "                        negative = negative.view(self.batch_size, -1)\n",
        "                        z_hat = self.E(anchor)\n",
        "                        X_hat = self.G(z)\n",
        "\n",
        "                        D_enc = self.D_(anchor, z_hat)\n",
        "                        D_gen = self.D_(X_hat, z)\n",
        "                                                \n",
        "\n",
        "                    D_loss = -torch.mean(log(D_enc) + log(1 - D_gen))\n",
        "                    G_loss = -torch.mean(log(D_gen) + log(1 - D_enc))\n",
        "                    triplet_fn = TripletLoss(self.triplet_margin)                  \n",
        "                    triplet_loss = triplet_fn(self.E(anchor), self.E(positive), self.E(negative))\n",
        "\n",
        "                    test_loss_G += G_loss.data\n",
        "                    test_loss_D += D_loss.data\n",
        "                    test_triplet_loss += triplet_loss.data\n",
        "\n",
        "                    pixel_norm = anchor -  self.G(z_hat)\n",
        "                    pixel_norm = pixel_norm.norm().data / float(self.X_dim)\n",
        "                    mean_pixel_norm += pixel_norm\n",
        "\n",
        "\n",
        "                    z_norm = z - self.E(X_hat)\n",
        "                    z_norm = z_norm.norm().data / float(self.z_dim)\n",
        "                    mean_z_norm += z_norm\n",
        "\n",
        "                    norm_counter += 1\n",
        "\n",
        "\n",
        "            print(\"Eval loss G:\", test_loss_G / norm_counter)\n",
        "            print(\"Eval loss D:\", test_loss_D / norm_counter)\n",
        "            print(\"Eval triplet loss:\", test_triplet_loss / norm_counter)\n",
        "\n",
        "            self.eval_hist['D_loss'].append(test_loss_D / norm_counter)\n",
        "            self.eval_hist['G_loss'].append(test_loss_G / norm_counter)\n",
        "            self.eval_hist['triplet_loss'].append(test_triplet_loss / norm_counter)\n",
        "\n",
        "            print(\"Pixel norm:\", mean_pixel_norm / norm_counter)\n",
        "            self.eval_hist['pixel_norm'].append( mean_pixel_norm / norm_counter )\n",
        "\n",
        "            with open('pixel_error_BIGAN.txt', 'a') as f:\n",
        "                f.writelines(str(mean_pixel_norm / norm_counter) + '\\n')\n",
        "\n",
        "            print(\"z norm:\", mean_z_norm / norm_counter)\n",
        "            self.eval_hist['z_norm'].append( mean_z_norm / norm_counter )\n",
        "\n",
        "            with open('z_error_BIGAN.txt', 'a') as f:\n",
        "                f.writelines(str(mean_z_norm / norm_counter) + '\\n')\n",
        "\n",
        "            ##### At the end of the epoch, save X and its reconstruction G(E(X))\n",
        "            samples = anchor.data.cpu().numpy()\n",
        "\n",
        "            fig = plt.figure(figsize=(10, 2))\n",
        "            gs = gridspec.GridSpec(2, 10)\n",
        "            gs.update(wspace=0.05, hspace=0.05)\n",
        "\n",
        "            for i, sample in enumerate(samples):\n",
        "                if i<10:\n",
        "                    ax = plt.subplot(gs[i])\n",
        "                    plt.axis('off')\n",
        "                    ax.set_xticklabels([])\n",
        "                    ax.set_yticklabels([])\n",
        "                    ax.set_aspect('equal')\n",
        "                    if self.network_type == 'FC':\n",
        "                        if self.dataset == 'mnist':\n",
        "                            sample = sample.reshape(28, 28)\n",
        "                            plt.imshow(sample, cmap='Greys_r')\n",
        "                        \n",
        "                        \n",
        "            X_hat = self.G(self.E(anchor).view(self.batch_size, self.z_dim))\n",
        "            samples = X_hat.data.cpu().numpy()\n",
        "\n",
        "\n",
        "            for i, sample in enumerate(samples):\n",
        "                if i<10:\n",
        "                    ax = plt.subplot(gs[10+i])\n",
        "                    plt.axis('off')\n",
        "                    ax.set_xticklabels([])\n",
        "                    ax.set_yticklabels([])\n",
        "                    ax.set_aspect('equal')\n",
        "                    if self.network_type == 'FC':\n",
        "                        if self.dataset == 'mnist':\n",
        "                            sample = sample.reshape(28, 28)\n",
        "                            plt.imshow(sample, cmap='Greys_r')\n",
        "                        \n",
        "\n",
        "            if not os.path.exists(self.result_dir + '/recons/'):\n",
        "                os.makedirs(self.result_dir + '/recons/')\n",
        "\n",
        "            filename = \"epoch_\" + str(epoch)\n",
        "            plt.savefig(self.result_dir + '/recons/{}.png'.format(filename), bbox_inches='tight')\n",
        "            plt.close()\n",
        "\n",
        "\n",
        "        save_plot_losses_triplet(self.train_hist['D_loss'], self.train_hist['G_loss'], self.train_hist['triplet_loss'], self.eval_hist['D_loss'], self.eval_hist['G_loss'], self.eval_hist['triplet_loss'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "        save_plot_pixel_norm(self.eval_hist['pixel_norm'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "        save_plot_z_norm(self.eval_hist['z_norm'], self.network_type, self.z_dim, self.epoch, self.learning_rate, self.batch_size)\n",
        "\n",
        "    def save_model(self):\n",
        "        if not os.path.exists(self.save_dir):\n",
        "            os.makedirs(self.save_dir)\n",
        "        \n",
        "        torch.save(self.G.state_dict(), self.save_dir + \"/G.pt\")\n",
        "        torch.save(self.E.state_dict(), self.save_dir + \"/E.pt\")\n",
        "        torch.save(self.D.state_dict(), self.save_dir + \"/D.pt\")\n",
        "\n",
        "    def load_model(self, kwargs):\n",
        "        if kwargs['network_type'] == 'FC':\n",
        "            # networks init\n",
        "            self.G = Generator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.D = Discriminator_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "            self.E = Encoder_FC(self.z_dim, self.h_dim, self.X_dim)\n",
        "        \n",
        "        self.G.load_state_dict(torch.load(\"/models/G.pt\"))\n",
        "        self.E.load_state_dict(torch.load(\"/models/E.pt\"))\n",
        "        self.D.load_state_dict(torch.load(\"/models/D.pt\"))\n",
        "        \n",
        "\n",
        "        if self.gpu_mode:\n",
        "            self.G.cuda()\n",
        "            self.D.cuda()\n",
        "            self.E.cuda()\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "T1MFMSY-jVZ6",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## TriBiGAN Training"
      ]
    },
    {
      "metadata": {
        "id": "ydW2K_dBCr60",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "\"\"\"config\"\"\"\n",
        "\n",
        "kwargs_triplet = {\n",
        "    'dataset': 'mnist',\n",
        "    'dataset_path': '/home/zelazny/Downloads/tooploox/data/',\n",
        "    'gpu_mode': True,\n",
        "    'save_dir': 'models',\n",
        "    'result_dir': 'results',\n",
        "    'log_dir': 'logs',\n",
        "    'epoch': 400,\n",
        "    'batch_size': 128,\n",
        "    'triplet_margin': 0.2,\n",
        "    'lr': 1e-4,\n",
        "    'lr_decay': (0.01)**(1/200), # exponential decay to 1e-6 over 200 epochs\n",
        "    'beta1': 0.5,\n",
        "    'beta2': 0.999,\n",
        "    # 'slope': 1e-2,\n",
        "    'decay': 2.5*1e-5,\n",
        "    # 'dropout': 0.2,\n",
        "    'network_type': 'FC',\n",
        "    'z_dim': 50,\n",
        "    'h_dim': 1024\n",
        "}\n",
        "\n",
        "\"\"\"check arguments\"\"\"\n",
        "\n",
        "def check_kwargs(kwargs_triplet):\n",
        "    # save_dir\n",
        "    if not os.path.exists(kwargs_triplet['save_dir']):\n",
        "        os.makedirs(kwargs_triplet['save_dir'])\n",
        "\n",
        "    # result_dir\n",
        "    if not os.path.exists(kwargs_triplet['result_dir']):\n",
        "        os.makedirs(kwargs_triplet['result_dir'])\n",
        "\n",
        "    # log_dir\n",
        "    if not os.path.exists(kwargs_triplet['log_dir']):\n",
        "        os.makedirs(kwargs_triplet['log_dir'])\n",
        "\n",
        "    # epoch\n",
        "    try:\n",
        "        assert kwargs_triplet['epoch'] >= 1\n",
        "    except:\n",
        "        print('number of epochs must be larger than or equal to one')\n",
        "\n",
        "    # batch_size\n",
        "    try:\n",
        "        assert kwargs_triplet['batch_size'] >= 1\n",
        "    except:\n",
        "        print('batch size must be larger than or equal to one')\n",
        "\n",
        "    return kwargs_triplet\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "8yjFjA4qC8w0",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "#import warnings\n",
        "\n",
        "\"\"\"main\"\"\"\n",
        "def main():\n",
        "    #warnings.simplefilter('error', UserWarning)\n",
        "    # check arguments\n",
        "    if kwargs_triplet is None:\n",
        "        exit()\n",
        "    else:\n",
        "        check_kwargs(kwargs_triplet)\n",
        "\n",
        "    triplet_bigan = TripletBIGAN(kwargs_triplet)\n",
        "\n",
        "    # wipe old files\n",
        "    with open('pixel_error_BIGAN.txt', 'w') as f:\n",
        "        f.writelines('')\n",
        "    with open('z_error_BIGAN.txt', 'w') as f:\n",
        "        f.writelines('')\n",
        "\n",
        "    triplet_bigan.train()\n",
        "    print(\" [*] Training finished!\")\n",
        "\n",
        "    triplet_bigan.save_model()\n",
        "    \n",
        "\n",
        "if __name__ == '__main__':\n",
        "    main()\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "GykcsrkhlWGz",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## TriBiGAN Evaluation"
      ]
    },
    {
      "metadata": {
        "id": "cwlyYuGAEWtF",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# load trained net\n",
        "\n",
        "triplet_bigan = TripletBIGAN(kwargs_triplet)\n",
        "triplet_bigan.load_model(kwargs_triplet)\n",
        "\n",
        "dataset = TripletMnist(128)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "WhVy7ge8lcIP",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate encodings for dataset\n",
        "\n",
        "def generate_encodings(dataset, train_test):\n",
        "    if train_test == 'train':\n",
        "        loader = dataset.train_loader\n",
        "    elif train_test == 'test':\n",
        "        loader = dataset.test_loader\n",
        "    \n",
        "    encodings = []\n",
        "    labels = []\n",
        "    \n",
        "    for batch_id, (data, target) in enumerate(loader):\n",
        "\n",
        "        X_data = Variable(data[0])\n",
        "\n",
        "        if triplet_bigan.gpu_mode:\n",
        "            X_data = X_data.cuda()\n",
        "        \n",
        "        if X_data.size(0) == triplet_bigan.batch_size:\n",
        "            X = X_data\n",
        "            X = X.view(triplet_bigan.batch_size, -1)\n",
        "            z_hat = triplet_bigan.E(X)\n",
        "            encodings.append(z_hat)\n",
        "            labels.append(target)\n",
        "\n",
        "\n",
        "    encodings = torch.cat(encodings).data.cpu().numpy()\n",
        "    labels = torch.cat(labels).data.cpu().numpy()\n",
        "\n",
        "    return encodings, labels\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "8MSxUFHdlcV1",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate encodings for test set\n",
        "\n",
        "encodings_train, labels_train = generate_encodings(dataset, 'train')\n",
        "encodings_test, labels_test = generate_encodings(dataset, 'test')\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "Mp6XPNg5hkcS",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# generate and visualize t-SNE encodings for test set\n",
        "\n",
        "tsne_encodings_test = TSNE().fit_transform(encodings_test)\n",
        "scatter(tsne_encodings_test, labels_test)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "OGN3kbHcsf80",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "# kNN classification accuracy on test set\n",
        "\n",
        "def knn_results(n):\n",
        "    knn = KNeighborsClassifier(n_neighbors=n)\n",
        "    knn.fit(encodings_train, labels_train)\n",
        "    labels_hat = knn.predict(encodings_test)\n",
        "\n",
        "    print('%sNN classification accuracy (%%)' %n, round(metrics.accuracy_score(labels_test, labels_hat)*100, 2))\n",
        "\n",
        "    \n",
        "for i in range(1, 11):\n",
        "    knn_results(i)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "nZwCWBS0nzsv",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "## Conclusion"
      ]
    },
    {
      "metadata": {
        "id": "8-5l0xSXmO7J",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "Now we can compare the three architectures.\n",
        "\n",
        "And we are done! Hope you have enjoyed it!"
      ]
    }
  ]
}