{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Variational AutoEncoder\n\n![](https://cdn-images-1.medium.com/max/2600/1*22cSCfmktNIwH5m__u2ffA.png)"},{"metadata":{},"cell_type":"markdown","source":"**VAE** can do generation of data that similar to it have seen before. But before diving into VAE, let's understand how simple **AutoEncoder** works. <br><br>\nAutoEncoder consist of two parts: **encoder** and **decoder** networks.<br>\n![](https://blog.keras.io/img/ae/autoencoder_schema.jpg)\n**Encoder** takes your data as input and produces some continuous representation (aka **latent variable**) of given samples. This representation should have smaller dimension than data. The reason is - AE must take only most important information about data and throw away non-important just like PCA, but in this case it works non-linearly. <br><br>\n**Decoder** takes this representation as input and tries to reconstruct the original data. <br><br>\nLoss of this network is some defined distance between original input and reconstructed output. In Images it is usually pixel-wise **Binary CrossEntropy Loss**."},{"metadata":{},"cell_type":"markdown","source":"## Variational AE\n\nThere is small change in latent representation.\nNow encoder will produce two vectors: vector of means and vector of standart deviations.\n![](https://cdn-images-1.medium.com/max/1600/1*CiVcrrPmpcB1YGMkTF7hzA.png)\nUsing this mean and std vectors we can generate some amount of new samples and propagate them through decoder."},{"metadata":{},"cell_type":"markdown","source":"## Kullback–Leibler divergence\nFor this model we need to define new loss. KL divergence measures how two distributions are different. Exact mathematical formula: \n![](https://cdn-images-1.medium.com/max/1600/0*opyFpDwDt0H8rfCv) <br>\nP and Q are **Probability Density Functions** (PDF) of distributions. We know that log(1)=0, so when Q and P are equal, KL distance is 0.<br><br>\nFor VAEs it is distance between hidden variables and standart normal distribution. This forces latent space to be distributed normally.<br> Derived formula for VAEs is:\n![](https://cdn-images-1.medium.com/max/1200/1*uEAxCmyVKxzZOJG6afkCCg.png)\nμ - mean and σ - std <br><br>\nIf we will use only KL Loss then encodings in hidden space will be distributed randomly and near the center of space. Decoder will not be able to reconstruct something from this noise. Out goal to differentiate different classes to different clusters. <br><br>\n**Reconstruction loss** will help us to separate classes. It is a distance between original data and generated. Typically, in Images reconstruction loss is pixel-wise BCE or MSE. <br><br>\n**Final loss** will be the combination of these two. <br>\n`Loss = KL + ReconstructLoss`"},{"metadata":{},"cell_type":"markdown","source":"## Reparametrization trick\nProbably, you wondered: How gradients from reconstruction loss goes to encoder through sampling? Sampling is not differentiable.<br> Here is some method called `reparametrization`.<br>\nThe key insight is that `N(μ, σ) == N(0, 1) * σ + μ`<br>\nEncoder predicts means and sigmas and combines them with standart normal noise, so that gradients now is available for encoder"},{"metadata":{},"cell_type":"markdown","source":"## Why we should use VAE, not AE?\nVAE is more complicated and requires knowledge of math, what's the sacral meaning of this? <br><br>\nAs I said, VAE makes strong assumption about distribution of latent variable. It is better for us because vectors of samples from the same class will be lie continuously. We can do better clusterization, interpolation. Generating of new data will be easier. If hidden space would have gaps between clusters then sampling from this space would produce bad results."},{"metadata":{},"cell_type":"markdown","source":"# Time to plunge into the code"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nimport matplotlib.pyplot as plt\n\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.autograd import Variable\n\nfrom PIL import Image\n\nfrom tqdm import tqdm_notebook as tqdm","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Custom dataset"},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"class DogDataset(Dataset):\n    def __init__(self, img_dir, transform1=None, transform2=None):\n    \n        self.img_dir = img_dir\n        self.img_names = os.listdir(img_dir)\n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.imgs = []\n        for img_name in self.img_names:\n            img = Image.open(os.path.join(img_dir, img_name))\n            \n            if self.transform1 is not None:\n                img = self.transform1(img)\n                \n            self.imgs.append(img)\n\n    def __getitem__(self, index):\n        img = self.imgs[index]\n        \n        if self.transform2 is not None:\n            img = self.transform2(img)\n        \n        return img\n\n    def __len__(self):\n        return len(self.imgs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# First preprocessing of data\ntransform1 = transforms.Compose([transforms.Resize(64),\n                                transforms.CenterCrop(64)])\n\n# Data augmentation and converting to tensors\nrandom_transforms = [transforms.RandomRotation(degrees=10)]\ntransform2 = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5),\n                                 transforms.RandomApply(random_transforms, p=0.3), \n                                 transforms.ToTensor(),\n                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n                                 \ntrain_dataset = DogDataset(img_dir='../input/all-dogs/all-dogs/',\n                           transform1=transform1,\n                           transform2=transform2)\n\ntrain_loader = DataLoader(dataset=train_dataset,\n                          batch_size=batch_size,\n                          shuffle=True,\n                          num_workers=4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Examples of data"},{"metadata":{"trusted":true},"cell_type":"code","source":"x = next(iter(train_loader))\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(x):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    \n    img = img.numpy().transpose(1, 2, 0)\n    plt.imshow((img+1.)/2.)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## VAE Model\nThe code below is based on https://github.com/atinghosh/VAE-pytorch"},{"metadata":{"trusted":true},"cell_type":"code","source":"class VAE(nn.Module):\n    def __init__(self, latent_dim=128, no_of_sample=10, batch_size=32, channels=3):\n        super(VAE, self).__init__()\n        \n        self.no_of_sample = no_of_sample\n        self.batch_size = batch_size\n        self.channels = channels\n        self.latent_dim = latent_dim\n        \n        \n        # Encoder\n        def convlayer_enc(n_input, n_output, k_size=4, stride=2, padding=1, bn=False):\n            block = [nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False)]\n            if bn:\n                block.append(nn.BatchNorm2d(n_output))\n            block.append(nn.LeakyReLU(0.2, inplace=True))\n            return block\n        \n        self.encoder = nn.Sequential(\n            *convlayer_enc(self.channels, 64, 4, 2, 2),               # (64, 32, 32)\n            *convlayer_enc(64, 128, 4, 2, 2),                         # (128, 16, 16)\n            *convlayer_enc(128, 256, 4, 2, 2, bn=True),               # (256, 8, 8)\n            *convlayer_enc(256, 512, 4, 2, 2, bn=True),               # (512, 4, 4)\n            nn.Conv2d(512, self.latent_dim*2, 4, 1, 1, bias=False),   # (latent_dim*2, 4, 4)\n            nn.LeakyReLU(0.2, inplace=True)\n        )\n        \n        \n        # Decoder\n        def convlayer_dec(n_input, n_output, k_size=4, stride=2, padding=0):\n            block = [\n                nn.ConvTranspose2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=False),\n                nn.BatchNorm2d(n_output),\n                nn.ReLU(inplace=True),\n            ]\n            return block\n        \n        self.decoder = nn.Sequential(\n            *convlayer_dec(self.latent_dim, 512, 4, 2, 1),           # (512, 8, 8)\n            *convlayer_dec(512, 256, 4, 2, 1),                       # (256, 16, 16)\n            *convlayer_dec(256, 128, 4, 2, 1),                       # (128, 32, 32)\n            *convlayer_dec(128, 64, 4, 2, 1),                        # (64, 64, 64)\n            nn.ConvTranspose2d(64, self.channels, 3, 1, 1),          # (3, 64, 64)\n            nn.Sigmoid()\n        )\n\n    def encode(self, x):\n        '''return mu_z and logvar_z'''\n        x = self.encoder(x)\n        return x[:, :self.latent_dim, :, :], x[:, self.latent_dim:, :, :]\n    \n    def decode(self, z):\n        z = self.decoder(z)\n        return z.view(-1, 3 * 64 * 64)\n\n    def reparameterize(self, mu, logvar):\n        if self.training:\n            # multiply log variance with 0.5, then in-place exponent\n            # yielding the standard deviation\n\n            sample_z = []\n            for _ in range(self.no_of_sample):\n                std = logvar.mul(0.5).exp_()\n                eps = Variable(std.data.new(std.size()).normal_())\n                sample_z.append(eps.mul(std).add_(mu))\n            return sample_z\n        \n        else:\n            return mu\n\n    def forward(self, x):\n        mu, logvar = self.encode(x)\n        z = self.reparameterize(mu, logvar)\n        \n        if self.training:\n            return [self.decode(z) for z in z], mu, logvar\n        else:\n            return self.decode(z), mu, logvar\n\n    def loss_function(self, recon_x, x, mu, logvar):\n\n        if self.training:\n            BCE = 0\n            for recon_x_one in recon_x:\n                BCE += F.binary_cross_entropy(recon_x_one, x.view(-1, 3 * 64 * 64))\n            BCE /= len(recon_x)\n        else:\n            BCE = F.binary_cross_entropy(recon_x, x.view(-1, 3 * 64 * 64))\n\n        KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n        KLD /= self.batch_size * 3 * 64 * 64\n\n        return BCE + KLD","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Input dim is `64 x 64 x 3 = 12288`<br>\nLatent dim is `4 x 4 x 32 = 512`<br>\nBottleneck is 24 times smaller than input image! Autoencoder should keep most important information"},{"metadata":{"trusted":true},"cell_type":"code","source":"lr = 0.001\nepochs = 50\nlatent_dim = 32\n\nmodel = VAE(latent_dim, batch_size=batch_size).to(device)\noptimizer = optim.Adam(model.parameters(), lr=lr)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Image for validation"},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.imshow((x[0].numpy().transpose(1, 2, 0)+1)/2)\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train loop"},{"metadata":{"trusted":true},"cell_type":"code","source":"for epoch in range(1, epochs+1):\n    model.train()\n    print(f'Epoch {epoch} start')\n    \n    for batch_idx, data in enumerate(train_loader):\n        data = data.to(device)\n        optimizer.zero_grad()\n\n        recon_batch, mu, logvar = model(data)\n        loss = model.loss_function(recon_batch, data, mu, logvar)\n\n        loss.backward()\n        optimizer.step()\n        \n    model.eval()\n    recon_img, _, _ = model(x[:1].to(device))\n    img = recon_img.view(3, 64, 64).detach().cpu().numpy().transpose(1, 2, 0)\n    \n    plt.imshow((img+1.)/2.)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Check how well VAE reconstruct images"},{"metadata":{"trusted":true},"cell_type":"code","source":"reconstructed, mu, _ = model(x.to(device))\nreconstructed = reconstructed.view(-1, 3, 64, 64).detach().cpu().numpy().transpose(0, 2, 3, 1)\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(reconstructed):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    plt.imshow((img+1.)/2.)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Walk in latent space from one dog to another"},{"metadata":{"trusted":true},"cell_type":"code","source":"first_dog_idx = 0\nsecond_dog_idx = 1\n\ndz = (mu[second_dog_idx] - mu[first_dog_idx]) / 31\nwalk = Variable(torch.randn(32, latent_dim, 4, 4)).to(device)\nwalk[0] = mu[first_dog_idx]\n\nfor i in range(1, 32):\n    walk[i] = walk[i-1] + dz\nwalk = model.decoder(walk).detach().cpu().numpy().transpose(0, 2, 3, 1)\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(walk):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    plt.imshow((img+1.)/2.)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Generate random noise and run decoder on this"},{"metadata":{"trusted":true},"cell_type":"code","source":"samples = Variable(torch.randn(32, latent_dim, 4, 4)).to(device)\nsamples = model.decoder(samples).detach().cpu().numpy().transpose(0, 2, 3, 1)\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(samples):\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    plt.imshow((img+1.)/2.)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I was inspired by:\n - https://towardsdatascience.com/intuitively-understanding-variational-autoencoders-1bfe67eb5daf\n - https://github.com/atinghosh/VAE-pytorch"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}