Structured_GAN_commented.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# NIPS Paper Implementation Challenge \n",
    "## PyTorch Code Implementation for NIPS Paper Structured Generative Adversarial Networks\n",
    "https://arxiv.org/abs/1711.00889\n",
    "\n",
    "Authors: Jonas Falkner; Stefan Krempel; Jun-Jun Wan (after alphabetical order, equal contributions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   INFO [12:46:29] __main__: PyTorch version: 0.2.0_3\n"
     ]
    }
   ],
   "source": [
    "### IMPORTS ###\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import numpy as np\n",
    "import time\n",
    "import utils\n",
    "# initialize logger\n",
    "import logging.config\n",
    "import yaml\n",
    "with open('./log_config.yaml') as file:\n",
    "    Dict = yaml.load(file)    # load config file\n",
    "    logging.config.dictConfig(Dict)    # import config\n",
    "    \n",
    "logger = logging.getLogger(__name__)\n",
    "logger.info('PyTorch version: ' + str(torch.__version__))\n",
    "\n",
    "# import SGAN utils\n",
    "from layers import rampup, rampdown\n",
    "from zca import ZCA\n",
    "from models import Generator, InferenceNet, ClassifierNet, DConvNet1, DConvNet2\n",
    "from trainGAN import pretrain_classifier, train_classifier, train_gan, eval_classifier\n",
    "\n",
    "from sklearn.metrics import accuracy_score\n",
    "from torch.autograd import Variable\n",
    "import logging\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.utils import weight_norm as wn\n",
    "\n",
    "from layers import conv_concat, mlp_concat, init_weights, Gaussian_NoiseLayer, MeanOnlyBatchNorm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Global Parameter Setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   INFO [12:46:29] __main__: Cuda = False\n"
     ]
    }
   ],
   "source": [
    "\n",
    "### GLOBAL PARAMS ###\n",
    "BATCH_SIZE = 200\n",
    "BATCH_SIZE_EVAL = 200\n",
    "NUM_CLASSES = 10\n",
    "NUM_LABELLED = 4000\n",
    "SSL_SEED = 1\n",
    "NP_SEED = 1234\n",
    "CUDA = torch.cuda.is_available()\n",
    "logger.info('Cuda = ' + str(CUDA))\n",
    "\n",
    "# data dependent\n",
    "IN_CHANNELS = 3\n",
    "\n",
    "# evaluation\n",
    "VIS_EPOCH = 1\n",
    "EVAL_EPOCH = 1\n",
    "\n",
    "# C\n",
    "SCALED_UNSUP_WEIGHT_MAX = 100.0\n",
    "\n",
    "# G\n",
    "N_Z = 100\n",
    "\n",
    "# optimization\n",
    "B1 = 0.5  # moment1 in Adam\n",
    "LR = 3e-4\n",
    "LR_CLA = 3e-3\n",
    "NUM_EPOCHS = 1000\n",
    "NUM_EPOCHS_PRE = 20\n",
    "ANNEAL_EPOCH = 200\n",
    "ANNEAL_EVERY_EPOCH = 1\n",
    "ANNEAL_FACTOR = 0.995\n",
    "ANNEAL_FACTOR_CLA = 0.99\n",
    "\n",
    "# results and checkpoints\n",
    "path_out = \"./results\"\n",
    "if not os.path.exists(path_out):\n",
    "    os.makedirs(path_out)\n",
    "\n",
    "checkpoint_directory = \"./checkpoints/\"\n",
    "if not os.path.exists(checkpoint_directory):\n",
    "    os.makedirs(checkpoint_directory)\n",
    "\n",
    "classifier_result = checkpoint_directory + \"classifier_pretrained.pth\"\n",
    "EPOCH_SAVE_CHECKPOINTS = 25"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Data Preprocessing\n",
    "- Download the cifar-10 dataset from 'http://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz'\n",
    "- Split the dataset into labelled training and test dataset, with length of 50000 and 10000 resprectively\n",
    "- Create an unlabelled dataset by copying the training dataset created above\n",
    "- Shuffle the labelled training dataset\n",
    "- Create a much smaller length of labelled training dataset with length of 4000 for our semi-supervised classification setting.\n",
    "- Calculate the number of minibatches for labelled training and test, and unlabelled training datasets respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   INFO [12:46:31] __main__: Loading data...\n"
     ]
    }
   ],
   "source": [
    "### DATA ###\n",
    "logger.info('Loading data...')\n",
    "train_x, train_y = utils.load('./cifar10/', 'train')\n",
    "eval_x, eval_y = utils.load('./cifar10/', 'test')\n",
    "\n",
    "train_y = np.int32(train_y)\n",
    "eval_y = np.int32(eval_y)\n",
    "x_unlabelled = train_x.copy()\n",
    "\n",
    "rng_data = np.random.RandomState(SSL_SEED)\n",
    "inds = rng_data.permutation(train_x.shape[0])\n",
    "train_x = train_x[inds]\n",
    "train_y = train_y[inds]\n",
    "x_labelled = []\n",
    "y_labelled = []\n",
    "\n",
    "for j in range(NUM_CLASSES):\n",
    "    x_labelled.append(train_x[train_y == j][:int(NUM_LABELLED / NUM_CLASSES)])\n",
    "    y_labelled.append(train_y[train_y == j][:int(NUM_LABELLED / NUM_CLASSES)])\n",
    "\n",
    "x_labelled = np.concatenate(x_labelled, axis=0)\n",
    "y_labelled = np.concatenate(y_labelled, axis=0)\n",
    "del train_x\n",
    "\n",
    "num_batches_l = int(x_labelled.shape[0] // BATCH_SIZE)\n",
    "num_batches_u = int(x_unlabelled.shape[0] // BATCH_SIZE)\n",
    "num_batches_e = int(eval_x.shape[0] // BATCH_SIZE_EVAL)\n",
    "rng = np.random.RandomState(NP_SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model Structures\n",
    "The model structures of structured GAN consist of one generator network, two discriminator networks, and two inference networks which are named as classifier and inferentor respectively. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Generator Network $p_g (x \\mid y, z)$\n",
    "The generator parametrizes the sampling process $x \\sim p_g(x \\mid y, z) = G(y, z)$ by taking some hidden structures $y$ of $x$ and other factors $z$ as inputs, and outputing generated samples $x$. \n",
    "![title](generator_paper.png)\n",
    "\n",
    "The generator architecture consists of following layers, activation functions, and parameters.\n",
    "![title](generator.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "### MODEL STRUCTURES ###\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, input_size, num_classes, dense_neurons, weight_init=True):\n",
    "        super(Generator, self).__init__()\n",
    "\n",
    "        self.logger = logging.getLogger(__name__)  # initialize logger\n",
    "\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "        self.Dense = nn.Linear(input_size, dense_neurons)\n",
    "        self.Relu = nn.ReLU()\n",
    "        self.Tanh = nn.Tanh()\n",
    "\n",
    "        self.Deconv2D_0 = nn.ConvTranspose2d(in_channels=522, out_channels=256,\n",
    "                                             kernel_size=5, stride=2, padding=2,\n",
    "                                             output_padding=1, bias=False)\n",
    "        self.Deconv2D_1 = nn.ConvTranspose2d(in_channels=266, out_channels=128,\n",
    "                                             kernel_size=5, stride=2, padding=2,\n",
    "                                             output_padding=1, bias=False)\n",
    "        self.Deconv2D_2 = wn(nn.ConvTranspose2d(in_channels=138, out_channels=3,\n",
    "                                                kernel_size=5, stride=2, padding=2,\n",
    "                                                output_padding=1, bias=False))\n",
    "\n",
    "        self.BatchNorm1D = nn.BatchNorm1d(dense_neurons)\n",
    "\n",
    "        self.BatchNorm2D_0 = nn.BatchNorm2d(256)\n",
    "        self.BatchNorm2D_1 = nn.BatchNorm2d(128)\n",
    "\n",
    "        if weight_init:\n",
    "            # initialize weights for all conv and lin layers\n",
    "            self.apply(init_weights)\n",
    "            # log network structure\n",
    "            self.logger.debug(self)\n",
    "\n",
    "    def forward(self, z, y):\n",
    "        x = mlp_concat(z, y, self.num_classes)\n",
    "\n",
    "        x = self.Dense(x)\n",
    "        x = self.Relu(x)\n",
    "        x = self.BatchNorm1D(x)\n",
    "\n",
    "        x = x.resize(z.size(0), 512, 4, 4)\n",
    "        x = conv_concat(x, y, self.num_classes)\n",
    "\n",
    "        x = self.Deconv2D_0(x)                    # output shape (256,8,8) = 8192 * 2\n",
    "        x = self.Relu(x)\n",
    "        x = self.BatchNorm2D_0(x)\n",
    "\n",
    "        x = conv_concat(x, y, self.num_classes)\n",
    "\n",
    "        x = self.Deconv2D_1(x)                    # output shape (128,16,16) = 8192 * 2 * 2\n",
    "        x = self.Relu(x)\n",
    "        x = self.BatchNorm2D_1(x)\n",
    "\n",
    "        x = conv_concat(x, y, self.num_classes)\n",
    "        x = self.Deconv2D_2(x)                    # output shape (3, 32, 32) = 3072\n",
    "        x = self.Tanh(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discriminator Network 1\n",
    "\n",
    "Discriminator 1 is trained to distinguish generated pairs $(x, y) \\sim p_g(x,y)$ by using generator introduced above from those come from real pairs $p(x,y)$. It takes $x$ and $y$ as inputs and outputs the classification decisions whether the joint pairs are drawn from $p_g(x,y)$ or $p(x,y)$.\n",
    "\n",
    "![title](discriminator1.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# discriminator xy2p: test a pair of input comes from p(x, y) instead of p_c or p_g\n",
    "class DConvNet1(nn.Module):\n",
    "    '''\n",
    "    1st convolutional discriminative net (discriminator xy2p)\n",
    "    --> does a pair of input come from p(x, y) instead of p_c or p_g ?\n",
    "    '''\n",
    "\n",
    "    def __init__(self, channel_in, num_classes, p_dropout=0.2, weight_init=True):\n",
    "        super(DConvNet1, self).__init__()\n",
    "\n",
    "        self.logger = logging.getLogger(__name__)  # initialize logger\n",
    "\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "        # general reusable layers:\n",
    "        self.LReLU = nn.LeakyReLU(negative_slope=0.2)  # leaky ReLU activation function\n",
    "        self.sgmd = nn.Sigmoid()  # sigmoid activation function\n",
    "        self.drop = nn.Dropout(p=p_dropout)  # dropout layer\n",
    "\n",
    "        # input -->\n",
    "        # drop\n",
    "        # ConvConcat\n",
    "\n",
    "        self.conv1 = wn(nn.Conv2d(in_channels=channel_in + num_classes, out_channels=32,\n",
    "                                  kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False))\n",
    "        # LReLU\n",
    "        # ConvConcat\n",
    "\n",
    "        self.conv2 = wn(nn.Conv2d(in_channels=32 + num_classes, out_channels=32,\n",
    "                                  kernel_size=(3, 3), stride=2, padding=1, bias=False))\n",
    "        # LReLU\n",
    "        # drop\n",
    "        # ConvConcat\n",
    "\n",
    "        self.conv3 = wn(nn.Conv2d(in_channels=32 + num_classes, out_channels=64,\n",
    "                                  kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False))\n",
    "        # LReLU\n",
    "        # ConvConcat\n",
    "\n",
    "        self.conv4 = wn(nn.Conv2d(in_channels=64 + num_classes, out_channels=64,\n",
    "                                  kernel_size=(3, 3), stride=2, padding=1, bias=False))\n",
    "        # LReLU\n",
    "        # drop\n",
    "        # ConvConcat\n",
    "\n",
    "        self.conv5 = wn(nn.Conv2d(in_channels=64 + num_classes, out_channels=128,\n",
    "                                  kernel_size=(3, 3), stride=(1, 1), padding=0, bias=False))\n",
    "        # LReLU\n",
    "        # ConvConcat\n",
    "\n",
    "        self.conv6 = wn(nn.Conv2d(in_channels=128 + num_classes, out_channels=128,\n",
    "                                  kernel_size=(3, 3), stride=(1, 1), padding=0, bias=False))\n",
    "        # LReLU\n",
    "\n",
    "        self.globalPool = nn.AdaptiveAvgPool2d(output_size=4)\n",
    "\n",
    "        # MLPConcat\n",
    "\n",
    "        self.lin = nn.Linear(in_features=128 * 4 * 4 + num_classes,\n",
    "                             out_features=1)\n",
    "        # smg\n",
    "\n",
    "        if weight_init:\n",
    "            # initialize weights for all conv and lin layers\n",
    "            self.apply(init_weights)\n",
    "            # log network structure\n",
    "            self.logger.debug(self)\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        # x: (bs, channel_in, dim_input)\n",
    "        # y: (bs, 1)\n",
    "\n",
    "        x0 = self.drop(x)\n",
    "        x0 = conv_concat(x0, y, self.num_classes)\n",
    "\n",
    "        x1 = self.LReLU(self.conv1(x0))\n",
    "        x1 = conv_concat(x1, y, self.num_classes)\n",
    "\n",
    "        x2 = self.LReLU(self.conv2(x1))\n",
    "        x2 = self.drop(x2)\n",
    "        x2 = conv_concat(x2, y, self.num_classes)\n",
    "\n",
    "        x3 = self.LReLU(self.conv3(x2))\n",
    "        x3 = conv_concat(x3, y, self.num_classes)\n",
    "\n",
    "        x4 = self.LReLU(self.conv4(x3))\n",
    "        x4 = self.drop(x4)\n",
    "        x4 = conv_concat(x4, y, self.num_classes)\n",
    "\n",
    "        x5 = self.LReLU(self.conv5(x4))\n",
    "        x5 = conv_concat(x5, y, self.num_classes)\n",
    "\n",
    "        x6 = self.LReLU(self.conv6(x5))\n",
    "\n",
    "        x_pool = self.globalPool(x6)\n",
    "\n",
    "        x_pool = x_pool.view(-1, 128 * 4 * 4)\n",
    "        x_out = mlp_concat(x_pool, y, self.num_classes)\n",
    "\n",
    "        out = self.sgmd(self.lin(x_out))\n",
    "\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Discriminator Network 2\n",
    "\n",
    "Discriminator 2 is trained to distinguish generated pairs $(x, z) \\sim p_g(x,z)$ by using generator introduced above from those come from inferenced pairs $p_i (x,z)$ by using inferentor introduced below. It takes $x$ and $z$ as inputs and outputs the classification decisions whether the joint pairs are drawn from $p_g(x,z)$ or $p_i(x,z)$.\n",
    "![title](discriminator2.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# discriminator xz\n",
    "class DConvNet2(nn.Module):\n",
    "    '''\n",
    "    2nd convolutional discriminative net (discriminator xz)\n",
    "    '''\n",
    "\n",
    "    def __init__(self, n_z, channel_in, num_classes, weight_init=True):\n",
    "        super(DConvNet2, self).__init__()\n",
    "\n",
    "        self.logger = logging.getLogger(__name__)  # initialize logger\n",
    "\n",
    "        self.num_classes = num_classes\n",
    "\n",
    "        # general reusable layers:\n",
    "        self.LReLU = nn.LeakyReLU(negative_slope=0.2)  # leaky ReLU activation function\n",
    "        self.sgmd = nn.Sigmoid()  # sigmoid activation function\n",
    "\n",
    "        # z input -->\n",
    "        self.lin_z0 = nn.Linear(in_features=n_z,\n",
    "                                out_features=512)\n",
    "        # LReLU\n",
    "\n",
    "        self.lin_z1 = nn.Linear(in_features=512,\n",
    "                                out_features=512)\n",
    "        # LReLU\n",
    "\n",
    "        # -------------------------------------\n",
    "\n",
    "        # x input -->\n",
    "        self.conv_x0 = nn.Conv2d(in_channels=channel_in, out_channels=128,\n",
    "                                 kernel_size=(5, 5), stride=2, padding=2, bias=False)\n",
    "        # LReLU\n",
    "\n",
    "        self.conv_x1 = nn.Conv2d(in_channels=128, out_channels=256,\n",
    "                                 kernel_size=(5, 5), stride=2, padding=2, bias=False)\n",
    "        # LReLU\n",
    "        self.bn1 = nn.BatchNorm2d(num_features=256)\n",
    "\n",
    "        self.conv_x2 = nn.Conv2d(in_channels=256, out_channels=512,\n",
    "                                 kernel_size=(5, 5), stride=2, padding=2, bias=False)\n",
    "        # LReLU\n",
    "        self.bn2 = nn.BatchNorm2d(num_features=512)\n",
    "\n",
    "        # -------------------------------------\n",
    "\n",
    "        # concat x & z -->\n",
    "        self.lin_f0 = nn.Linear(in_features=8704,\n",
    "                                out_features=1024)\n",
    "        # LReLU\n",
    "\n",
    "        self.lin_f1 = nn.Linear(in_features=1024,\n",
    "                                out_features=1)\n",
    "        # smg\n",
    "\n",
    "        if weight_init:\n",
    "            # initialize weights for all conv and lin layers\n",
    "            self.apply(init_weights)\n",
    "            # log network structure\n",
    "            self.logger.debug(self)\n",
    "\n",
    "    def forward(self, z, x):\n",
    "        # x: (bs, channel_in, dim_input)\n",
    "        # z: (bs, n_z)\n",
    "\n",
    "        z0 = self.LReLU(self.lin_z0(z))\n",
    "        z_out = self.LReLU(self.lin_z1(z0))\n",
    "\n",
    "        x0 = self.LReLU(self.conv_x0(x))\n",
    "        x1 = self.LReLU(self.conv_x1(x0))\n",
    "        x1 = self.bn1(x1)\n",
    "        x_out = self.LReLU(self.conv_x2(x1))\n",
    "        x_out = self.bn2(x_out)\n",
    "\n",
    "        dims = x_out.size()\n",
    "        fusion = torch.cat([x_out.view(dims[0], -1).squeeze(-1).squeeze(-1), z_out], dim=1)\n",
    "\n",
    "        f_out = self.LReLU(self.lin_f0(fusion))\n",
    "        out = self.sgmd(self.lin_f1(f_out))\n",
    "\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Inference Networks\n",
    "Two inference networks define two distributions $p_i(z\\mid x)$ and $p_c(y\\mid x)$ that are used to approximate the true posterior $p(z\\mid x)$ and $p(y \\mid x)$ using two different adversarial games, which are named InferenceNet and ClassifierNet accordingly. \n",
    "![title](inferencenets.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class InferenceNet(nn.Module):\n",
    "    def __init__(self, in_channels, n_z, weight_init=True):\n",
    "        super(InferenceNet, self).__init__()\n",
    "\n",
    "        self.logger = logging.getLogger(__name__)  # initialize logger\n",
    "\n",
    "        self.inf02 = nn.Conv2d(in_channels=in_channels, out_channels=64, kernel_size=4,\n",
    "                               stride=2, padding=1)\n",
    "        self.inf03 = nn.BatchNorm2d(64)\n",
    "        self.inf11 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4,\n",
    "                               stride=2, padding=1)\n",
    "        self.inf12 = nn.BatchNorm2d(128)\n",
    "        self.inf21 = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4,\n",
    "                               stride=2, padding=1)\n",
    "        self.inf22 = nn.BatchNorm2d(256)\n",
    "        self.inf31 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4,\n",
    "                               stride=2, padding=1)\n",
    "        self.inf32 = nn.BatchNorm2d(512)\n",
    "        self.inf4 = nn.Linear(in_features=512*2*2, out_features=n_z)\n",
    "\n",
    "        if weight_init:\n",
    "            # initialize weights for all conv and lin layers\n",
    "            self.apply(init_weights)\n",
    "            # log network structure\n",
    "            self.logger.debug(self)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.leaky_relu(self.inf03(self.inf02(x)))\n",
    "        x = F.leaky_relu(self.inf12(self.inf11(x)))\n",
    "        x = F.leaky_relu(self.inf22(self.inf21(x)))\n",
    "        x = F.leaky_relu(self.inf32(self.inf31(x)))\n",
    "        x = x.view(-1, 512*2*2)\n",
    "        x = self.inf4(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Classifier\n",
    "The classifier is a inference network $C: x \\rightarrow y$ which approximates the posterior $p(y\\mid x)$ as $y \\sim p_c(y \\mid x) =C(x)$. In the case of using cifar-10 dataset, it is a 10-way classifier. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# classifier module\n",
    "class ClassifierNet(nn.Module):\n",
    "    def __init__(self, in_channels, weight_init=True):\n",
    "        super(ClassifierNet, self).__init__()\n",
    "\n",
    "        self.logger = logging.getLogger(__name__)  # initialize logger\n",
    "\n",
    "        self.gaussian = Gaussian_NoiseLayer()\n",
    "\n",
    "        self.conv1a = nn.Conv2d(in_channels=in_channels, out_channels=128, kernel_size=3,\n",
    "                                stride=1, padding=1)\n",
    "        self.convWN1 = MeanOnlyBatchNorm([1, 128, 32, 32])\n",
    "        self.conv1b = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3,\n",
    "                                stride=1, padding=1)\n",
    "        self.conv_relu = nn.LeakyReLU(negative_slope=0.1)\n",
    "        self.conv1c = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=3,\n",
    "                                stride=1, padding=1)\n",
    "        self.conv_maxpool1 = nn.MaxPool2d(kernel_size=2)\n",
    "        self.dropout1 = nn.Dropout2d(p=0.5)\n",
    "        self.conv2a = nn.Conv2d(in_channels=128, out_channels=256, kernel_size=3,\n",
    "                                stride=1, padding=1)\n",
    "        self.convWN2 = MeanOnlyBatchNorm([1, 256, 16, 16])\n",
    "        self.conv2b = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3,\n",
    "                                stride=1, padding=1)\n",
    "        self.conv2c = nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3,\n",
    "                                stride=1, padding=1)\n",
    "        self.conv_maxpool2 = nn.MaxPool2d(kernel_size=2)\n",
    "        self.dropout2 = nn.Dropout2d(p=0.5)\n",
    "        self.conv3a = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3,\n",
    "                                stride=1, padding=0)  # output[6,6]\n",
    "        self.convWN3a = MeanOnlyBatchNorm([1, 512, 6, 6])\n",
    "        self.conv3b = nn.Conv2d(in_channels=512, out_channels=256, kernel_size=1,\n",
    "                                stride=1, padding=0)\n",
    "        self.convWN3b = MeanOnlyBatchNorm([1, 256, 6, 6])\n",
    "        self.conv3c = nn.Conv2d(in_channels=256, out_channels=128, kernel_size=1,\n",
    "                                stride=1, padding=0)\n",
    "        self.convWN3c = MeanOnlyBatchNorm([1, 128, 6, 6])\n",
    "\n",
    "        self.conv_globalpool = nn.AdaptiveAvgPool2d(6)\n",
    "\n",
    "        self.dense = nn.Linear(in_features=128 * 6 * 6, out_features=10)\n",
    "        self.smx = nn.Softmax()\n",
    "        #self.WNfinal = MeanOnlyBatchNorm([1, 128, 6, 6])\n",
    "\n",
    "        if weight_init:\n",
    "            # initialize weights for all conv and lin layers\n",
    "            self.apply(init_weights)\n",
    "            # log network structure\n",
    "            self.logger.debug(self)\n",
    "\n",
    "    def forward(self, x, cuda, deterministic=False):\n",
    "        x = self.gaussian(x, cuda=cuda, deterministic=deterministic)\n",
    "        x = self.convWN1(self.conv_relu(self.conv1a(x)))\n",
    "        x = self.convWN1(self.conv_relu(self.conv1b(x)))\n",
    "        x = self.convWN1(self.conv_relu(self.conv1c(x)))\n",
    "        x = self.conv_maxpool1(x)\n",
    "        x = self.dropout1(x)\n",
    "        x = self.convWN2(self.conv_relu(self.conv2a(x)))\n",
    "        x = self.convWN2(self.conv_relu(self.conv2b(x)))\n",
    "        x = self.convWN2(self.conv_relu(self.conv2c(x)))\n",
    "        x = self.conv_maxpool2(x)\n",
    "        x = self.dropout2(x)\n",
    "        x = self.convWN3a(self.conv_relu(self.conv3a(x)))\n",
    "        x = self.convWN3b(self.conv_relu(self.conv3b(x)))\n",
    "        x = self.convWN3c(self.conv_relu(self.conv3c(x)))\n",
    "        x = self.conv_globalpool(x)\n",
    "        x = x.view(-1, 128 * 6 * 6)\n",
    "        #x = self.WNfinal(self.smx(self.dense(x)))\n",
    "        x = self.smx(self.dense(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Process\n",
    "The training process of structured GAN consists of pretraining of classifier, and training classifier, discriminator, inferentor, and generator in succession. It envolves training two adversarial games $\\mathcal{L}_{xy}$ and $\\mathcal{L}_{xz}$, and two collorative games $\\mathcal{R}_y$ and $\\mathcal{R}_z$. In following subsections, training details for different networks, such loss functions and optimizers used, are described. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pretraining\n",
    "Pretraining aims to obtain a relative good prior weights for later training phase. It is achieved by minimizing the reconstruction error of $y$ in terms of $C$, on both labeled data $X_l$ and generated data:  $$\\min_{C,G} \\mathcal{R}_y = - \\mathbb{E}_{(x,y)\\sim p(x,y)}[\\log p_c(y\\mid x)]$$\n",
    "\n",
    "Adam optimizer is used for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def pretrain_classifier(x_labelled, x_unlabelled, y_labelled, eval_x, eval_y, num_batches_l,\n",
    "                        batch_size, num_batches_u, classifier, whitener, losses, rng, cuda):\n",
    "    '''pre-trains the classifier model\n",
    "\n",
    "    Args:\n",
    "        x_labelled: labelled input data\n",
    "        x_unlabelled: unlabelled input data\n",
    "        y_labelled: target labels\n",
    "        eval_x: evaluation input data\n",
    "        eval_y: evaluation targets\n",
    "        num_batches_l: number of labelled batches\n",
    "        batch_size: size of mini-batch\n",
    "        num_batches_u: number of unlabelled batches\n",
    "        classifier: classifier model instance\n",
    "        whitener: zca whitener instance\n",
    "        losses: loss functions\n",
    "        rng: random generator\n",
    "        cuda: cuda flag\n",
    "\n",
    "    Returns: pre-trained classifier\n",
    "\n",
    "    '''\n",
    "\n",
    "    # randomly permute data and labels\n",
    "    permutation_labelled = rng.permutation(x_labelled.shape[0])\n",
    "    x_labelled = x_labelled[permutation_labelled]\n",
    "    y_labelled = y_labelled[permutation_labelled]\n",
    "    permutation_unlabelled = rng.permutation(x_unlabelled.shape[0]).astype('int32')\n",
    "\n",
    "    x_labelled = Variable(torch.from_numpy(x_labelled))\n",
    "    y_labelled = Variable(torch.from_numpy(y_labelled))\n",
    "\n",
    "    eval_x = Variable(torch.from_numpy(eval_x))\n",
    "    eval_y = Variable(torch.from_numpy(eval_y))\n",
    "    x_unlabelled = Variable(torch.from_numpy(x_unlabelled))\n",
    "\n",
    "    if cuda:\n",
    "        x_labelled, y_labelled, eval_x, eval_y, x_unlabelled = \\\n",
    "            x_labelled.cuda(), y_labelled.cuda(), eval_x.cuda(), eval_y.cuda(), x_unlabelled.cuda()\n",
    "\n",
    "    for i in range(num_batches_u):\n",
    "        i_c = i % num_batches_l\n",
    "        x_l = x_labelled[i_c * batch_size:(i_c + 1) * batch_size]\n",
    "        x_l_zca = whitener.apply(x_l)\n",
    "        y = y_labelled[i_c * batch_size:(i_c + 1) * batch_size]\n",
    "        y = y.type(torch.LongTensor)\n",
    "        if cuda:\n",
    "            y = y.cuda()\n",
    "\n",
    "        # classify input\n",
    "        cla_out_y_l = classifier(x_l_zca, cuda=cuda)\n",
    "        # calculate loss\n",
    "        cla_cost_l = losses['ce'](cla_out_y_l, y)\n",
    "\n",
    "        batch_slicer = torch.from_numpy(permutation_unlabelled[i * batch_size:(i + 1) * batch_size]).type(torch.LongTensor)\n",
    "        if cuda:\n",
    "            batch_slicer = batch_slicer.cuda()\n",
    "        x_u_rep = x_unlabelled[batch_slicer]\n",
    "        x_u_rep_zca = whitener.apply(x_u_rep)\n",
    "        # classify input\n",
    "        cla_out_y_rep = classifier(x_u_rep_zca, cuda=cuda)\n",
    "        target = cla_out_y_rep.detach()\n",
    "        del cla_out_y_rep\n",
    "\n",
    "        x_u = x_unlabelled[batch_slicer]\n",
    "        x_u_zca = whitener.apply(x_u)\n",
    "        # classify input\n",
    "        cla_out_y = classifier(x_u_zca, cuda=cuda)\n",
    "\n",
    "        # calculate loss\n",
    "        cla_cost_u = 100 * losses['mse'](cla_out_y, target)\n",
    "\n",
    "        # sum losses\n",
    "        pretrain_cost = cla_cost_l + cla_cost_u\n",
    "\n",
    "        # run optimization and update weights\n",
    "        cla_optimizer.zero_grad()\n",
    "        cla_optimizer = optim.Adam(classifier.parameters(), betas=(0.9, 0.999), lr=3e-3)  # they implement robust adam\n",
    "        pretrain_cost.backward()\n",
    "        cla_optimizer.step()\n",
    "\n",
    "    return classifier"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training Classifier\n",
    "Training classifier is achieved by minimizing the reconstruction error of $y$ in terms of $C$ and $G$, on both labeled data $X_l$ and generated data $(x,y)\\sim p_g(x,y)$:  $$\\min_{C,G} \\mathcal{R}_y = - \\mathbb{E}_{(x,y)\\sim p(x,y)}[\\log p_c(y\\mid x)] - \\mathbb{E}_{(x,y)\\sim p_g(x,y)}[\\log p_c(y\\mid x)]$$\n",
    "\n",
    "Adam optimizer is used for training with adaptive beta1 parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_classifier(x_labelled, y_labelled, x_unlabelled, num_batches_u, eval_epoch,\n",
    "                     size_l, size_u, size_g, n_z, whitener, classifier, p_u,\n",
    "                     unsup_weight, losses, generator, w_g, cla_lr, rng, b1_c, cuda):\n",
    "    '''does batch training for the classifier\n",
    "\n",
    "    Args:\n",
    "        x_labelled: labelled input data\n",
    "        y_labelled: target labels\n",
    "        x_unlabelled: unlabelled data\n",
    "        num_batches_u: number of unlabelled batches\n",
    "        eval_epoch: number of epochs to evaluate classifier\n",
    "        n_z: number of noise samples\n",
    "        whitener: ZCA whitener instance\n",
    "        classifier: Classifier Net instance\n",
    "        p_u: unlabelled sample data\n",
    "        unsup_weight: mse loss weight\n",
    "        losses(dict): dictionary containing respective loss instances (BCE, MSE, CE)\n",
    "        generator: generator model instance\n",
    "        w_g: cross entropy loss weight\n",
    "        cla_lr: learning rate for adam optimizer\n",
    "        rng: random generator\n",
    "        b1_c: beta1 for adam optimizer\n",
    "\n",
    "    Returns: classifier batch loss\n",
    "\n",
    "    '''\n",
    "\n",
    "    running_cla_cost = 0.0\n",
    "    epochs = num_batches_u * eval_epoch\n",
    "\n",
    "    for i in range(epochs):\n",
    "\n",
    "        i_l = i % (x_labelled.shape[0] // size_l)\n",
    "        i_u = i % (x_unlabelled.shape[0] // size_u)\n",
    "\n",
    "        y_real = np.int32(np.random.randint(10, size=size_g))\n",
    "        z_real = np.random.uniform(size=(size_g, n_z)).astype(np.float32)\n",
    "\n",
    "        x_l = x_labelled[i_l * size_l:(i_l + 1) * size_l]\n",
    "        y = y_labelled[i_l * size_l:(i_l + 1) * size_l]\n",
    "        x_l_zca = whitener.apply(x_l)\n",
    "\n",
    "        slice_x_u_c = p_u[i_u*size_u:(i_u+1)*size_u]\n",
    "        x_u_rep = x_unlabelled[slice_x_u_c]  # copy x_u_zca? double assigned variable??\n",
    "        x_u = x_unlabelled[slice_x_u_c]\n",
    "        x_u_rep_zca = whitener.apply(x_u_rep)\n",
    "        x_u_zca = whitener.apply(x_u)\n",
    "\n",
    "        # convert to torch tensor variable\n",
    "        y_real = Variable(torch.from_numpy(y_real))\n",
    "        z_real = Variable(torch.from_numpy(z_real))\n",
    "        y = Variable(torch.from_numpy(y).type(torch.LongTensor))\n",
    "        x_l_zca = Variable(torch.from_numpy(x_l_zca))\n",
    "        x_u_rep_zca = Variable(torch.from_numpy(x_u_rep_zca))\n",
    "        x_u_zca = Variable(torch.from_numpy(x_u_zca))\n",
    "        if cuda:\n",
    "            y_real, z_real, y = y_real.cuda(), z_real.cuda(), y.cuda()\n",
    "            x_l_zca, x_u_rep_zca, x_u_zca = x_l_zca.cuda(), x_u_rep_zca.cuda(), x_u_zca.cuda()\n",
    "\n",
    "        # classify input\n",
    "        cla_out_y_l = classifier(x_l_zca, cuda=cuda)\n",
    "        # calculate loss\n",
    "        cla_cost_l = losses['ce'](cla_out_y_l, y)  # size_average in pytorch is by default\n",
    "\n",
    "        # classify input for target\n",
    "        cla_out_y_rep = classifier(x_u_rep_zca, cuda=cuda)\n",
    "        target = cla_out_y_rep.detach()\n",
    "        del cla_out_y_rep\n",
    "\n",
    "        # classify input\n",
    "        cla_out_y = classifier(x_u_zca, cuda=cuda)\n",
    "        # calculate loss\n",
    "        cla_cost_u = unsup_weight * losses['mse'](cla_out_y, target)\n",
    "\n",
    "        y_m = y_real.type(torch.LongTensor)\n",
    "        z_m = z_real\n",
    "        if cuda:\n",
    "            y_m, z_m = y_m.cuda(), z_m.cuda()\n",
    "        gen_out_x_m = generator(z=z_m, y=y_m)\n",
    "        gen_out_x_m_zca = whitener.apply(gen_out_x_m)\n",
    "\n",
    "        # classify input\n",
    "        cla_out_y_m = classifier(gen_out_x_m_zca, cuda=cuda)\n",
    "        # calculate loss\n",
    "        cla_cost_g = losses['ce'](cla_out_y_m, y_m) * float(w_g)\n",
    "\n",
    "        # sum individual losses for backward\n",
    "        cla_cost = cla_cost_l + cla_cost_u + cla_cost_g\n",
    "\n",
    "        cla_optimizer = optim.Adam(classifier.parameters(), betas=(b1_c, 0.999), lr=cla_lr)\n",
    "        # zero the parameter gradients, optimize and update parameters\n",
    "        cla_optimizer.zero_grad()\n",
    "        cla_cost.backward()\n",
    "        cla_optimizer.step()\n",
    "\n",
    "        # update batch permutations\n",
    "        if i_l == ((x_labelled.shape[0] // size_l) - 1):\n",
    "            p_l = rng.permutation(x_labelled.shape[0])\n",
    "            x_labelled = x_labelled[p_l]\n",
    "            y_labelled = y_labelled[p_l]\n",
    "        if i_u == (num_batches_u - 1):\n",
    "            p_u = rng.permutation(x_unlabelled.shape[0]).astype('int32')\n",
    "\n",
    "        running_cla_cost += cla_cost.cpu().data.numpy().mean()\n",
    "\n",
    "    return running_cla_cost/epochs\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training Discriminators\n",
    "In the setting of a GAN, the discriminator is trained to decide whether a sample is from $p_{data}(x)$ or the generator distribution $p_g(x)$. In the present cases, two adversarial games are played. Discriminator 1 is used to maximize the loss $\\mathcal{L}_{xy}$:\n",
    "$$\\max_{D_{xy}} \\mathcal{L}_{xy} = \\mathbb{E}_{(x,y)\\sim p(x,y)} [\\log (D_{xy}(x_m,y_m))] + \\mathbb{E}_{y \\sim p(y), z \\sim p(z)} [\\log (1-D_{xy}(G(y,z),y_g))],$$ so that it is trained to discriminate $(x,y)\\sim p(x,y)$ from $(x,y)\\sim p_g(x,y)$.\n",
    "\n",
    "Discriminator 2 is used to distinguish pairs $(x, z) \\sim p_g(x,z)$ from those come from $p_i(x,z)$ by maximizing the loss $\\mathcal{L}_{xz}$:\n",
    "$$\\max_{D_{xz}} \\mathcal{L}_{xz} = \\mathbb{E}_{(x)\\sim p(x)} [\\log (D_{xz}(x_u, I(x_u)))] + \\mathbb{E}_{z \\sim p(z), y \\sim p(y)} [\\log (1-D_{xy}(G(y,z),z))]$$\n",
    "\n",
    "The total cost of both discriminator 1 and discriminator 2 is backpropagated and Adam optimizer is used."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_discriminator(discriminator1, discriminator2, generator, inferentor, classificator, whitener,\n",
    "                        x_labelled, x_unlabelled, y_labelled,\n",
    "                        slice_x_dis, y_real, z_real, slice_x_inf, sample_y, z_rand,\n",
    "                        batch_size, optimizer, loss, cuda):\n",
    "    '''does batch training for the discriminators\n",
    "\n",
    "    Args:\n",
    "        discriminator1(DConvNet1): Discriminator instance xy\n",
    "        discriminator2(DConvNet2): Discriminator instance xz\n",
    "        generator(Generator): Generator instance\n",
    "        inferentor(InferenceNet): Inference Net instance\n",
    "        classificator(ClassifierNet): Classifier Net instance\n",
    "        whitener(ZCA): ZCA  whitener instance\n",
    "        x_labelled: labelled input data\n",
    "        x_unlabelled: unlabelled input data\n",
    "        y_labelled: corresponding target labels\n",
    "        slice_x_dis: indexes to select unlabelled data for discriminator\n",
    "        y_real: class labels\n",
    "        z_real: generator_x_m noise input\n",
    "        slice_x_inf: indexes to select unlabelled data for inference net\n",
    "        sample_y: sampled labels\n",
    "        z_rand: generator_x noise input\n",
    "        batch_size(int): size of mini-batch\n",
    "        optimizer(torch.optim): optimizer instance for discriminators\n",
    "        loss(torch.nn.Loss): loss instance for discriminators (BCE)\n",
    "        cuda(bool): cuda flag (GPU)\n",
    "\n",
    "    Returns: total discriminator batch loss\n",
    "\n",
    "    '''\n",
    "\n",
    "\n",
    "    # get respective data slices for batch\n",
    "    unlabel_dis = x_unlabelled[slice_x_dis]  # original: sym_x_u_d\n",
    "    unlabel_dis_zca = whitener.apply(unlabel_dis)  # original: sym_x_u_d_zca\n",
    "    unlabel_inf = x_unlabelled[slice_x_inf]  # original: sym_x_u_i\n",
    "\n",
    "    # convert data ndarrays to pytorch tensor variables\n",
    "    x_labelled = Variable(torch.from_numpy(x_labelled))\n",
    "    y_labelled = Variable(torch.from_numpy(y_labelled))\n",
    "    unlabel_dis = Variable(torch.from_numpy(unlabel_dis))\n",
    "    unlabel_dis_zca = Variable(torch.from_numpy(unlabel_dis_zca))\n",
    "    unlabel_inf = Variable(torch.from_numpy(unlabel_inf))\n",
    "\n",
    "    if cuda:\n",
    "        x_labelled, y_labelled = x_labelled.cuda(), y_labelled.cuda()\n",
    "        unlabel_dis, unlabel_dis_zca, unlabel_inf = unlabel_dis.cuda(), unlabel_dis_zca.cuda(), unlabel_inf.cuda()\n",
    "\n",
    "    # generate samples\n",
    "    gen_out_x = generator(z=z_rand, y=sample_y)\n",
    "    gen_out_x_m = generator(z=z_real, y=y_real)\n",
    "\n",
    "    # compute inference\n",
    "    inf_z = inferentor(unlabel_inf)\n",
    "\n",
    "    # classify\n",
    "    cla_out = classificator(unlabel_dis_zca, cuda=cuda)\n",
    "    cla_out_val, cla_out_idx = cla_out.max(dim=1)\n",
    "\n",
    "    # concatenate inputs\n",
    "    x_in = torch.cat([x_labelled, unlabel_dis, gen_out_x_m], dim=0)[:batch_size]\n",
    "\n",
    "    y_labelled = y_labelled.long()\n",
    "    y_real = y_real.long()\n",
    "    y_in = torch.cat([y_labelled, cla_out_idx, y_real], dim=0)[:batch_size]\n",
    "\n",
    "    # calculate probabilities by discriminators\n",
    "    dis1_out_p = discriminator1(x=x_in, y=y_in)\n",
    "    dis1_out_pg = discriminator1(x=gen_out_x, y=sample_y)\n",
    "\n",
    "    dis2_out_p = discriminator2(z=inf_z, x=unlabel_inf)\n",
    "    dis2_out_pg = discriminator2(z=z_rand, x=gen_out_x)\n",
    "\n",
    "    # create discriminator labels\n",
    "    p_label_d1 = Variable(torch.ones(dis1_out_p.size()))\n",
    "    pg_label_d1 = Variable(torch.zeros(dis1_out_pg.size()))\n",
    "    p_label_d2 = Variable(torch.ones(dis2_out_p.size()))\n",
    "    pg_label_d2 = Variable(torch.zeros(dis2_out_pg.size()))\n",
    "\n",
    "    if cuda:\n",
    "        p_label_d1, pg_label_d1, \\\n",
    "        p_label_d2, pg_label_d2 = p_label_d1.cuda(), pg_label_d1.cuda(), \\\n",
    "                                  p_label_d2.cuda(), pg_label_d2.cuda()\n",
    "\n",
    "    # compute loss\n",
    "    dis1_cost_p = loss(dis1_out_p, p_label_d1)\n",
    "    dis1_cost_pg = loss(dis1_out_pg, pg_label_d1)\n",
    "    dis2_cost_p = loss(dis2_out_p, p_label_d2)\n",
    "    dis2_cost_pg = loss(dis2_out_pg, pg_label_d2)\n",
    "\n",
    "    # sum individual losses\n",
    "    dis1_cost = dis1_cost_p + dis1_cost_pg  # for report\n",
    "    dis2_cost = dis2_cost_p + dis2_cost_pg  # for report\n",
    "    total_cost = dis1_cost + dis2_cost\n",
    "\n",
    "    # optimization routines and weight updates\n",
    "    optimizer.zero_grad()\n",
    "    total_cost.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    return total_cost.cpu().data.numpy().mean()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training Inferentor\n",
    "On the one hand, the inferentor plays the adversarial game $\\mathcal{L}_{xz}$ by minimizing the loss $\\mathcal{L}_{xz}$ the discriminator 2 is trying to maximize as above described, on the other hand, the inferentor plays the collaborative game $\\mathcal{R}_z$ for the purpose of enforcing any other unstructured information not of interest to be fully captured in $z$, without being entangled with $y$, by minimizing \n",
    "$$\\min_{I, G} R_z = -\\mathbb{E}_{(x,z)\\sim p_g(x,z)}[\\log p_i (z \\mid x_g)]$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_inferentor(x_unlabelled, sample_y, generator, z_rand, discriminator2, inferentor,\n",
    "                     mse, bce, slice_x_u_i, optimizer, cuda):\n",
    "    '''does batch training for the inference model\n",
    "\n",
    "    Args:\n",
    "        x_unlabelled: unlabelled input data\n",
    "        sample_y: sampled target labels\n",
    "        generator: generator model instance\n",
    "        z_rand: random noise input for generator\n",
    "        discriminator2: discriminator2 (xz) model instance\n",
    "        inferentor: inference model instance\n",
    "        mse: MSE loss function\n",
    "        bce: BCE loss functions\n",
    "        slice_x_u_i: indexes to select unlabelled data\n",
    "        optimizer: optimizer instance\n",
    "        cuda: cuda flag\n",
    "\n",
    "    Returns: inferentor batch loss\n",
    "\n",
    "    '''\n",
    "\n",
    "    x_u_i = x_unlabelled[slice_x_u_i]\n",
    "    x_u_i = Variable(torch.from_numpy(x_u_i))\n",
    "\n",
    "    if cuda:\n",
    "        x_u_i = x_u_i.cuda()\n",
    "\n",
    "    y_g = sample_y\n",
    "    gen_out_x = generator(z_rand, y_g)\n",
    "    inf_z = inferentor(x_u_i)\n",
    "    inf_z_g = inferentor(gen_out_x)\n",
    "    disxz_out_p = discriminator2(z=inf_z, x=x_u_i)\n",
    "    target = inf_z_g.detach()\n",
    "    rz = mse(z_rand, target)\n",
    "\n",
    "    target = Variable(torch.zeros(disxz_out_p.size()))\n",
    "    if cuda:\n",
    "        target = target.cuda()\n",
    "    inf_cost_p_i = bce(disxz_out_p, target)\n",
    "    inf_cost = inf_cost_p_i + rz\n",
    "\n",
    "    optimizer.zero_grad()\n",
    "    inf_cost.backward()\n",
    "    optimizer.step()\n",
    "    return inf_cost.cpu().data.numpy().mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training Generator\n",
    "The generator jointly minimizing all four costs of adversarial games $\\mathcal{L}_{xz}$ and $\\mathcal{L}_{xy}$,  and collaborative games $\\mathcal{R}_{y}$ and $\\mathcal{R}_{z}$. In the setting of adversarial games, the generator attempts to fake the generated samples from the real ones. In the setting of collaborative games, the generator attempts to minimize the error of reconstructing either $y$ or $z$ for the purpose of increasing its controllability. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_generator(whitener, optimizer, BCE_loss, MSE_loss, cross_entropy_loss,\n",
    "                    discriminator1, discriminator2, inferentor, generator, classifier, sample_y, z_rand, cuda):\n",
    "    '''does batch training for the generator model\n",
    "\n",
    "    Args:\n",
    "        whitener(ZCA): ZCA whitener instance\n",
    "        optimizer: optimizer for generator\n",
    "        BCE_loss: binary cross entropy loss function\n",
    "        MSE_loss: mean squared error loss function\n",
    "        cross_entropy_loss: cross entropy loss function\n",
    "        discriminator1(DConvNet1): Discriminator instance xy\n",
    "        discriminator2(DConvNet2): Discriminator instance xz\n",
    "        inferentor: Inference model instance\n",
    "        generator: Generator model instance\n",
    "        classifier: Classifier model instance\n",
    "        sample_y: sampled labels\n",
    "        z_rand: noise sample for generator\n",
    "\n",
    "\n",
    "    Returns: generator batch loss\n",
    "\n",
    "    '''\n",
    "    # compute loss\n",
    "    gen_out_x = generator(z_rand, sample_y)\n",
    "    inf_z_g = inferentor(gen_out_x)\n",
    "    gen_out_x_zca = whitener.apply(gen_out_x)\n",
    "    cla_out_y_g = classifier(gen_out_x_zca, cuda=cuda)\n",
    "    rz = MSE_loss(inf_z_g, z_rand)\n",
    "    sample_y = sample_y.long()\n",
    "    ry = cross_entropy_loss(cla_out_y_g, sample_y)\n",
    "    dis_out_p_g = discriminator1(x=gen_out_x, y=sample_y)\n",
    "    disxz_out_p_g = discriminator2(z=z_rand, x=gen_out_x)\n",
    "\n",
    "    target1, target2 = Variable(torch.ones(dis_out_p_g.size())), Variable(torch.ones(disxz_out_p_g.size()))\n",
    "    if cuda:\n",
    "        target1, target2 = target1.cuda(), target2.cuda()\n",
    "\n",
    "    gen_cost_p_g_1 = BCE_loss(dis_out_p_g, target1)\n",
    "    gen_cost_p_g_2 = BCE_loss(disxz_out_p_g, target2)\n",
    "\n",
    "    generator_cost = gen_cost_p_g_1 + gen_cost_p_g_2 + rz + ry\n",
    "\n",
    "    # optimization routines and weight updates\n",
    "    optimizer.zero_grad()\n",
    "    generator_cost.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    return generator_cost.cpu().data.numpy().mean()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def train_gan(discriminator1, discriminator2, generator, inferentor, classifier, whitener,\n",
    "              x_labelled, x_unlabelled, y_labelled, p_u_d, p_u_i,\n",
    "              num_classes, batch_size, num_batches_u,\n",
    "              batch_c, batch_l, batch_g,\n",
    "              n_z, optimizers, losses, rng, cuda=False):\n",
    "\n",
    "    '''does one GAN training epoch for generator, inferentor and discriminator\n",
    "\n",
    "    Args:\n",
    "        discriminator1(DConvNet1): Discriminator instance xy\n",
    "        discriminator2(DConvNet2): Discriminator instance xz\n",
    "        generator(Generator): Generator instance\n",
    "        inferentor(InferenceNet): Inference Net instance\n",
    "        classifier(ClassifierNet): Classifier Net instance\n",
    "        whitener(ZCA): ZCA instance\n",
    "        x_labelled: batch of labelled input data\n",
    "        x_unlabelled: batch of unlabelled input data\n",
    "        y_labelled: batch of corresponding labels\n",
    "        p_u_d: data slice object (idx)\n",
    "        p_u_i: data slice object (idx)\n",
    "        num_classes(int): number of target classes\n",
    "        batch_size(int): size of mini-batch\n",
    "        num_batches_u:\n",
    "        n_z: number of noise samples\n",
    "        optimizers(dict): py dictionary containing optimizer instances for all respective nets (dis, gen, inf)\n",
    "        losses(dict): py dictionary containing respective loss instances (bce, mse, ce)\n",
    "        b1: beta1 in Adam\n",
    "        cuda(bool): cuda flag\n",
    "\n",
    "    Returns: py dictionary containing losses of all respective nets (dis, gen, inf)\n",
    "\n",
    "    '''\n",
    "\n",
    "    for i in range(num_batches_u):\n",
    "            i_l = i % (x_labelled.shape[0] // batch_l)\n",
    "\n",
    "            from_u_i = i*batch_size  # unlabelled inferentor slice\n",
    "            to_u_i = (i+1)*batch_size\n",
    "            from_u_d = i*batch_c    # unlabelled discriminator slice\n",
    "            to_u_d = (i+1) * batch_c\n",
    "            from_l = i_l*batch_l    # labelled\n",
    "            to_l = (i_l+1)*batch_l\n",
    "\n",
    "            # create samples and labels\n",
    "            sample_y = torch.from_numpy(np.int32(np.repeat(np.arange(num_classes), int(batch_size/num_classes))))\n",
    "            y_real = torch.from_numpy(np.int32(np.random.randint(10, size=batch_g)))\n",
    "            z_real = torch.from_numpy(np.random.uniform(size=(batch_g, n_z)).astype(np.float32))\n",
    "            z_rand = torch.rand((batch_size*n_z)).view(batch_size, n_z)\n",
    "\n",
    "            sample_y, y_real, z_real, z_rand = Variable(sample_y), Variable(y_real), Variable(z_real), Variable(z_rand)\n",
    "            if cuda:\n",
    "                sample_y, y_real, z_real, z_rand = sample_y.cuda(), y_real.cuda(), z_real.cuda(), z_rand.cuda()\n",
    "\n",
    "            dis_losses = train_discriminator(discriminator1=discriminator1,\n",
    "                                             discriminator2=discriminator2,\n",
    "                                             generator=generator,\n",
    "                                             inferentor=inferentor,\n",
    "                                             classificator=classifier,\n",
    "                                             whitener=whitener,\n",
    "                                             x_labelled=x_labelled[from_l:to_l],  # sym_x_l\n",
    "                                             x_unlabelled=x_unlabelled,\n",
    "                                             y_labelled=y_labelled[from_l:to_l],  # sym_y\n",
    "                                             slice_x_dis=p_u_d[from_u_d:to_u_d],  # slice_x_u_d\n",
    "                                             y_real=y_real,  # sym_y_m\n",
    "                                             z_real=z_real,  # sym_z_m\n",
    "                                             slice_x_inf=p_u_i[from_u_i:to_u_i],  # slice_x_u_i\n",
    "                                             sample_y=sample_y,  # sym_y_g\n",
    "                                             z_rand=z_rand,\n",
    "                                             batch_size=batch_size,\n",
    "                                             optimizer=optimizers['dis'],\n",
    "                                             loss=losses['bce'],\n",
    "                                             cuda=cuda)\n",
    "\n",
    "            inf_losses = train_inferentor(x_unlabelled=x_unlabelled,\n",
    "                                          sample_y=sample_y,\n",
    "                                          generator=generator,\n",
    "                                          z_rand=z_rand,\n",
    "                                          discriminator2=discriminator2,\n",
    "                                          inferentor=inferentor,\n",
    "                                          mse=losses['mse'],\n",
    "                                          bce=losses['bce'],\n",
    "                                          slice_x_u_i=p_u_i[from_u_i:to_u_i],\n",
    "                                          optimizer=optimizers['inf'],\n",
    "                                          cuda=cuda)\n",
    "\n",
    "            gen_losses = train_generator(whitener=whitener,\n",
    "                                         optimizer=optimizers['gen'],\n",
    "                                         BCE_loss=losses['bce'],\n",
    "                                         MSE_loss=losses['mse'],\n",
    "                                         cross_entropy_loss=losses['ce'],\n",
    "                                         discriminator1=discriminator1,\n",
    "                                         discriminator2=discriminator2,\n",
    "                                         inferentor=inferentor,\n",
    "                                         generator=generator,\n",
    "                                         classifier=classifier,\n",
    "                                         sample_y=sample_y,\n",
    "                                         z_rand=z_rand,\n",
    "                                         cuda=cuda)\n",
    "\n",
    "\n",
    "            gan_loss = {\n",
    "                'dis': dis_losses,\n",
    "                'inf': inf_losses,\n",
    "                'gen': gen_losses\n",
    "            }\n",
    "\n",
    "            return gan_loss\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Accuracy Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def eval_classifier(num_batches_e, eval_x, eval_y, batch_size, whitener, classifier, cuda):\n",
    "    '''evaluates the classifier\n",
    "\n",
    "    Args:\n",
    "        num_batches_e: number of evaluation batches\n",
    "        eval_x: evaluation input data\n",
    "        eval_y: evaluation target labels\n",
    "        batch_size: size of mini-batch\n",
    "        whitener: ZCA whitener instance\n",
    "        classifier: classifier model instance\n",
    "        cuda: cuda flag\n",
    "\n",
    "    Returns: estimated accurracy of the classifier\n",
    "\n",
    "    '''\n",
    "\n",
    "    accurracy = []\n",
    "    for i in range(num_batches_e):\n",
    "        x_eval = eval_x[i*batch_size:(i+1)*batch_size]\n",
    "        y_eval = eval_y[i*batch_size:(i+1)*batch_size]\n",
    "        x_eval_zca = whitener.apply(x_eval)\n",
    "        x_eval_zca = Variable(torch.from_numpy(x_eval_zca))\n",
    "        if cuda:\n",
    "            x_eval_zca = x_eval_zca.cuda()\n",
    "\n",
    "        cla_out_y_eval = classifier(x_eval_zca, cuda=cuda, deterministic=True)\n",
    "\n",
    "        pred = cla_out_y_eval.cpu().data.numpy()\n",
    "        pred = np.argmax(pred, axis=1)\n",
    "\n",
    "        accurracy_batch = accuracy_score(y_eval, pred)\n",
    "        accurracy.append(accurracy_batch)\n",
    "\n",
    "    return np.mean(accurracy)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialization of Different Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "### INITS ###\n",
    "\n",
    "# GENRATOR\n",
    "generator = Generator(input_size=110, num_classes=NUM_CLASSES, dense_neurons=(4 * 4 * 512))\n",
    "\n",
    "# INFERENCE\n",
    "inference = InferenceNet(in_channels=IN_CHANNELS, n_z=N_Z)\n",
    "\n",
    "# CLASSIFIER\n",
    "classifier = ClassifierNet(in_channels=IN_CHANNELS)\n",
    "\n",
    "# DISCRIMINATOR\n",
    "discriminator1 = DConvNet1(channel_in=IN_CHANNELS, num_classes=NUM_CLASSES)\n",
    "discriminator2 = DConvNet2(n_z=N_Z, channel_in=IN_CHANNELS, num_classes=NUM_CLASSES)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Cuda Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# put on GPU\n",
    "if CUDA:\n",
    "    generator.cuda()\n",
    "    inference.cuda()\n",
    "    classifier.cuda()\n",
    "    discriminator1.cuda()\n",
    "    discriminator2.cuda()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# ZCA\n",
    "whitener = ZCA(x=x_unlabelled)\n",
    "\n",
    "# LOSS FUNCTIONS\n",
    "if CUDA:\n",
    "    losses = {\n",
    "        'bce': nn.BCELoss().cuda(),\n",
    "        'mse': nn.MSELoss().cuda(),\n",
    "        'ce': nn.CrossEntropyLoss().cuda()\n",
    "    }\n",
    "else:\n",
    "    losses = {\n",
    "        'bce': nn.BCELoss(),\n",
    "        'mse': nn.MSELoss(),\n",
    "        'ce': nn.CrossEntropyLoss()\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 1: Pretraining for 20 Epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "### PRETRAIN CLASSIFIER ###\n",
    "if os.path.isfile(classifier_result):\n",
    "    logger.info('Load pretrained classifier from disk')\n",
    "    classifier.load_state_dict(torch.load(classifier_result))\n",
    "else:\n",
    "    logger.info('Start pretraining...')\n",
    "    for epoch in range(1, 1+NUM_EPOCHS_PRE):\n",
    "\n",
    "        # pretrain classifier net\n",
    "        classifier = pretrain_classifier(x_labelled, x_unlabelled, y_labelled, eval_x, eval_y, num_batches_l,\n",
    "                                         BATCH_SIZE, num_batches_u, classifier, whitener, losses, rng, CUDA)\n",
    "\n",
    "        # evaluate\n",
    "        accurracy = eval_classifier(num_batches_e, eval_x, eval_y, BATCH_SIZE_EVAL, whitener, classifier, CUDA)\n",
    "\n",
    "        logger.info(str(epoch) + ':Pretrain error_rate: ' + str(1 - accurracy))\n",
    "\n",
    "        torch.save(classifier.state_dict(), classifier_result)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Step 2: Training for 1000 Epochs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "### GAN TRAINING ###\n",
    "\n",
    "# assign start values\n",
    "lr_cla = LR_CLA\n",
    "lr = LR\n",
    "start_full = time.time()\n",
    "\n",
    "logger.info(\"Start GAN training...\")\n",
    "for epoch in range(1, 1+NUM_EPOCHS):\n",
    "\n",
    "    # OPTIMIZERS\n",
    "    optimizers = {\n",
    "        'dis': optim.Adam(list(discriminator1.parameters()) + list(discriminator2.parameters()), betas=(B1, 0.999), lr=lr),\n",
    "        'gen': optim.Adam(generator.parameters(), betas=(B1, 0.999), lr=lr),\n",
    "        'inf': optim.Adam(inference.parameters(), betas=(B1, 0.999), lr=lr)\n",
    "    }\n",
    "\n",
    "    # randomly permute data and labels each epoch\n",
    "    p_l = rng.permutation(x_labelled.shape[0])\n",
    "    x_labelled = x_labelled[p_l]\n",
    "    y_labelled = y_labelled[p_l]\n",
    "\n",
    "    # permuted slicer objects\n",
    "    p_u = rng.permutation(x_unlabelled.shape[0]).astype('int32')\n",
    "    p_u_d = rng.permutation(x_unlabelled.shape[0]).astype('int32')\n",
    "    p_u_i = rng.permutation(x_unlabelled.shape[0]).astype('int32')\n",
    "\n",
    "    # set epoch dependent values\n",
    "    if epoch < (NUM_EPOCHS/2):\n",
    "        if epoch % 50 == 1:\n",
    "            batch_l = 200 - (epoch // 50 + 1) * 16\n",
    "            batch_c = (epoch // 50 + 1) * 16\n",
    "            batch_g = 1\n",
    "    elif epoch < NUM_EPOCHS and epoch % 100 == 0:\n",
    "        batch_l = 50\n",
    "        batch_c = 140 - 10 * (epoch-500)/100\n",
    "        batch_g = 10 + 10 * (epoch-500)/100\n",
    "\n",
    "    # if current epoch is an evaluation epoch, train classifier and report results\n",
    "    if epoch % EVAL_EPOCH == 0:\n",
    "\n",
    "        logger.info('Train classifier...')\n",
    "\n",
    "        rampup_value = rampup(epoch-1)\n",
    "        rampdown_value = rampdown(epoch-1)\n",
    "        b1_c = rampdown_value * 0.9 + (1.0 - rampdown_value) * 0.5\n",
    "        unsup_weight = rampup_value * SCALED_UNSUP_WEIGHT_MAX if epoch > 1 else 0.0\n",
    "        w_g = np.float32(min(float(epoch) / 300.0, 1.0))\n",
    "\n",
    "        size_l = 100\n",
    "        size_g = 100\n",
    "        size_u = 100\n",
    "\n",
    "        cla_losses = train_classifier(x_labelled=x_labelled,\n",
    "                                      y_labelled=y_labelled,\n",
    "                                      x_unlabelled=x_unlabelled,\n",
    "                                      num_batches_u=num_batches_u,\n",
    "                                      eval_epoch=EVAL_EPOCH,\n",
    "                                      size_l=size_l,\n",
    "                                      size_u=size_u,\n",
    "                                      size_g=size_g,\n",
    "                                      n_z=N_Z,\n",
    "                                      whitener=whitener,\n",
    "                                      classifier=classifier,\n",
    "                                      p_u=p_u,\n",
    "                                      unsup_weight=unsup_weight,\n",
    "                                      losses=losses,\n",
    "                                      generator=generator,\n",
    "                                      w_g=w_g,\n",
    "                                      cla_lr=lr_cla,\n",
    "                                      rng=rng,\n",
    "                                      b1_c=b1_c,\n",
    "                                      cuda=CUDA)\n",
    "\n",
    "        # evaluate & report\n",
    "        accurracy = eval_classifier(num_batches_e, eval_x, eval_y, BATCH_SIZE_EVAL, whitener, classifier, CUDA)\n",
    "\n",
    "        logger.info('Evaluation error_rate: %.5f\\n' % (1 - accurracy))\n",
    "\n",
    "    logger.info('Train generator, inference and discriminator model...')\n",
    "    # train GAN model\n",
    "    for i in range(num_batches_u):\n",
    "        gan_losses = train_gan(discriminator1=discriminator1,\n",
    "                               discriminator2=discriminator2,\n",
    "                               generator=generator,\n",
    "                               inferentor=inference,\n",
    "                               classifier=classifier,\n",
    "                               whitener=whitener,\n",
    "                               x_labelled=x_labelled,\n",
    "                               x_unlabelled=x_unlabelled,\n",
    "                               y_labelled=y_labelled,\n",
    "                               p_u_d=p_u_d,\n",
    "                               p_u_i=p_u_i,\n",
    "                               num_classes=NUM_CLASSES,\n",
    "                               batch_size=BATCH_SIZE,\n",
    "                               num_batches_u=num_batches_u,\n",
    "                               batch_c=batch_c,\n",
    "                               batch_l=batch_l,\n",
    "                               batch_g=batch_g,\n",
    "                               n_z=N_Z,\n",
    "                               optimizers=optimizers,\n",
    "                               losses=losses,\n",
    "                               rng=rng,\n",
    "                               cuda=CUDA)\n",
    "\n",
    "    # anneal the learning rates\n",
    "    if (epoch >= ANNEAL_EPOCH) and (epoch % ANNEAL_EVERY_EPOCH == 0):\n",
    "        lr = lr * ANNEAL_FACTOR\n",
    "        lr_cla *= ANNEAL_FACTOR_CLA\n",
    "\n",
    "    # report and log training info\n",
    "    t = time.time() - start_full\n",
    "    line = \"*Epoch=%d Time=%.2f LR=%.5f\\n\" % (epoch, t, lr) + \"DisLosses: \" + str(gan_losses['dis']) + \"\\nGenLosses: \" + \\\n",
    "           str(gan_losses['gen']) + \"\\nInfLosses: \" + str(gan_losses['inf']) + \"\\nClaLosses: \" + str(cla_losses)\n",
    "    logger.info(line)\n",
    "\n",
    "    loss_logger.info(line)\n",
    "\n",
    "    # save checkpoints\n",
    "    if epoch % EPOCH_SAVE_CHECKPOINTS == 0 or epoch == NUM_EPOCHS:\n",
    "        utils.save_checkpoint(generator, \"generator\", checkpoint_directory, epoch)\n",
    "        utils.save_checkpoint(discriminator1, \"discriminator1\", checkpoint_directory, epoch)\n",
    "        utils.save_checkpoint(discriminator2, \"discriminator2\", checkpoint_directory, epoch)\n",
    "        utils.save_checkpoint(inference, \"inferencer\", checkpoint_directory, epoch)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}