SSVAE_with_Pyro_MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision.datasets as dset\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.datasets import MNIST\n",
    "\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "import pyro.contrib.examples.util  # patches torchvision\n",
    "from pyro.infer import SVI, JitTrace_ELBO, JitTraceEnum_ELBO, Trace_ELBO, TraceEnum_ELBO, config_enumerate\n",
    "from pyro.optim import Adam\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from utils.datasets import SupervisedDataset\n",
    "from utils.utils import One_Hot\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, z_dim, hidden_dim, out_dim):\n",
    "        super().__init__()\n",
    "        # setup the three linear transformations used\n",
    "        self.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(4,4), stride=2)# 28x28 -> 24/2+1x24/2+1\n",
    "        self.conv2 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=(3,3),stride=2)# 13x13 -> 10/2+1x10/2+1\n",
    "        self.fc1 = nn.Linear(6*6*32+out_dim,hidden_dim)\n",
    "        self.fc21 = nn.Linear(hidden_dim, z_dim)\n",
    "        self.fc22 = nn.Linear(hidden_dim, z_dim)\n",
    "        # setup the non-linearities\n",
    "        self.softplus = nn.Softplus()\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        #print(x.shape)\n",
    "        #print(y.shape)\n",
    "        # define the forward computation on the image x\n",
    "        # first shape the mini-batch to have pixels in the rightmost dimension\n",
    "        x = self.relu(self.conv1(x))\n",
    "        x = self.relu(self.conv2(x))\n",
    "        x = x.view(-1, 6*6*32)\n",
    "        x = torch.cat([x,y],axis=1)\n",
    "        # then compute the hidden units\n",
    "        hidden = self.relu(self.fc1(x))\n",
    "        # then return a mean vector and a (positive) square root covariance\n",
    "        # each of size batch_size x z_dim\n",
    "        z_loc = self.fc21(hidden)\n",
    "        z_scale = self.softplus(self.fc22(hidden))\n",
    "        return z_loc, z_scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, z_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        # setup the two linear transformations used\n",
    "        self.conv2 = nn.ConvTranspose2d(in_channels=64, out_channels=1, kernel_size=(4,4), stride=2)# 28x28 -> 24/2+1x24/2+1\n",
    "        self.conv1 = nn.ConvTranspose2d(in_channels=32, out_channels=64, kernel_size=(3,3), stride=2)# 13x13 -> 10/2+1x10/2+1\n",
    "        self.fc1 = nn.Linear(z_dim, hidden_dim)\n",
    "        self.fc2 = nn.Linear(hidden_dim, 32*6*6)\n",
    "        # setup the non-linearities\n",
    "        self.relu = nn.ReLU()\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, z):\n",
    "        # define the forward computation on the latent z\n",
    "        # first compute the hidden units\n",
    "        hidden = self.relu(self.fc1(z))\n",
    "        hidden = self.relu(self.fc2(hidden))\n",
    "        hidden = hidden.view(-1,32,6,6)\n",
    "        hidden = self.relu(self.conv1(hidden))\n",
    "        hidden = self.sigmoid(self.conv2(hidden))\n",
    "        # return the parameter for the output Bernoulli\n",
    "        # each is of size batch_size x 784\n",
    "        loc_img = hidden.view(-1,784)\n",
    "        return loc_img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Classifier(nn.Module):\n",
    "    def __init__(self, y_dim, hidden_dim):\n",
    "        super().__init__()\n",
    "        # setup the three linear transformations used\n",
    "        self.conv1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=(4,4), stride=2)# 28x28 -> 24/2+1x24/2+1\n",
    "        self.conv2 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=(3,3),stride=2)# 13x13 -> 10/2+1x10/2+1\n",
    "        self.fc1 = nn.Linear(6*6*32, hidden_dim)\n",
    "        self.fc21 = nn.Linear(hidden_dim, y_dim)\n",
    "        # setup the non-linearities\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        # define the forward computation on the image x\n",
    "        # first shape the mini-batch to have pixels in the rightmost dimension\n",
    "        x = self.relu(self.conv1(x))\n",
    "        x = self.relu(self.conv2(x))\n",
    "        x = x.view(-1, 6*6*32)\n",
    "        # then compute the hidden units\n",
    "        hidden = self.relu(self.fc1(x))\n",
    "        # then return a mean vector and a (positive) square root covariance\n",
    "        # each of size batch_size x z_dim\n",
    "        return F.softmax(self.fc21(hidden),dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SSVAE(nn.Module):\n",
    "    \"\"\"\n",
    "    This class encapsulates the parameters (neural networks) and models & guides needed to train a\n",
    "    semi-supervised variational auto-encoder on the MNIST image dataset\n",
    "    :param output_size: size of the tensor representing the class label (10 for MNIST since\n",
    "                        we represent the class labels as a one-hot vector with 10 components)\n",
    "    :param input_size: size of the tensor representing the image (28*28 = 784 for our MNIST dataset\n",
    "                       since we flatten the images and scale the pixels to be in [0,1])\n",
    "    :param z_dim: size of the tensor representing the latent random variable z\n",
    "                  (handwriting style for our MNIST dataset)\n",
    "    :param hidden_layers: a tuple (or list) of MLP layers to be used in the neural networks\n",
    "                          representing the parameters of the distributions in our model\n",
    "    :param use_cuda: use GPUs for faster training\n",
    "    :param aux_loss_multiplier: the multiplier to use with the auxiliary loss\n",
    "    \"\"\"\n",
    "    def __init__(self, output_size=10, input_size=784, z_dim=50, hidden_layers=500,\n",
    "                 config_enum=None, use_cuda=False, aux_loss_multiplier=None):\n",
    "\n",
    "        super().__init__()\n",
    "\n",
    "        # initialize the class with all arguments provided to the constructor\n",
    "        self.output_size = output_size\n",
    "        self.input_size = input_size\n",
    "        self.z_dim = z_dim\n",
    "        self.hidden_layers = hidden_layers\n",
    "        self.allow_broadcast = config_enum == 'parallel'\n",
    "        self.use_cuda = use_cuda\n",
    "        self.aux_loss_multiplier = aux_loss_multiplier\n",
    "\n",
    "        # define and instantiate the neural networks representing\n",
    "        # the paramters of various distributions in the model\n",
    "        self.setup_networks()\n",
    "\n",
    "    def setup_networks(self):\n",
    "\n",
    "        self.encoder_y = Classifier(self.output_size, self.hidden_layers)\n",
    "\n",
    "        self.encoder_z = Encoder(self.z_dim, self.hidden_layers, self.output_size)\n",
    "\n",
    "        self.decoder = Decoder(self.z_dim+self.output_size, self.hidden_layers)\n",
    "\n",
    "        if self.use_cuda:\n",
    "            self.cuda()\n",
    "\n",
    "    def model(self, xs, ys=None):\n",
    "        \"\"\"\n",
    "        The model corresponds to the following generative process:\n",
    "        p(z) = normal(0,I)              # handwriting style (latent)\n",
    "        p(y|x) = categorical(I/10.)     # which digit (semi-supervised)\n",
    "        p(x|y,z) = bernoulli(loc(y,z))   # an image\n",
    "        loc is given by a neural network  `decoder`\n",
    "        :param xs: a batch of scaled vectors of pixels from an image\n",
    "        :param ys: (optional) a batch of the class labels i.e.\n",
    "                   the digit corresponding to the image(s)\n",
    "        :return: None\n",
    "        \"\"\"\n",
    "        # register this pytorch module and all of its sub-modules with pyro\n",
    "        pyro.module(\"ss_vae\", self)\n",
    "\n",
    "        batch_size = xs.size(0)\n",
    "        options = dict(dtype=xs.dtype, device=xs.device)\n",
    "        with pyro.plate(\"data\"):\n",
    "\n",
    "            # sample the handwriting style from the constant prior distribution\n",
    "            prior_loc = torch.zeros(batch_size, self.z_dim, **options)\n",
    "            prior_scale = torch.ones(batch_size, self.z_dim, **options)\n",
    "            zs = pyro.sample(\"z\", dist.Normal(prior_loc, prior_scale).to_event(1))\n",
    "\n",
    "            # if the label y (which digit to write) is supervised, sample from the\n",
    "            # constant prior, otherwise, observe the value (i.e. score it against the constant prior)\n",
    "            alpha_prior = torch.ones(batch_size, self.output_size, **options) / (1.0 * self.output_size)\n",
    "            ys = pyro.sample(\"y\", dist.OneHotCategorical(alpha_prior), obs=ys)\n",
    "\n",
    "            # finally, score the image (x) using the handwriting style (z) and\n",
    "            # the class label y (which digit to write) against the\n",
    "            # parametrized distribution p(x|y,z) = bernoulli(decoder(y,z))\n",
    "            # where `decoder` is a neural network\n",
    "            loc = self.decoder.forward(torch.cat([zs, ys], axis=1))\n",
    "            pyro.sample(\"x\", dist.Bernoulli(loc).to_event(1), obs=xs.reshape(-1, 784))\n",
    "            # return the loc so we can visualize it later\n",
    "            return loc\n",
    "\n",
    "    def guide(self, xs, ys=None):\n",
    "        \"\"\"\n",
    "        The guide corresponds to the following:\n",
    "        q(y|x) = categorical(alpha(x))              # infer digit from an image\n",
    "        q(z|x,y) = normal(loc(x,y),scale(x,y))       # infer handwriting style from an image and the digit\n",
    "        loc, scale are given by a neural network `encoder_z`\n",
    "        alpha is given by a neural network `encoder_y`\n",
    "        :param xs: a batch of scaled vectors of pixels from an image\n",
    "        :param ys: (optional) a batch of the class labels i.e.\n",
    "                   the digit corresponding to the image(s)\n",
    "        :return: None\n",
    "        \"\"\"\n",
    "        # inform Pyro that the variables in the batch of xs, ys are conditionally independent\n",
    "        with pyro.plate(\"data\"):\n",
    "\n",
    "            # if the class label (the digit) is not supervised, sample\n",
    "            # (and score) the digit with the variational distribution\n",
    "            # q(y|x) = categorical(alpha(x))\n",
    "            if ys is None:\n",
    "                alpha = self.encoder_y.forward(xs)\n",
    "                ys = pyro.sample(\"y\", dist.OneHotCategorical(alpha))\n",
    "\n",
    "            # sample (and score) the latent handwriting-style with the variational\n",
    "            # distribution q(z|x,y) = normal(loc(x,y),scale(x,y))\n",
    "            loc, scale = self.encoder_z.forward(xs, ys)\n",
    "            pyro.sample(\"z\", dist.Normal(loc, scale).to_event(1))\n",
    "\n",
    "    def classifier(self, xs):\n",
    "        \"\"\"\n",
    "        classify an image (or a batch of images)\n",
    "        :param xs: a batch of scaled vectors of pixels from an image\n",
    "        :return: a batch of the corresponding class labels (as one-hots)\n",
    "        \"\"\"\n",
    "        # use the trained model q(y|x) = categorical(alpha(x))\n",
    "        # compute all class probabilities for the image(s)\n",
    "        alpha = self.encoder_y.forward(xs)\n",
    "\n",
    "        # get the index (digit) that corresponds to\n",
    "        # the maximum predicted class probability\n",
    "        res, ind = torch.topk(alpha, 1)\n",
    "\n",
    "        # convert the digit(s) to one-hot tensor(s)\n",
    "        ys = torch.zeros_like(alpha).scatter_(1, ind, 1.0)\n",
    "        return ys\n",
    "\n",
    "    def model_classify(self, xs, ys=None):\n",
    "        \"\"\"\n",
    "        this model is used to add an auxiliary (supervised) loss as described in the\n",
    "        Kingma et al., \"Semi-Supervised Learning with Deep Generative Models\".\n",
    "        \"\"\"\n",
    "        # register all pytorch (sub)modules with pyro\n",
    "        pyro.module(\"ss_vae\", self)\n",
    "\n",
    "        # inform Pyro that the variables in the batch of xs, ys are conditionally independent\n",
    "        with pyro.plate(\"data\"):\n",
    "            # this here is the extra term to yield an auxiliary loss that we do gradient descent on\n",
    "            if ys is not None:\n",
    "                alpha = self.encoder_y.forward(xs)\n",
    "                with pyro.poutine.scale(scale=self.aux_loss_multiplier):\n",
    "                    pyro.sample(\"y_aux\", dist.OneHotCategorical(alpha), obs=ys)\n",
    "\n",
    "    def guide_classify(self, xs, ys=None):\n",
    "        \"\"\"\n",
    "        dummy guide function to accompany model_classify in inference\n",
    "        \"\"\"\n",
    "        pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_accuracy(data_loader, classifier_fn, batch_size):\n",
    "    \"\"\"\n",
    "    compute the accuracy over the supervised training set or the testing set\n",
    "    \"\"\"\n",
    "    predictions, actuals = [], []\n",
    "\n",
    "    # use the appropriate data loader\n",
    "    for (xs, ys) in data_loader:\n",
    "        # use classification function to compute all predictions for each batch\n",
    "        predictions.append(classifier_fn(xs.cuda()))\n",
    "        actuals.append(ys.cuda())\n",
    "\n",
    "    # compute the number of accurate predictions\n",
    "    accurate_preds = 0\n",
    "    for pred, act in zip(predictions, actuals):\n",
    "        for i in range(pred.size(0)):\n",
    "            v = torch.sum(pred[i] == act[i])\n",
    "            accurate_preds += (v.item() == 10)\n",
    "\n",
    "    # calculate the accuracy between 0 and 1\n",
    "    accuracy = (accurate_preds * 1.0) / (len(predictions) * batch_size)\n",
    "    return accuracy\n",
    "\n",
    "\n",
    "def visualize(ss_vae, viz, test_loader):\n",
    "    if viz:\n",
    "        plot_conditional_samples_ssvae(ss_vae, viz)\n",
    "        mnist_test_tsne_ssvae(ssvae=ss_vae, test_loader=test_loader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_inference_for_epoch(data_loaders, losses, periodic_interval_batches):\n",
    "    \"\"\"\n",
    "    runs the inference algorithm for an epoch\n",
    "    returns the values of all losses separately on supervised and unsupervised parts\n",
    "    \"\"\"\n",
    "    num_losses = len(losses)\n",
    "\n",
    "    # compute number of batches for an epoch\n",
    "    sup_batches = len(data_loaders[\"sup\"])\n",
    "    unsup_batches = len(data_loaders[\"unsup\"])\n",
    "    batches_per_epoch = sup_batches + unsup_batches\n",
    "\n",
    "    # initialize variables to store loss values\n",
    "    epoch_losses_sup = [0.] * num_losses\n",
    "    epoch_losses_unsup = [0.] * num_losses\n",
    "\n",
    "    # setup the iterators for training data loaders\n",
    "    sup_iter = iter(data_loaders[\"sup\"])\n",
    "    unsup_iter = iter(data_loaders[\"unsup\"])\n",
    "\n",
    "    # count the number of supervised batches seen in this epoch\n",
    "    ctr_sup = 0\n",
    "    for i in range(batches_per_epoch):\n",
    "\n",
    "        # whether this batch is supervised or not\n",
    "        is_supervised = (i % periodic_interval_batches == 1) and ctr_sup < sup_batches\n",
    "\n",
    "        # extract the corresponding batch\n",
    "        if is_supervised:\n",
    "            (xs, ys) = next(sup_iter)\n",
    "            ctr_sup += 1\n",
    "        else:\n",
    "            xs= next(unsup_iter)\n",
    "\n",
    "        # run the inference for each loss with supervised or un-supervised\n",
    "        # data as arguments\n",
    "        for loss_id in range(num_losses):\n",
    "            if is_supervised:\n",
    "                #print(\"supervised\")\n",
    "                new_loss = losses[loss_id].step(xs.cuda(), ys.cuda())\n",
    "                epoch_losses_sup[loss_id] += new_loss\n",
    "            else:\n",
    "                #print(\"unsupervised\")\n",
    "                new_loss = losses[loss_id].step(xs.cuda())\n",
    "                #print(type(new_loss))\n",
    "                epoch_losses_unsup[loss_id] += new_loss\n",
    "\n",
    "    # return the values of all losses\n",
    "    return epoch_losses_sup, epoch_losses_unsup\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_set = MNIST(root='./data', train=True, transform=transforms.ToTensor(),\n",
    "                           download=True)\n",
    "test_set = MNIST(root='./data', train=False, transform=transforms.ToTensor())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train = train_set.data.float()/255\n",
    "y_train = One_Hot(10)(train_set.targets)\n",
    "\n",
    "X_test = test_set.data.float()/255\n",
    "y_test = One_Hot(10)(test_set.targets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3000, 10])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y_train[:3000].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Defining model and inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ss_vae = SSVAE(output_size=10, input_size=784, z_dim=50, hidden_layers=500,\n",
    "                 config_enum=\"parallel\", use_cuda=True, aux_loss_multiplier=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "adam_params = {\"lr\": 0.001}\n",
    "optimizer = Adam(adam_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ELBO of main loss $\\mathcal{J} = \\sum_{sup} \\mathcal{L}(\\vec{x},y) + \\sum_{unsup} \\mathcal{U}(\\vec{x})$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "guide = config_enumerate(ss_vae.guide, \"sequential\", expand=True) #config_enumerate(ss_vae.guide, \"parallel\", expand=True)\n",
    "elbo = TraceEnum_ELBO(max_plate_nesting=1) #JitTraceEnum_ELBO\n",
    "loss_basic = SVI(ss_vae.model, guide, optimizer, loss=elbo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = [loss_basic]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## adding auxillary loss $\\mathbb{E}_{sup}[- \\log q(y|\\vec{x})]$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "elbo = Trace_ELBO()\n",
    "loss_aux = SVI(ss_vae.model_classify, ss_vae.guide_classify, optimizer, loss=elbo)\n",
    "losses.append(loss_aux)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "sup_num= 10000\n",
    "periodic_interval_batches = int(50000 / (1.0 * sup_num))\n",
    "unsup_num = 50000-sup_num\n",
    "best_valid_acc, corresponding_test_acc = 0.0, 0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_loaders = {}\n",
    "data_loaders[\"sup\"] = torch.utils.data.DataLoader(SupervisedDataset(X_train[:sup_num].view(-1, 1,28,28), y_train[:sup_num]), \n",
    "                                                  batch_size=batch_size, shuffle=True)\n",
    "data_loaders[\"unsup\"] = torch.utils.data.DataLoader(X_train[sup_num:].view(-1, 1,28,28), batch_size=batch_size, shuffle=True)\n",
    "data_loaders[\"valid\"] = torch.utils.data.DataLoader(SupervisedDataset(X_train[-5000:].view(-1, 1,28,28), y_train[-5000:]), \n",
    "                                                  batch_size=batch_size)\n",
    "data_loaders[\"test\"] = torch.utils.data.DataLoader(SupervisedDataset(X_test.view(-1, 1,28,28), y_test), batch_size=batch_size)                                                  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 epoch: avg losses 229.6227453125 1.7494861633300782 269.85746126708983 0.0 validation accuracy 0.6806 test accuracy 0.6338\n",
      "1 epoch: avg losses 157.86091171875 0.7604765869140625 192.82245803222656 0.0 validation accuracy 0.8422 test accuracy 0.8023\n",
      "2 epoch: avg losses 130.988081640625 0.5768449447631836 161.7880133605957 0.0 validation accuracy 0.8596 test accuracy 0.8395\n",
      "3 epoch: avg losses 118.9936546875 0.5568067138671875 147.9520163330078 0.0 validation accuracy 0.889 test accuracy 0.8712\n",
      "4 epoch: avg losses 113.320178515625 0.44935596160888674 140.98933096313476 0.0 validation accuracy 0.8974 test accuracy 0.8809\n",
      "5 epoch: avg losses 109.839412890625 0.38662865295410154 137.13567880249025 0.0 validation accuracy 0.9144 test accuracy 0.8993\n",
      "6 epoch: avg losses 107.64010390625 0.3421112747192383 134.33126782836914 0.0 validation accuracy 0.9232 test accuracy 0.9129\n",
      "7 epoch: avg losses 105.805780078125 0.28438831253051755 132.3639467895508 0.0 validation accuracy 0.9426 test accuracy 0.9338\n",
      "8 epoch: avg losses 104.449015234375 0.2426683380126953 130.74628061523438 0.0 validation accuracy 0.9478 test accuracy 0.9406\n",
      "9 epoch: avg losses 103.592869140625 0.21584069900512695 129.51352569274903 0.0 validation accuracy 0.9502 test accuracy 0.9453\n",
      "10 epoch: avg losses 102.681369140625 0.18890862884521484 128.3063039855957 0.0 validation accuracy 0.9548 test accuracy 0.9483\n",
      "11 epoch: avg losses 101.72181875 0.15774828491210938 127.40493998413086 0.0 validation accuracy 0.961 test accuracy 0.9519\n",
      "12 epoch: avg losses 101.203494140625 0.15751935577392578 126.61635982666016 0.0 validation accuracy 0.9644 test accuracy 0.9567\n",
      "13 epoch: avg losses 100.515605078125 0.13562992401123047 125.83640650024414 0.0 validation accuracy 0.9662 test accuracy 0.9564\n",
      "14 epoch: avg losses 99.91789921875 0.1271669387817383 125.11068839111329 0.0 validation accuracy 0.9718 test accuracy 0.961\n",
      "15 epoch: avg losses 99.47421796875 0.11929012699127198 124.56831834716797 0.0 validation accuracy 0.9688 test accuracy 0.9617\n",
      "16 epoch: avg losses 99.079314453125 0.10292992115020752 124.05335877685548 0.0 validation accuracy 0.9734 test accuracy 0.9674\n",
      "17 epoch: avg losses 98.90828828125 0.09186105737686157 123.62649692382813 0.0 validation accuracy 0.9744 test accuracy 0.9673\n",
      "18 epoch: avg losses 98.202844140625 0.09245932331085205 123.1244541809082 0.0 validation accuracy 0.9732 test accuracy 0.967\n",
      "19 epoch: avg losses 98.0227171875 0.08455673370361329 122.7360176574707 0.0 validation accuracy 0.978 test accuracy 0.9709\n",
      "20 epoch: avg losses 97.784248046875 0.07482936449050903 122.43190952758789 0.0 validation accuracy 0.9792 test accuracy 0.9733\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-19-1ff922642819>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[1;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;36m100\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 2\u001b[1;33m     \u001b[0mepoch_losses_sup\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mepoch_losses_unsup\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mrun_inference_for_epoch\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdata_loaders\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mlosses\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mperiodic_interval_batches\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      3\u001b[0m     \u001b[1;31m# compute average epoch losses i.e. losses per example\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      5\u001b[0m     \u001b[0mavg_epoch_losses_sup\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;32mlambda\u001b[0m \u001b[0mv\u001b[0m\u001b[1;33m:\u001b[0m \u001b[0mv\u001b[0m \u001b[1;33m/\u001b[0m \u001b[0msup_num\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mepoch_losses_sup\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m<ipython-input-7-8462fb606f85>\u001b[0m in \u001b[0;36mrun_inference_for_epoch\u001b[1;34m(data_loaders, losses, periodic_interval_batches)\u001b[0m\n\u001b[0;32m     42\u001b[0m             \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     43\u001b[0m                 \u001b[1;31m#print(\"unsupervised\")\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 44\u001b[1;33m                 \u001b[0mnew_loss\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mlosses\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mloss_id\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mxs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     45\u001b[0m                 \u001b[1;31m#print(type(new_loss))\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     46\u001b[0m                 \u001b[0mepoch_losses_unsup\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mloss_id\u001b[0m\u001b[1;33m]\u001b[0m \u001b[1;33m+=\u001b[0m \u001b[0mnew_loss\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Miniconda3\\envs\\Pytorch\\lib\\site-packages\\pyro\\infer\\svi.py\u001b[0m in \u001b[0;36mstep\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    133\u001b[0m         \u001b[1;31m# actually perform gradient steps\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    134\u001b[0m         \u001b[1;31m# torch.optim objects gets instantiated for any params that haven't been seen yet\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 135\u001b[1;33m         \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moptim\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mparams\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    136\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    137\u001b[0m         \u001b[1;31m# zero gradients\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Miniconda3\\envs\\Pytorch\\lib\\site-packages\\pyro\\optim\\optim.py\u001b[0m in \u001b[0;36m__call__\u001b[1;34m(self, params, *args, **kwargs)\u001b[0m\n\u001b[0;32m     76\u001b[0m                 \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moptim_objs\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mp\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moptimizer\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     77\u001b[0m             \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 78\u001b[1;33m                 \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0moptim_objs\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mp\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m*\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;33m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     79\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     80\u001b[0m     \u001b[1;32mdef\u001b[0m \u001b[0mget_state\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Miniconda3\\envs\\Pytorch\\lib\\site-packages\\torch\\optim\\adam.py\u001b[0m in \u001b[0;36mstep\u001b[1;34m(self, closure)\u001b[0m\n\u001b[0;32m     93\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     94\u001b[0m                 \u001b[1;31m# Decay the first and second moment running average coefficient\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 95\u001b[1;33m                 \u001b[0mexp_avg\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmul_\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mbeta1\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0madd_\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m \u001b[1;33m-\u001b[0m \u001b[0mbeta1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgrad\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     96\u001b[0m                 \u001b[0mexp_avg_sq\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmul_\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mbeta2\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0maddcmul_\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;36m1\u001b[0m \u001b[1;33m-\u001b[0m \u001b[0mbeta2\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgrad\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mgrad\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     97\u001b[0m                 \u001b[1;32mif\u001b[0m \u001b[0mamsgrad\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "for i in range(0, 100):\n",
    "    epoch_losses_sup, epoch_losses_unsup = run_inference_for_epoch(data_loaders, losses, periodic_interval_batches)\n",
    "    # compute average epoch losses i.e. losses per example\n",
    "    \n",
    "    avg_epoch_losses_sup = map(lambda v: v / sup_num, epoch_losses_sup)\n",
    "    avg_epoch_losses_unsup = map(lambda v: v / unsup_num, epoch_losses_unsup)\n",
    "    \n",
    "    # store the loss and validation/testing accuracies in the logfile\n",
    "    str_loss_sup = \" \".join(map(str, avg_epoch_losses_sup))\n",
    "    str_loss_unsup = \" \".join(map(str, avg_epoch_losses_unsup))\n",
    "\n",
    "    str_print = \"{} epoch: avg losses {}\".format(i, \"{} {}\".format(str_loss_sup, str_loss_unsup))\n",
    "    \n",
    "    validation_accuracy = get_accuracy(data_loaders[\"valid\"], ss_vae.classifier, batch_size)\n",
    "    str_print += \" validation accuracy {}\".format(validation_accuracy)\n",
    "    \n",
    "    # this test accuracy is only for logging, this is not used\n",
    "    # to make any decisions during training\n",
    "    test_accuracy = get_accuracy(data_loaders[\"test\"], ss_vae.classifier, batch_size)\n",
    "    str_print += \" test accuracy {}\".format(test_accuracy)\n",
    "    \n",
    "    print(str_print)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_test_accuracy = get_accuracy(data_loaders[\"test\"], ss_vae.classifier, batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ss_vae.encoder_z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = next(iter(data_loaders[\"unsup\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp[1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#ss_vae.encoder_z(tmp.view(-1,1,28,28).cuda(),torch.cat(200*[torch.tensor([[0,0,0,0,0,0,0,0,0,1]]).float().cuda()]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ss_vae.encoder_y(tmp.view(-1,1,28,28).cuda()).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}