construct_network.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "_uuid": "05778375fb32b8b3a0a2e5b72822f6db8061e7fb",
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "64 1 1\n",
      "CNN_VAE(\n",
      "  (Encoder_conv): ModuleList(\n",
      "    (0): Conv2d(1, 8, kernel_size=(5, 5), stride=(1, 1))\n",
      "    (1): BatchNorm2d(8, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (2): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
      "    (3): Conv2d(8, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
      "    (4): BatchNorm2d(32, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (5): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
      "    (6): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (7): BatchNorm2d(64, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (8): MaxPool2d(kernel_size=3, stride=3, padding=0, dilation=1, ceil_mode=False)\n",
      "  )\n",
      "  (Encoder_FC): ModuleList(\n",
      "    (0): Linear(in_features=64, out_features=1000, bias=True)\n",
      "    (1): BatchNorm1d(1000, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (2): Linear(in_features=1000, out_features=200, bias=True)\n",
      "    (3): BatchNorm1d(200, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (Additional_layer): ModuleList(\n",
      "    (0): Linear(in_features=202, out_features=256, bias=True)\n",
      "    (1): BatchNorm1d(256, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (Decoder_FC): ModuleList(\n",
      "    (0): Linear(in_features=130, out_features=200, bias=True)\n",
      "    (1): BatchNorm1d(200, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (2): Linear(in_features=200, out_features=1000, bias=True)\n",
      "    (3): BatchNorm1d(1000, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (4): Linear(in_features=1000, out_features=64, bias=True)\n",
      "    (5): BatchNorm1d(64, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (Decoder_conv): ModuleList(\n",
      "    (0): ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
      "    (1): BatchNorm2d(32, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (2): ConvTranspose2d(32, 8, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))\n",
      "    (3): BatchNorm2d(8, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (4): ConvTranspose2d(8, 2, kernel_size=(5, 5), stride=(1, 1))\n",
      "    (5): BatchNorm2d(2, eps=1e-05, momentum=0.2, affine=True, track_running_stats=True)\n",
      "  )\n",
      "  (Classifier): ModuleList(\n",
      "    (0): Linear(in_features=200, out_features=1000, bias=True)\n",
      "    (1): BatchNorm1d(1000, eps=0.0001, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (2): Linear(in_features=1000, out_features=200, bias=True)\n",
      "    (3): BatchNorm1d(200, eps=0.0001, momentum=0.2, affine=True, track_running_stats=True)\n",
      "    (4): Linear(in_features=200, out_features=2, bias=True)\n",
      "  )\n",
      ")\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/site-packages/ipykernel_launcher.py:336: UserWarning: nn.init.xavier_normal is now deprecated in favor of nn.init.xavier_normal_.\n"
     ]
    }
   ],
   "source": [
    "warnings.filterwarnings(\"ignore\")\n",
    "\n",
    "## Define size variables\n",
    "print_shapes = False\n",
    "height = IMG_SIZE\n",
    "width = IMG_SIZE\n",
    "channels = 1\n",
    "num_features = height*width*channels\n",
    "\n",
    "# Regulization\n",
    "L2_reg = 1e-6\n",
    "DROPOUT = True\n",
    "do_p = 0.05 # do_p for conv \n",
    "do_p2 = 0.1 # do_p for linear   NB: Classifier dropout is set manuel to 0.3\n",
    "batchnorm_eps = 1e-5\n",
    "batchnorm_momentum = 0.2\n",
    "\n",
    "# Conv Layers\n",
    "conv_out_channels = [8, 32, 64]\n",
    "conv_kernel = [5, 5, 3]\n",
    "conv_padding = [0, 2, 1]\n",
    "conv_stride = [1, 1, 1]\n",
    "\n",
    "# MaxPool Layers\n",
    "pool_kernel = 3\n",
    "pool_padding = 0\n",
    "pool_stride = 3\n",
    "\n",
    "# Fully connected layers\n",
    "lin_layer = [1000, 200]\n",
    "\n",
    "# auxillary parameters\n",
    "aux_layer = [200, 200]\n",
    "aux_variables = 0\n",
    "aux_in = 2 # layer no. where a is included in encoder\n",
    "aux_decoder_layers = [200,200]\n",
    "\n",
    "# classifier parameters\n",
    "classifier_layer = [1000,200]\n",
    "No_classes = len(classes)\n",
    "\n",
    "# No. of layes\n",
    "NUM_CONV = len(conv_out_channels)\n",
    "NUM_LIN = len(lin_layer)\n",
    "NUM_AUX = len(aux_layer)\n",
    "NUM_CLASS = len(classifier_layer)\n",
    "NUM_AUX_DECODER = len(aux_decoder_layers)\n",
    "\n",
    "# Calculating the dimensions \n",
    "def compute_conv_dim(height, width, kernel_size, padding_size, stride_size):\n",
    "    height_new = int((height - kernel_size + 2 * padding_size) / stride_size + 1)\n",
    "    width_new =  int((width  - kernel_size + 2 * padding_size) / stride_size + 1)\n",
    "    return [height_new, width_new]\n",
    "\n",
    "def compute_final_dimension(height, width, last_num_channels, num_layers):\n",
    "    # First conv layer\n",
    "    CNN_height = height\n",
    "    CNN_width = width\n",
    "    for i in range(num_layers):\n",
    "        # conv layer\n",
    "        CNN_height, CNN_width = compute_conv_dim(CNN_height, CNN_width, conv_kernel[i], conv_padding[i], conv_stride[i])\n",
    "        # maxpool layer\n",
    "        CNN_height, CNN_width = compute_conv_dim(CNN_height, CNN_width, pool_kernel, pool_padding, pool_stride)\n",
    "    final_dim = CNN_height * CNN_width * last_num_channels\n",
    "    # print(final_dim,CNN_height,CNN_width)\n",
    "    return [final_dim, CNN_height, CNN_width]\n",
    "\n",
    "def normalize(x):\n",
    "    tmp = x-torch.min( torch.min(x,dim = 2, keepdim = True)[0] ,dim = 3, keepdim = True)[0]\n",
    "    if torch.sum(torch.isnan(tmp))>0:\n",
    "        print(\"nan of tmp\",torch.sum(torch.isnan(tmp)))\n",
    "    return tmp/(torch.max( torch.max(tmp,dim = 2, keepdim = True)[0] ,dim = 3, keepdim = True)[0] + 1e-8)  \n",
    "\n",
    "def gaussian_sample(mu,log_var, num_samples, latent_features):    \n",
    "    # Don't propagate gradients through randomness\n",
    "    with torch.no_grad():\n",
    "        batch_size = mu.size(0)\n",
    "        epsilon = torch.randn(batch_size, num_samples, latent_features)\n",
    "            \n",
    "    if cuda:\n",
    "        epsilon = epsilon.cuda()\n",
    "        \n",
    "    sigma = torch.exp(log_var/2)\n",
    "        \n",
    "    # We will need to unsqueeze to turn\n",
    "    # (batch_size, latent_dim) -> (batch_size, 1, latent_dim)\n",
    "    if len(mu.shape) == 2:\n",
    "        z = mu.unsqueeze(1) + epsilon * sigma.unsqueeze(1)\n",
    "    else:\n",
    "        z = mu + epsilon * sigma\n",
    "    return z\n",
    "\n",
    "def output_recon(x):\n",
    "    # Shape of x_mean: [batch_size, num_samples, channel, height, width]\n",
    "    x_mean, x_log_var = torch.chunk(x, 2, dim=2) # the mean and log_var reconstructions from the decoder\n",
    "    \n",
    "    # The original digits are on the scale [0, 1] \n",
    "    x_hat = x_mean[:,1,].unsqueeze(1)\n",
    "    #x_hat = normalize(x_mean[:,1,].unsqueeze(1))# to scale for showing an image\n",
    "    #x_hat = normalize(x_mean)\n",
    "    x_log_var = softplus(x_log_var)\n",
    "    \n",
    "    # Mean over samples\n",
    "    #x_hat = torch.mean(x_hat, dim=1)\n",
    "    x_log_var= torch.mean(x_log_var, dim=1)\n",
    "    x_mean = torch.mean(x_mean,dim=1) # used for the loss\n",
    "    \n",
    "    # Resize x_hat from [batch_size, no_features] to [batch_size, channels, height, width]\n",
    "    x_hat = x_hat.view( batch_size, 1, height, width)\n",
    "    x_log_var = x_log_var.view( batch_size, 1, height, width)\n",
    "    return x_hat, x_log_var, x_mean\n",
    "\n",
    "######## Image has to be: (num, channels, height, width)!!!! #########\n",
    "class CNN_VAE(nn.Module):\n",
    "    def encoder(self,x):\n",
    "        # Convolutional layers of encoder\n",
    "        for i in range(0,len(self.Encoder_conv),3):\n",
    "            x = self.Encoder_conv[i](x) # Convolutional layer\n",
    "            self.layer_size.append(x.shape[-1])\n",
    "            x = self.Encoder_conv[i+1](x) # Batchnorm layer\n",
    "            x = relu(x)\n",
    "            if DROPOUT:\n",
    "                x = dropout2d(x, p=do_p)   \n",
    "            x = self.Encoder_conv[i+2](x) # Maxpool Layer\n",
    "        x = x.view(batch_size, -1) # Prepare x for linear layers\n",
    "        \n",
    "        # Fully connected layers of encoder\n",
    "        for i in range(0,len(self.Encoder_FC),2): \n",
    "            x = self.Encoder_FC[i](x) # Linear layer\n",
    "            x = self.Encoder_FC[i+1](x) # Batchnorm\n",
    "            x = relu(x)\n",
    "            if DROPOUT:\n",
    "                x = dropout(x, p=do_p2)\n",
    "        return x\n",
    "\n",
    "    \n",
    "    def decoder(self,z,y):\n",
    "        x = torch.cat([z,y],dim=-1)\n",
    "        # Fully connected layers of decoder\n",
    "        for i in range(0,len(self.Decoder_FC),2):\n",
    "            x = self.Decoder_FC[i](x)\n",
    "            x = x.permute(0,2,1)\n",
    "            x = self.Decoder_FC[i+1](x)\n",
    "            x = x.permute(0,2,1)\n",
    "            x = relu(x)\n",
    "            x = dropout(x,p= do_p2)\n",
    "        x = x.view(-1, self.Decoder_conv[0].in_channels, self.final_dim[1], self.final_dim[2])\n",
    "        \n",
    "        # Convolutional layers of decoder\n",
    "        curr_layer = len(self.Decoder_conv)//2-1\n",
    "        for i in range(0,len(self.Decoder_conv),2):\n",
    "            x = interpolate(x,size = [self.layer_size[curr_layer],self.layer_size[curr_layer]],\n",
    "                                      mode = 'bilinear', \n",
    "                                      align_corners = False)\n",
    "            curr_layer -=1\n",
    "            x = self.Decoder_conv[i](x) # Convolutional layers\n",
    "            x = self.Decoder_conv[i+1](x) # BatchNorm\n",
    "            x = relu(x)\n",
    "            if DROPOUT:\n",
    "                x = dropout2d(x, p=do_p)\n",
    "        return x.view(batch_size,-1,channels*2,height,width)\n",
    "            \n",
    "    def encoder_aux(self,a):\n",
    "        for i in range(0,len(self.Encoder_aux),2):\n",
    "            a = self.Encoder_aux[i](a)\n",
    "            a = self.Encoder_aux[i+1](a)\n",
    "            a = relu(a)\n",
    "            if DROPOUT:\n",
    "                a = dropout(a, p=do_p2)\n",
    "        q_a_mu, q_a_log_var = torch.chunk(a, 2, dim=-1) # divide to mu and sigma\n",
    "        return q_a_mu, q_a_log_var\n",
    "    \n",
    "    def decoder_aux(self,xz,y):\n",
    "        a = torch.cat([xz,y],dim=-1)\n",
    "        for i in range(0,len(self.Decoder_aux),2):\n",
    "            a = self.Decoder_aux[i](a)\n",
    "            a = a.permute(0,2,1)\n",
    "            a = self.Decoder_aux[i+1](a)\n",
    "            a = a.permute(0,2,1)\n",
    "            a = relu(a)\n",
    "            if DROPOUT:\n",
    "                a = dropout(a, p=do_p2)  \n",
    "        return a\n",
    "    \n",
    "    def additional_layer(self,xa,y):\n",
    "        z = torch.cat([xa,y],dim=-1)\n",
    "        z = self.Additional_layer[0](z)\n",
    "        z = z.permute(0,2,1)\n",
    "        z = self.Additional_layer[1](z)\n",
    "        z = z.permute(0,2,1)\n",
    "        z = relu(z)\n",
    "        if DROPOUT:\n",
    "            z = dropout(z, p=do_p2) \n",
    "        return z\n",
    "    \n",
    "    def classifier(self,xa):\n",
    "        for i in range(0,len(self.Classifier),2):\n",
    "            xa = self.Classifier[i](xa)\n",
    "            if i < len(self.Classifier)-1:\n",
    "                if aux_variables > 0:\n",
    "                    xa = xa.permute(0,2,1)\n",
    "                    xa = self.Classifier[i+1](xa)\n",
    "                    xa = xa.permute(0,2,1)\n",
    "                else:\n",
    "                    xa = self.Classifier[i+1](xa)\n",
    "                xa = relu(xa)\n",
    "                if DROPOUT:\n",
    "                    xa = dropout(xa, p=0.3)\n",
    "        return softmax(xa,dim=-1)\n",
    "    \n",
    "    def sample_y(self,batch_size,num_samples,no_classes,i):\n",
    "        tmp = Variable(torch.zeros(no_classes))\n",
    "        tmp[i] = 1\n",
    "        if cuda:\n",
    "            tmp = tmp.cuda()\n",
    "        return tmp.repeat(batch_size,num_samples,1)\n",
    "    \n",
    "    def sample_from_latent(self,x):\n",
    "        x_UL = []\n",
    "        for j in range(No_classes):\n",
    "            tmp = self.decoder(x.unsqueeze(1).repeat(1,num_samples,1), self.sample_y(batch_size,num_samples,No_classes,j))\n",
    "            x_UL.append(tmp)\n",
    "        x_hat, _, _ = output_recon(sum(x_UL))\n",
    "        return x_hat\n",
    "  \n",
    "    \n",
    "    def __init__(self, latent_features, num_samples):\n",
    "        super(CNN_VAE, self).__init__()\n",
    "        \n",
    "        self.latent_features = latent_features\n",
    "        self.num_samples = num_samples\n",
    "        \n",
    "        # Calculate final size of the CNN\n",
    "        self.final_dim = compute_final_dimension(height,width,conv_out_channels[-1],NUM_CONV)\n",
    "        \n",
    "        ## Convolutional layers of the encoder\n",
    "        input_channels = channels\n",
    "        Encoder_conv = nn.ModuleList()\n",
    "        for i in range(NUM_CONV):\n",
    "            Encoder_conv.append(Conv2d( in_channels=input_channels,\n",
    "                                            out_channels=conv_out_channels[i],\n",
    "                                            kernel_size=conv_kernel[i],\n",
    "                                            stride=conv_stride[i],\n",
    "                                            padding=conv_padding[i]))\n",
    "            Encoder_conv.append(BatchNorm2d(conv_out_channels[i], eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "            Encoder_conv.append(MaxPool2d(  kernel_size=pool_kernel, \n",
    "                                        stride=pool_stride,\n",
    "                                        padding=pool_padding,\n",
    "                                        return_indices = False))\n",
    "            input_channels = conv_out_channels[i]\n",
    "        self.add_module(\"Encoder_conv\",Encoder_conv)\n",
    "        \n",
    "        # Fully connected layers of encoder\n",
    "        Encoder_FC = nn.ModuleList()\n",
    "        in_weights = self.final_dim[0]\n",
    "        for i in range(NUM_LIN):\n",
    "            Encoder_FC.append(Linear(in_features=in_weights, out_features=lin_layer[i]))\n",
    "            Encoder_FC.append(BatchNorm1d(lin_layer[i], eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "            in_weights = lin_layer[i]\n",
    "        self.add_module(\"Encoder_FC\",Encoder_FC)\n",
    "        \n",
    "        # map to latent space\n",
    "        Additional_layer = nn.ModuleList()\n",
    "        Additional_layer.append(Linear(in_features=lin_layer[-1]+aux_variables+No_classes, out_features=latent_features*2))\n",
    "        Additional_layer.append(BatchNorm1d(latent_features*2, eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "        self.add_module(\"Additional_layer\",Additional_layer)\n",
    "        \n",
    "        # Auxillary network\n",
    "        if aux_variables > 0:\n",
    "            # Auxillary encoder\n",
    "            Encoder_aux = nn.ModuleList()\n",
    "            in_weights = lin_layer[-1]\n",
    "            for i in range(NUM_AUX):\n",
    "                Encoder_aux.append(Linear(in_features=in_weights, out_features=aux_layer[i]))\n",
    "                Encoder_aux.append(BatchNorm1d(aux_layer[i], eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "                in_weights = aux_layer[i]\n",
    "            Encoder_aux.append(Linear(in_features=aux_layer[-1], out_features=aux_variables*2))\n",
    "            Encoder_aux.append(BatchNorm1d(aux_variables*2, eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "            self.add_module(\"Encoder_aux\", Encoder_aux)\n",
    "            # Auxillary decoder\n",
    "            Decoder_aux = nn.ModuleList()\n",
    "            for i in range(NUM_AUX_DECODER):\n",
    "                if i == 0:\n",
    "                    in_weights = self.latent_features + lin_layer[-1] + No_classes\n",
    "                else:\n",
    "                    in_weights = aux_decoder_layers[i-1]\n",
    "                Decoder_aux.append(Linear(in_features=in_weights, out_features=aux_decoder_layers[i]))\n",
    "                Decoder_aux.append(BatchNorm1d(aux_decoder_layers[i], eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "            Decoder_aux.append(Linear(in_features=aux_decoder_layers[-1], out_features=aux_variables*2))\n",
    "            Decoder_aux.append(BatchNorm1d(aux_variables*2, eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "            self.add_module(\"Decoder_aux\", Decoder_aux)    \n",
    "        \n",
    "        # Initialize fully connected layers from latent space to convolutional layers\n",
    "        Decoder_FC = nn.ModuleList()\n",
    "        Decoder_FC.append(Linear(in_features=latent_features+No_classes, out_features=lin_layer[-1]))\n",
    "        Decoder_FC.append(BatchNorm1d(lin_layer[-1], eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "        for i in reversed(range(NUM_LIN)):\n",
    "            if i == 0:\n",
    "                out_weights = self.final_dim[0]\n",
    "            else:\n",
    "                out_weights = lin_layer[i-1]\n",
    "            Decoder_FC.append(Linear(in_features=lin_layer[i], out_features=out_weights))\n",
    "            Decoder_FC.append(BatchNorm1d(out_weights, eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "        self.add_module(\"Decoder_FC\",Decoder_FC)\n",
    "        \n",
    "        # Convolutional layers of the decoder\n",
    "        Decoder_conv = nn.ModuleList()\n",
    "        for i in reversed(range(NUM_CONV)):\n",
    "            if i == 0:\n",
    "                output_channels = channels*2\n",
    "            else:\n",
    "                output_channels = conv_out_channels[i-1] \n",
    "            Decoder_conv.append(ConvTranspose2d(in_channels=conv_out_channels[i],\n",
    "                                                out_channels=output_channels,\n",
    "                                                kernel_size=conv_kernel[i],\n",
    "                                                stride=conv_stride[i],\n",
    "                                                padding=conv_padding[i]))\n",
    "            Decoder_conv.append(BatchNorm2d(output_channels, eps = batchnorm_eps, momentum = batchnorm_momentum))\n",
    "        self.add_module(\"Decoder_conv\",Decoder_conv)\n",
    "\n",
    "        # Fully connected layers from convolutional layers to classification\n",
    "        Classifier = nn.ModuleList()\n",
    "        if aux_variables > 0:\n",
    "            in_weights = lin_layer[-1]+aux_variables\n",
    "        else:\n",
    "            in_weights = lin_layer[-1]\n",
    "        for i in range(NUM_CLASS):\n",
    "            Classifier.append(Linear(in_features=in_weights, out_features=classifier_layer[i]))\n",
    "            Classifier.append(BatchNorm1d(classifier_layer[i], eps = 1e-4, momentum = batchnorm_momentum))\n",
    "            in_weights = classifier_layer[i]\n",
    "        Classifier.append(Linear(in_features=classifier_layer[-1], out_features = No_classes))\n",
    "        self.add_module(\"Classifier\", Classifier)\n",
    "        \n",
    "        \n",
    "        # Initialize weight of layers\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Linear) or isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):\n",
    "                nn.init.xavier_normal(m.weight.data)\n",
    "                if m.bias is not None:\n",
    "                    m.bias.data.zero_()\n",
    "                \n",
    "### Forward ####\n",
    "    def forward(self, x, y=None):\n",
    "        outputs = {}\n",
    "        self.indices = []\n",
    "        self.layer_size = []\n",
    "        x = self.encoder(x)\n",
    "        if aux_variables > 0:\n",
    "            q_a_mu, q_a_log_var = self.encoder_aux(x)\n",
    "            q_a = gaussian_sample(q_a_mu,q_a_log_var,num_samples,aux_variables) # sample auxillary variables\n",
    "            outputs[\"q_a\"] = q_a # Assign to outputs\n",
    "            xa = torch.cat([x.unsqueeze(1).repeat(1,num_samples,1),q_a],dim=2) # Create combined vector of x and q_a\n",
    "        else:\n",
    "            xa = x\n",
    "        \n",
    "        # Run trough classifier\n",
    "        logits = self.classifier(xa)\n",
    "        \n",
    "        if aux_variables <= 0:\n",
    "            logits = logits.unsqueeze(1)\n",
    "            logits = logits.repeat(1,num_samples,1)\n",
    "            xa = xa.unsqueeze(1).repeat(1,num_samples,1) \n",
    "            \n",
    "        # Map x, y, a to latent space\n",
    "        if y is None:\n",
    "            x_UL = []\n",
    "            for j in range(No_classes):\n",
    "                z = self.additional_layer(xa,self.sample_y(batch_size,num_samples,No_classes,j))\n",
    "                activation = z\n",
    "                x_UL.append(z)\n",
    "            lat_in = sum(x_UL)\n",
    "            del x_UL\n",
    "        else:\n",
    "            lat_in = self.additional_layer(xa,y.unsqueeze(1).repeat(1,num_samples,1))\n",
    "#             activation = lat_in\n",
    "            \n",
    "        # Split into mu and log_var\n",
    "        mu, log_var = torch.chunk(lat_in, 2, dim=-1)\n",
    "        # Make sure that the log variance is positive\n",
    "        log_var = softplus(log_var)\n",
    "        # Sample from latent space\n",
    "        z = gaussian_sample(mu,log_var,num_samples,latent_features)\n",
    "                \n",
    "        # aux. decoder\n",
    "        if aux_variables > 0:\n",
    "            xz = torch.cat([z, x.unsqueeze(1).repeat(1,z.shape[1],1)],dim = -1)\n",
    "            if y is None:\n",
    "                a_UL = []\n",
    "                for j in range(No_classes):\n",
    "                    a_UL.append(self.decoder_aux(xz,self.sample_y(batch_size,num_samples,No_classes,j)))\n",
    "                a_log_var = []\n",
    "                a_mean = []\n",
    "                for j in range(No_classes):\n",
    "                    tmp1, tmp2 = torch.chunk(a_UL[j], 2, dim=-1) # the mean and log_var reconstructions from the decoder\n",
    "                    a_mean.append(tmp1)\n",
    "                    a_log_var.append(softplus(tmp2))\n",
    "                del a_UL, tmp1, tmp2\n",
    "            else:\n",
    "                a = self.decoder_aux(xz,y.unsqueeze(1).repeat(1,num_samples,1))\n",
    "                a_mean, a_log_var = torch.chunk(a, 2, dim=-1) # the mean and log_var reconstructions from the decoder\n",
    "                a_log_var = softplus(a_log_var)    \n",
    "            \n",
    "        # Decoder     \n",
    "        if y is None:\n",
    "            x_UL = []\n",
    "            for j in range(No_classes):\n",
    "                tmp = self.decoder(z, self.sample_y(batch_size,num_samples,No_classes,j))\n",
    "                x_UL.append(tmp)\n",
    "            x_log_var = []\n",
    "            x_mean = []\n",
    "            x_hat = []\n",
    "            for j in range(No_classes):\n",
    "                tmp1, tmp2, tmp3 = output_recon(x_UL[j])\n",
    "                x_hat.append(tmp1)\n",
    "                x_log_var.append(tmp2)\n",
    "                x_mean.append(tmp3)\n",
    "            del x_UL, tmp1, tmp2, tmp3\n",
    "        else:\n",
    "            x = self.decoder(z, y.unsqueeze(1).repeat(1,num_samples,1))\n",
    "            x_hat, x_log_var, x_mean = output_recon(x)\n",
    "        \n",
    "        # Assign variables\n",
    "        outputs[\"x_hat\"] = x_hat # This is used for visulizations only \n",
    "        outputs[\"z\"] = z\n",
    "        outputs[\"mu\"] = mu\n",
    "        outputs[\"log_var\"] = log_var\n",
    "        \n",
    "        # image recontructions (notice they are outputted as matrices)\n",
    "        outputs[\"x_mean\"] =  x_mean #x_hat  # mean reconstructions (for loss!!!)\n",
    "        outputs[\"x_log_var\"] = x_log_var #torch.reshape(x_log_var,(-1,height,width)) # log var reconstructions (for loss!!!)\n",
    "        \n",
    "        # auxillary outputs\n",
    "        if aux_variables > 0:            \n",
    "            outputs[\"q_a_mu\"] = q_a_mu\n",
    "            outputs[\"q_a_log_var\"] = q_a_log_var\n",
    "            outputs[\"p_a_mu\"] = a_mean\n",
    "            outputs[\"p_a_log_var\"] = a_log_var\n",
    "        \n",
    "        # classifier outputs \n",
    "        outputs[\"y_hat\"] = logits\n",
    "        \n",
    "        # Activation of latent features\n",
    "#         outputs[\"activation\"] = activation\n",
    "\n",
    "        return outputs\n",
    "\n",
    "# The number of samples used then initialising the VAE, \n",
    "# is number of samples drawn from the distribution\n",
    "num_samples = 5\n",
    "latent_features = 32\n",
    "\n",
    "net = CNN_VAE(latent_features, num_samples)\n",
    "\"\"\"\n",
    "print(net)\n",
    "\"\"\"\n",
    "\n",
    "# Transfer model to GPU ifavailable\n",
    "if cuda:\n",
    "    net = net.cuda()\n",
    "\n",
    "# Test\n",
    "if 1 == 0:\n",
    "    x = torch.randn(batch_size,1,width, height)\n",
    "    x = Variable(x)\n",
    "    if cuda:\n",
    "        x = x.cuda()\n",
    "        y = None\n",
    "    y = net(x)\n",
    "    print(y['x_hat'][0].shape)\n",
    "\n",
    "if cuda:\n",
    "    print('before: ',torch.cuda.memory_allocated(device=0))\n",
    "    import gc\n",
    "    #del y,x\n",
    "    # gc.collect()\n",
    "    print('after: ',torch.cuda.memory_allocated(device=0))\n",
    "\n",
    "#for parameter in net.parameters():\n",
    "#    print(parameter.shape)\n",
    "#epsilon = torch.randn(batch_size, latent_features).cuda\n",
    "#samples = torch.sigmoid(net.decoder(net.latent_to_CNN(epsilon))).detach()"
   ]
  }
 ],
 "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}