Testing_normal-loss-2x.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "X06qixDbbZO8"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import tarfile\n",
    "from torchvision.datasets.utils import download_url\n",
    "from torch.utils.data import random_split\n",
    "import os\n",
    "from torchvision.datasets import ImageFolder\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from itertools import chain\n",
    "from tqdm.autonotebook import tqdm\n",
    "import numpy as np\n",
    "from torch import optim\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "A0LwqWD7s6Ww"
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor(),\n",
    "     transforms.Resize(32)])\n",
    "\n",
    "batch_size = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 66,
     "referenced_widgets": [
      "21ca9bd0e04245a0ba6c283caea226ef",
      "896980c70dbd4b0aafdc9b8ed11ca093",
      "c0321ef49e4848efb7a6696e54b18914",
      "4d0e339661be4c47b543e5ec327c13f6",
      "7ad1a00b889f4dd6babb75f7357ddb43",
      "35062c0b25e04cd6b7d8f7d1dc6e64ec",
      "ac6012f09e6e4be0bf4d9c8d8ef5edfb",
      "ddcf1143b7314adcb64d6f64cfbbe462",
      "37d32865a623432cb46c24ac1a044386",
      "d8337d8fffdd46dbb29793335c7a907d",
      "c5e3583afecc49eebabf5957e445e3aa"
     ]
    },
    "id": "u-mC9pqSs7Ef",
    "outputId": "154c1118-e4c1-4950-c2ea-1f17f2c10604"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using downloaded and verified file: .\\cifar10.tgz\n"
     ]
    }
   ],
   "source": [
    "dataset_url = \"https://s3.amazonaws.com/fast-ai-imageclas/cifar10.tgz\"\n",
    "download_url(dataset_url, '.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "QOpZcIuRs-iY"
   },
   "outputs": [],
   "source": [
    "# Extract from archive\n",
    "with tarfile.open('./cifar10.tgz', 'r:gz') as tar:\n",
    "    tar.extractall(path='./data')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eguz40rzs-d4",
    "outputId": "cd335c7d-771c-4986-812f-187d1fa3aa54"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['test', 'train']\n",
      "['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']\n"
     ]
    }
   ],
   "source": [
    "data_dir = './data/cifar10'\n",
    "\n",
    "print(os.listdir(data_dir))\n",
    "classes = os.listdir(data_dir + \"/train\")\n",
    "print(classes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "7XPsTkRBs-Vc"
   },
   "outputs": [],
   "source": [
    "trainset = ImageFolder(data_dir+'/train', transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)\n",
    "\n",
    "testset = ImageFolder(data_dir+'/test', transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "mbFsNabXso7y"
   },
   "outputs": [],
   "source": [
    "# final\n",
    "class Encoder(nn.Module):\n",
    "    '''\n",
    "    simple encoder with a single hidden dense layer (ReLU activation)\n",
    "    and linear projections to the diag-Gauss parameters\n",
    "    '''\n",
    "    def __init__(self, latent_size=256, nc=3):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.latent_size = latent_size\n",
    "        self.nc = nc\n",
    "\n",
    "        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv2 = nn.Conv2d(16, out_channels=16, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv3 = nn.Conv2d(16, out_channels=32, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv4 = nn.Conv2d(32, out_channels=32, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv5 = nn.Conv2d(32, out_channels=64, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv6 = nn.Conv2d(64, out_channels=64, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.enc_linear = nn.Sequential(\n",
    "                nn.Linear(1024, 256),\n",
    "                nn.LeakyReLU(),\n",
    "                nn.Linear(256, 512),\n",
    "                nn.LeakyReLU())\n",
    "        \n",
    "        self.fc_mu = nn.Linear(512, latent_size)\n",
    "        self.fc_var = nn.Linear(512, latent_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # block 1\n",
    "        x = self.conv1(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv2(x)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        #block 2\n",
    "        x = self.conv3(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv4(x)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        # block 3\n",
    "        x = self.conv5(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv6(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = x.view(x.shape[0], -1)\n",
    "        x = self.enc_linear(x)\n",
    "\n",
    "        # Split the result into mu and var components\n",
    "        # of the latent Gaussian distribution\n",
    "        mu = self.fc_mu(x)\n",
    "        log_var = self.fc_var(x)\n",
    "\n",
    "        return mu, log_var\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "vhxjvVVEso5M"
   },
   "outputs": [],
   "source": [
    "# final\n",
    "class Decoder(nn.Module):\n",
    "    '''\n",
    "    simple decoder: single dense hidden layer (ReLU activation) followed by \n",
    "    output layer with a sigmoid to squish values\n",
    "    '''\n",
    "    def __init__(self, latent_size=256, nc=3):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.latent_size = latent_size\n",
    "        self.nc = nc\n",
    "        self.decoder_input = nn.Sequential(\n",
    "            nn.Linear(latent_size, 256),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(256, 512),\n",
    "            nn.LeakyReLU(),\n",
    "        )\n",
    "\n",
    "        self.activation = nn.Sigmoid()\n",
    "\n",
    "        self.conv1U = nn.ConvTranspose2d(32, out_channels=64, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
    "        self.conv2U = nn.ConvTranspose2d(64, out_channels=32, kernel_size=3, stride=1, padding=1, output_padding=0)\n",
    "        self.conv3U = nn.ConvTranspose2d(32, out_channels=32, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
    "        self.conv4U = nn.ConvTranspose2d(32, out_channels=16, kernel_size=3, stride=1, padding=1, output_padding=0)\n",
    "        self.conv5U = nn.ConvTranspose2d(16, out_channels=16, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
    "        self.conv6U = nn.ConvTranspose2d(16, out_channels=3, kernel_size=3, stride=1, padding=1, output_padding=0)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = self.decoder_input(z)\n",
    "        x = x.view(-1, 32, 4, 4)\n",
    "\n",
    "        # block 1\n",
    "        x = self.conv1U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv2U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        # block 2\n",
    "        x = self.conv3U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv4U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        # block 3\n",
    "        x = self.conv5U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv6U(x)\n",
    "\n",
    "        return self.activation(x)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "Dst3HRFJso1m"
   },
   "outputs": [],
   "source": [
    "def sample(mu, logvar):\n",
    "    std = torch.exp(0.5*logvar)  # e^(1/2 * log(std^2))\n",
    "    eps = torch.randn_like(std)  # random ~ N(0, 1)\n",
    "    return eps.mul(std).add_(mu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aa95RL3YuXrh"
   },
   "outputs": [],
   "source": [
    "device = \"cuda:0\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "list_fw1x_enc = ['C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta1.weights',\n",
    "                 'C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta16.weights',\n",
    "                 'C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta64.weights',\n",
    "                 'C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta256.weights'\n",
    "                ]\n",
    "\n",
    "list_fw1x_dec = ['C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta1.weights',\n",
    "                 'C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta16.weights',\n",
    "                 'C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta64.weights',\n",
    "                 'C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta256.weights'\n",
    "                ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "3Cd2B11fsoz0"
   },
   "outputs": [],
   "source": [
    "# build the model and load state\n",
    "enc = Encoder()\n",
    "enc.load_state_dict(torch.load('C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x1-width/encCIFAR10beta1_x1.weights', map_location=torch.device('cpu')))\n",
    "\n",
    "dec = Decoder()\n",
    "dec.load_state_dict(torch.load('C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x1-width/decCIFAR10beta1_x1.weights', map_location=torch.device('cpu')))\n",
    "\n",
    "# put model in eval mode\n",
    "enc = enc.eval() \n",
    "dec = dec.eval() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta1.weights C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta1.weights\n",
      "C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta16.weights C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta16.weights\n",
      "C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta64.weights C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta64.weights\n",
      "C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10beta256.weights C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10beta256.weights\n"
     ]
    }
   ],
   "source": [
    "for i, j in zip(list_fw1x_enc, list_fw1x_dec):\n",
    "    print(i, j)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta = 1\n",
    "loss_for_beta = []\n",
    "for i, j in zip(list_fw1x_enc, list_fw1x_dec):\n",
    "    enc = Encoder()\n",
    "    enc.load_state_dict(torch.load(i, map_location=torch.device('cpu')))\n",
    "\n",
    "    dec = Decoder()\n",
    "    dec.load_state_dict(torch.load(j, map_location=torch.device('cpu')))\n",
    "\n",
    "    # put model in eval mode\n",
    "    enc = enc.eval() \n",
    "    dec = dec.eval() \n",
    "    losses = []\n",
    "    for i, data in enumerate(testloader, 0):\n",
    "        inputs, _ = data\n",
    "    #     inputs, _ = inputs.to(device), _.to(device)\n",
    "\n",
    "        mu, log_sigma2 = enc(inputs)\n",
    "        z = sample(mu, log_sigma2)\n",
    "        outputs = dec(z)\n",
    "\n",
    "        recon = F.binary_cross_entropy(outputs, inputs, reduction='sum') / inputs.shape[0]\n",
    "\n",
    "        kl_diverge = 0.5 * torch.mean(torch.pow(mu, 2) + torch.pow(log_sigma2, 2) - torch.log(torch.pow(log_sigma2, 2)) - 1)\n",
    "\n",
    "        loss = recon + beta*kl_diverge \n",
    "\n",
    "        # keep track of the loss and update the stats\n",
    "        losses.append(loss.item())\n",
    "    loss_for_beta.append(np.mean(losses))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1724.9300258974486, 1730.2648106828522, 1742.8041945831685, 1770.117901379549]\n"
     ]
    }
   ],
   "source": [
    "print(loss_for_beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FiLMBlock(nn.Module):\n",
    "    def __init__(self, out_channels = None):\n",
    "        super(FiLMBlock, self).__init__()\n",
    "        self.out_channels = out_channels\n",
    "        self.mu = nn.Linear(256, out_channels)\n",
    "        self.sigma = nn.Linear(256, out_channels)\n",
    "        self.activation = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x):\n",
    "        mu = self.mu(x)\n",
    "        sigma = self.sigma(x)\n",
    "        return self.activation(mu), self.activation(sigma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FiLMImplement(nn.Module):\n",
    "    def __init__(self, window):\n",
    "        super(FiLMImplement, self).__init__()\n",
    "        self.window = window\n",
    "\n",
    "    def broadcast_2d(self, x):\n",
    "        return torch.broadcast_to(x, (self.window, self.window))\n",
    "\n",
    "    def forward(self, x, mu, sigma):\n",
    "        mu_broadcast = torch.stack(tuple(map(self.broadcast_2d, mu.squeeze(0))), dim=0)\n",
    "        sigma_broadcast = torch.stack(tuple(map(self.broadcast_2d, sigma.squeeze(0))), dim=0)\n",
    "        x = sigma_broadcast * x + mu_broadcast\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# final\n",
    "class Encoder1(nn.Module):\n",
    "    '''\n",
    "    simple encoder with a single hidden dense layer (ReLU activation)\n",
    "    and linear projections to the diag-Gauss parameters\n",
    "    '''\n",
    "    def __init__(self, latent_size=256, nc=3):\n",
    "        super(Encoder1, self).__init__()\n",
    "        self.latent_size = latent_size\n",
    "        self.nc = nc\n",
    "\n",
    "        self.film1 = FiLMBlock(16)\n",
    "        self.film2 = FiLMBlock(32)\n",
    "        self.film3 = FiLMBlock(64)\n",
    "        self.filmimp1 = FiLMImplement(16)\n",
    "        self.filmimp2 = FiLMImplement(8)\n",
    "        self.filmimp3 = FiLMImplement(4)\n",
    "\n",
    "        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv2 = nn.Conv2d(16, out_channels=16, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv3 = nn.Conv2d(16, out_channels=32, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv4 = nn.Conv2d(32, out_channels=32, kernel_size=3, stride=1, padding=1)\n",
    "        self.conv5 = nn.Conv2d(32, out_channels=64, kernel_size=3, stride=2, padding=1)\n",
    "        self.conv6 = nn.Conv2d(64, out_channels=64, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.enc_linear = nn.Sequential(\n",
    "                nn.Linear(1024, 256),\n",
    "                nn.LeakyReLU(),\n",
    "                nn.Linear(256, 512),\n",
    "                nn.LeakyReLU())\n",
    "        \n",
    "        self.fc_mu = nn.Linear(512, latent_size)\n",
    "        self.fc_var = nn.Linear(512, latent_size)\n",
    "\n",
    "    def forward(self, x, helo=None):\n",
    "        # block 1\n",
    "        x = self.conv1(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv2(x)\n",
    "        mu, sigma = self.film1( helo )\n",
    "        x = self.filmimp1(x, mu, sigma)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        #block 2\n",
    "        x = self.conv3(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv4(x)\n",
    "        mu2, sigma2 = self.film2( helo )\n",
    "        x = self.filmimp2(x, mu2, sigma2)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        # block 3\n",
    "        x = self.conv5(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv6(x)\n",
    "        mu3, sigma3 = self.film3( helo )\n",
    "        x = self.filmimp3(x, mu3, sigma3)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = x.view(x.shape[0], -1)\n",
    "        x = self.enc_linear(x)\n",
    "\n",
    "        # Split the result into mu and var components\n",
    "        # of the latent Gaussian distribution\n",
    "        mu = self.fc_mu(x)\n",
    "        log_var = self.fc_var(x)\n",
    "\n",
    "        return mu, log_var\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# final\n",
    "class Decoder1(nn.Module):\n",
    "    '''\n",
    "    simple decoder: single dense hidden layer (ReLU activation) followed by \n",
    "    output layer with a sigmoid to squish values\n",
    "    '''\n",
    "    def __init__(self, latent_size=256, nc=3):\n",
    "        super(Decoder1, self).__init__()\n",
    "        self.latent_size = latent_size\n",
    "        self.nc = nc\n",
    "        self.decoder_input = nn.Sequential(\n",
    "            nn.Linear(latent_size, 256),\n",
    "            nn.LeakyReLU(),\n",
    "            nn.Linear(256, 512),\n",
    "            nn.LeakyReLU(),\n",
    "        )\n",
    "\n",
    "        self.activation = nn.Sigmoid()\n",
    "\n",
    "        self.film1U = FiLMBlock(32)\n",
    "        self.film2U = FiLMBlock(16)\n",
    "        self.film3U = FiLMBlock(3)\n",
    "        self.filmimp1U = FiLMImplement(8)\n",
    "        self.filmimp2U = FiLMImplement(16)\n",
    "        self.filmimp3U = FiLMImplement(32)\n",
    "\n",
    "        self.conv1U = nn.ConvTranspose2d(32, out_channels=64, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
    "        self.conv2U = nn.ConvTranspose2d(64, out_channels=32, kernel_size=3, stride=1, padding=1, output_padding=0)\n",
    "        self.conv3U = nn.ConvTranspose2d(32, out_channels=32, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
    "        self.conv4U = nn.ConvTranspose2d(32, out_channels=16, kernel_size=3, stride=1, padding=1, output_padding=0)\n",
    "        self.conv5U = nn.ConvTranspose2d(16, out_channels=16, kernel_size=3, stride=2, padding=1, output_padding=1)\n",
    "        self.conv6U = nn.ConvTranspose2d(16, out_channels=3, kernel_size=3, stride=1, padding=1, output_padding=0)\n",
    "    \n",
    "    def forward(self, x , helo=None):\n",
    "        x = self.decoder_input(z)\n",
    "        x = x.view(-1, 32, 4, 4)\n",
    "\n",
    "        # block 1\n",
    "        x = self.conv1U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv2U(x)\n",
    "        mu11, sigma11 = self.film1U( helo )\n",
    "        x = self.filmimp1U(x, mu11, sigma11)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        # block 2\n",
    "        x = self.conv3U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv4U(x)\n",
    "        mu22, sigma22 = self.film2U( helo )\n",
    "        x = self.filmimp2U(x, mu22, sigma22)\n",
    "        x = F.leaky_relu(x)\n",
    "\n",
    "        # block 3\n",
    "        x = self.conv5U(x)\n",
    "        x = F.leaky_relu(x)\n",
    "        x = self.conv6U(x)\n",
    "        mu33, sigma33 = self.film3U( helo )\n",
    "        x = self.filmimp3U(x, mu33, sigma33)\n",
    "        return self.activation(x)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc1 = Encoder1()\n",
    "enc1.load_state_dict(torch.load('C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/encCIFAR10yoto.weights', map_location=torch.device('cpu')))\n",
    "\n",
    "dec1 = Decoder1()\n",
    "dec1.load_state_dict(torch.load('C:/Users/Hanu/Downloads/DL_cw/CIFAR10-x2-width/decCIFAR10yoto.weights', map_location=torch.device('cpu')))\n",
    "\n",
    "# put model in eval mode\n",
    "enc1 = enc1.eval() \n",
    "dec1 = dec1.eval() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "beta_vals = [1,16,64,256]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "16\n",
      "64\n",
      "256\n"
     ]
    }
   ],
   "source": [
    "loss_for_beta_yoto = []\n",
    "for beta in beta_vals:\n",
    "    print(beta)\n",
    "    losses = []\n",
    "    for i, data in enumerate(testloader, 0):\n",
    "        inputs, _ = data\n",
    "    #     inputs, _ = inputs.to(device), _.to(device)\n",
    "        \n",
    "        beta2 = torch.tensor([float(beta)], requires_grad=False)\n",
    "        beta2 = torch.broadcast_to(beta2, (1, 256))\n",
    "        \n",
    "        mu, log_sigma2 = enc1(inputs, beta2)\n",
    "        z = sample(mu, log_sigma2)\n",
    "        outputs = dec1(z, beta2)\n",
    "\n",
    "        recon = F.binary_cross_entropy(outputs, inputs, reduction='sum') / inputs.shape[0]\n",
    "\n",
    "        kl_diverge = 0.5 * torch.mean(torch.pow(mu, 2) + torch.pow(log_sigma2, 2) - torch.log(torch.pow(log_sigma2, 2)) - 1)\n",
    "\n",
    "        loss = recon + beta*kl_diverge \n",
    "\n",
    "        # keep track of the loss and update the stats\n",
    "        losses.append(loss.item())\n",
    "    loss_for_beta_yoto.append(np.mean(losses))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1773.4991640501385, 1784.2135318804392, 1802.9030421776108, 1833.0603320930577]\n"
     ]
    }
   ],
   "source": [
    "print(loss_for_beta_yoto)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "Testing_normal.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "21ca9bd0e04245a0ba6c283caea226ef": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_896980c70dbd4b0aafdc9b8ed11ca093",
       "IPY_MODEL_c0321ef49e4848efb7a6696e54b18914",
       "IPY_MODEL_4d0e339661be4c47b543e5ec327c13f6"
      ],
      "layout": "IPY_MODEL_7ad1a00b889f4dd6babb75f7357ddb43"
     }
    },
    "35062c0b25e04cd6b7d8f7d1dc6e64ec": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "37d32865a623432cb46c24ac1a044386": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "4d0e339661be4c47b543e5ec327c13f6": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d8337d8fffdd46dbb29793335c7a907d",
      "placeholder": "​",
      "style": "IPY_MODEL_c5e3583afecc49eebabf5957e445e3aa",
      "value": " 135108608/? [00:02&lt;00:00, 49467579.31it/s]"
     }
    },
    "7ad1a00b889f4dd6babb75f7357ddb43": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "896980c70dbd4b0aafdc9b8ed11ca093": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_35062c0b25e04cd6b7d8f7d1dc6e64ec",
      "placeholder": "​",
      "style": "IPY_MODEL_ac6012f09e6e4be0bf4d9c8d8ef5edfb",
      "value": ""
     }
    },
    "ac6012f09e6e4be0bf4d9c8d8ef5edfb": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "c0321ef49e4848efb7a6696e54b18914": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_ddcf1143b7314adcb64d6f64cfbbe462",
      "max": 135107811,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_37d32865a623432cb46c24ac1a044386",
      "value": 135107811
     }
    },
    "c5e3583afecc49eebabf5957e445e3aa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "d8337d8fffdd46dbb29793335c7a907d": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "ddcf1143b7314adcb64d6f64cfbbe462": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}