{"cells":[{"metadata":{},"cell_type":"markdown","source":"The script is base on the author of the VQ-VAE.\nI only change the dataset."},{"metadata":{},"cell_type":"markdown","source":"# [VQ-VAE](https://arxiv.org/abs/1711.00937) by  [Aäron van den Oord](https://twitter.com/avdnoord) et al. in PyTorch\n\n## Introduction\n\nVariational Auto Encoders (VAEs) can be thought of as what all but the last layer of a neural network is doing, namely feature extraction or seperating out the data. Thus given some data we can think of using a neural network for representation generation. \n\nRecall that the goal of a generative model is to estimate the probability distribution of high dimensional data such as images, videos, audio or even text by learning the underlying structure in the data as well as the dependencies between the different elements of the data. This is very useful since we can then use this representation to generate new data with similar properties. This way we can also learn useful features from the data in an unsupervised fashion.\n\nThe VQ-VAE uses a discrete latent representation mostly because many important real-world objects are discrete. For example in images we might have categories like \"Cat\", \"Car\", etc. and it might not make sense to interpolate between these categories. Discrete representations are also easier to model since each category has a single value whereas if we had a continous latent space then we will need to normalize this density function and learn the dependencies between the different variables which could be very complex.\n\n### Code\n\nI have followed the code from the TensorFlow implementation by the author which you can find here [vqvae.py](https://github.com/deepmind/sonnet/blob/master/sonnet/python/modules/nets/vqvae.py) and [vqvae_example.ipynb](https://github.com/deepmind/sonnet/blob/master/sonnet/examples/vqvae_example.ipynb). \n\nAnother PyTorch implementation is found at [pytorch-vqvae](https://github.com/ritheshkumar95/pytorch-vqvae).\n\n\n## Basic Idea\n\nThe overall architecture is summarized in the diagram below:\n\n![](images/vq-vae.png)"},{"metadata":{},"cell_type":"markdown","source":"We start by defining a latent embedding space of dimension `[K, D]` where `K` are the number of embeddings and `D` is the dimensionality of each latent embeddng vector, i.e. $e_i \\in \\mathbb{R}^{D}$. The model is comprised of an encoder and a decoder. The encoder will map the input to a sequence of discrete latent variables, whereas the decoder will try to reconstruct the input from these latent sequences. \n\nMore preciesly, the model will take in batches of RGB images,  say $x$, each of size 32x32 for our example, and pass it through a ConvNet encoder producing some output $E(x)$, where we make sure the channels are the same as the dimensionality of the latent embedding vectors. To calculate the discrete latent variable we find the nearest embedding vector and output it's index. \n\nThe input to the decoder is the embedding vector corresponding to the index which is passed through the decoder to produce the reconstructed image. \n\nSince the nearest neighbour lookup has no real gradient in the backward pass we simply pass the gradients from the decoder to the encoder  unaltered. The intuition is that since the output representation of the encoder and the input to the decoder share the same `D` channel dimensional space, the gradients contain useful information for how the encoder has to change its output to lower the reconstruction loss.\n\n## Loss\n\nThe total loss is actually composed of three components\n\n1. **reconstruction loss**: which optimizes the decoder and encoder\n1. **codebook loss**: due to the fact that gradients bypass the embedding, we use a dictionary learning algorithm  which uses an $l_2$  error to move the embedding vectors $e_i$ towards the encoder output\n1. **commitment loss**:  since the volume of the embedding space is dimensionless, it can grow arbirtarily if the embeddings $e_i$ do not train as fast as  the encoder parameters, and thus we add a commitment loss to make sure that the encoder commits to an embedding"},{"metadata":{},"cell_type":"markdown","source":"# VQ-VAE Introduction\nThe script is based on the https://github.com/ritheshkumar95/pytorch-vqvae. In this script, I will told a story about VQ-VAE(Vector-Quantized VAEs).\n\nVQ-VAE is similar to VAE, but it has a VQ between the encoder and decoder.\n\nThe components of VQ-VAE are:\n- encoder\n- VQ\n- decoder\n\nFor encoder, \n- the input is $(b, c, w, h)$, where b is the batch_size, c is channel, w is width, h is height. \n- the output is $(b, c_e, w_e, h_e)$, where b is the batch_size, $c_e$ is the output channel of the encoder, $w_e$ is the width, $h_e$ is the height.\n\nFor VQ, we set a code book $e \\in R ^{K\\times D}$, $K$ is the number of latent vector, and $D$ is the dimension of the latent vector, and $D = c_e$\n\n![VQ-VAE](https://ss.csdn.net/p?https://image.jiqizhixin.com/uploads/wangeditor/6a9585fb-0da3-404a-a86c-8b5f9b267a72/4739720171107141110.png)\n\n$$q(z = k | x) =\n\\begin{cases}\n1 & if \\ k = arg min_j ||z_e(x) - e_j||_2 \\\\\n0 & otherwise\n\\end{cases}$$\n\nFor decoder,\n- the input is the VQ which is affected by the output of encoder\n- the output is $(b, c, w, h)$, which is the same size with the input of encoder.\n\n\n**VQ-VAE Loss**\n$$\nL = \\log p(x | z_q(x)) + || sg[z_e(x)] - e ||_2^2 + \\beta|| z_e(x) - sg[e] ||_2^2\n$$\n\nThe symbol $sg$ means that the gradient of the variable will not be calculated and backward."},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import print_function\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom scipy.signal import savgol_filter\n\nfrom six.moves import xrange\n\nimport umap\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nimport torch.optim as optim\n\nimport torchvision.datasets as datasets\nimport torchvision.transforms as transforms\nfrom torchvision.utils import make_grid","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\ndevice","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Load Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"ComputeLB = True\nDogsOnly = True\n\nimport numpy as np, pandas as pd, os\nimport xml.etree.ElementTree as ET \nimport matplotlib.pyplot as plt, zipfile \nfrom PIL import Image \nfrom tqdm import tqdm_notebook\n\nROOT = '../input/generative-dog-images/'\nif not ComputeLB: ROOT = '../input/'\nIMAGES = os.listdir(ROOT + 'all-dogs/all-dogs/')\nbreeds = os.listdir(ROOT + 'annotation/Annotation/') \n\nidxIn = 0; namesIn = []\nimagesIn = np.zeros((25000,64,64,3))\n\n# CROP WITH BOUNDING BOXES TO GET DOGS ONLY\n# https://www.kaggle.com/paulorzp/show-annotations-and-breeds\nif DogsOnly:\n    for breed in tqdm_notebook(breeds):\n        for dog in os.listdir(ROOT+'annotation/Annotation/'+breed):\n            try: img = Image.open(ROOT+'all-dogs/all-dogs/'+dog+'.jpg') \n            except: continue           \n            tree = ET.parse(ROOT+'annotation/Annotation/'+breed+'/'+dog)\n            root = tree.getroot()\n            objects = root.findall('object')\n            for o in objects:\n                bndbox = o.find('bndbox') \n                xmin = int(bndbox.find('xmin').text)\n                ymin = int(bndbox.find('ymin').text)\n                xmax = int(bndbox.find('xmax').text)\n                ymax = int(bndbox.find('ymax').text)\n                w_, h_ = img.size\n                w = np.max((xmax - xmin, ymax - ymin))\n                img2 = img.crop((xmin, ymin, min(xmin+w, w_), min(ymin+w, h_)))\n                img2 = img2.resize((64,64), Image.ANTIALIAS)\n                imagesIn[idxIn,:,:,:] = np.asarray(img2)\n                #if idxIn%1000==0: print(idxIn)\n                namesIn.append(breed)\n                idxIn += 1\n    idx = np.arange(idxIn)\n    np.random.shuffle(idx)\n    imagesIn = imagesIn[idx,:,:,:]\n    namesIn = np.array(namesIn)[idx]\n    \n# RANDOMLY CROP FULL IMAGES\nelse:\n    x = np.random.choice(np.arange(20579),10000)\n    for k in tqdm_notebook(range(len(x))):\n        img = Image.open(ROOT + 'all-dogs/all-dogs/' + IMAGES[x[k]])\n        w = img.size[0]\n        h = img.size[1]\n        sz = np.min((w,h))\n        a=0; b=0\n        if w<h: b = (h-sz)//2\n        else: a = (w-sz)//2\n        img = img.crop((0+a, 0+b, sz+a, sz+b))  \n        img = img.resize((64,64), Image.ANTIALIAS)   \n        imagesIn[idxIn,:,:,:] = np.asarray(img)\n        namesIn.append(IMAGES[x[k]])\n        if idxIn%1000==0: print(idxIn)\n        idxIn += 1\n    \n# DISPLAY CROPPED IMAGES\nx = np.random.randint(0,idxIn,25)\nfor k in range(5):\n    plt.figure(figsize=(15,3))\n    for j in range(5):\n        plt.subplot(1,5,j+1)\n        img = Image.fromarray( imagesIn[x[k*5+j],:,:,:].astype('uint8') )\n        plt.axis('off')\n        if not DogsOnly: plt.title(namesIn[x[k*5+j]],fontsize=11)\n        else: plt.title(namesIn[x[k*5+j]].split('-')[1],fontsize=11)\n        plt.imshow(img)\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import TensorDataset, DataLoader\nfrom collections import defaultdict\nfrom sklearn.model_selection import train_test_split\nprint(f'The shape of image is {imagesIn.shape}, the shape of imagename is {namesIn.shape}')\nimagesIntorch = np.array([np.array(image).transpose(2, 1, 0) for image in imagesIn])\nprint(f'The shape of reshaped image is {imagesIntorch.shape}')\ndogs = list(set(namesIn))\nlen_dogs = len(dogs)\nprint(f'the number of dogs is {len_dogs}')\ndog2id = {dogs[i]:i for i in range(len(dogs))}\nid2dog = {v : k for k, v in dog2id.items()}\n# print(dog2id, id2dog)\nidIn = [dog2id[name] for name in namesIn]\n\ntrain_X, validation_X, train_y, validation_y = train_test_split(imagesIntorch, idIn, test_size=0.2, random_state=620402)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"np.array(train_X).shape, np.array(validation_X).shape, np.array(train_y).shape, np.array(validation_y).shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# train_X = train_X/255.\n# validation_X = validation_X/255.","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\ntraining_data = TensorDataset(torch.Tensor(train_X), torch.Tensor(train_y))\nvalidation_data = TensorDataset(torch.Tensor(validation_X), torch.Tensor(validation_y))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_variance = np.var(train_X)\ndata_variance","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Vector Quantizer Layer\n\nThis layer takes a tensor to be quantized. The channel dimension will be used as the space in which to quantize. All other dimensions will be flattened and will be seen as different examples to quantize.\n\nThe output tensor will have the same shape as the input.\n\nAs an example for a `BCHW` tensor of shape `[16, 64, 32, 32]`, we will first convert it to an `BHWC` tensor of shape `[16, 32, 32, 64]` and then reshape it into `[16384, 64]` and all `16384` vectors of size `64`  will be quantized independently. In otherwords, the channels are used as the space in which to quantize. All other dimensions will be flattened and be seen as different examples to quantize, `16384` in this case."},{"metadata":{"trusted":true},"cell_type":"code","source":"class VectorQuantizer(nn.Module):\n    def __init__(self, num_embeddings, embedding_dim, commitment_cost):\n        super(VectorQuantizer, self).__init__()\n        \n        self._embedding_dim = embedding_dim\n        self._num_embeddings = num_embeddings\n        \n        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)\n        self._embedding.weight.data.uniform_(-1/self._num_embeddings, 1/self._num_embeddings)\n        self._commitment_cost = commitment_cost\n\n    def forward(self, inputs):\n        # convert inputs from BCHW -> BHWC\n        inputs = inputs.permute(0, 2, 3, 1).contiguous()\n        input_shape = inputs.shape\n        \n        # Flatten input\n        flat_input = inputs.view(-1, self._embedding_dim)\n        \n        # Calculate distances\n        distances = (torch.sum(flat_input**2, dim=1, keepdim=True) \n                    + torch.sum(self._embedding.weight**2, dim=1)\n                    - 2 * torch.matmul(flat_input, self._embedding.weight.t()))\n            \n        # Encoding\n        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)\n        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings).to(device)\n        encodings.scatter_(1, encoding_indices, 1)\n        \n        # Quantize and unflatten\n        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)\n        \n        # Loss\n        e_latent_loss = torch.mean((quantized.detach() - inputs)**2)\n        q_latent_loss = torch.mean((quantized - inputs.detach())**2)\n        loss = q_latent_loss + self._commitment_cost * e_latent_loss\n        \n        quantized = inputs + (quantized - inputs).detach()\n        avg_probs = torch.mean(encodings, dim=0)\n        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))\n        \n        # convert quantized from BHWC -> BCHW\n        return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test input and output\nnum_embeddings, embedding_dim, commitment_cost = 64, 512, 0.25\ntestVectorQuantizer = VectorQuantizer(num_embeddings, embedding_dim, commitment_cost)\ntestVectorQuantizer.cuda()\ninput_tensor = torch.Tensor(np.random.normal(size = [32, 64, 4, 4]))\nprint(input_tensor.shape)\n_, output_tensor, perplexity, encodings = testVectorQuantizer(input_tensor.cuda())\nprint(output_tensor.shape)\nprint(encodings.shape)\nprint(perplexity.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"encodings, perplexity","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"We will also implement a slightly modified version  which will use exponential moving averages to update the embedding vectors instead of an auxillary loss. This has the advantage that the embedding updates are independent of the choice of optimizer for the encoder, decoder and other parts of the architecture. For most experiments the EMA version trains faster than the non-EMA version."},{"metadata":{"trusted":true},"cell_type":"code","source":"class VectorQuantizerEMA(nn.Module):\n    def __init__(self, num_embeddings, embedding_dim, commitment_cost, decay, epsilon=1e-5):\n        super(VectorQuantizerEMA, self).__init__()\n        \n        self._embedding_dim = embedding_dim\n        self._num_embeddings = num_embeddings\n        \n        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)\n        self._embedding.weight.data.normal_()\n        self._commitment_cost = commitment_cost\n        \n        self.register_buffer('_ema_cluster_size', torch.zeros(num_embeddings))\n        self._ema_w = nn.Parameter(torch.Tensor(num_embeddings, self._embedding_dim))\n        self._ema_w.data.normal_()\n        \n        self._decay = decay\n        self._epsilon = epsilon\n\n    def forward(self, inputs):\n        # convert inputs from BCHW -> BHWC\n        inputs = inputs.permute(0, 2, 3, 1).contiguous()\n        input_shape = inputs.shape\n        \n        # Flatten input\n        flat_input = inputs.view(-1, self._embedding_dim)\n        \n        # Calculate distances\n        distances = (torch.sum(flat_input**2, dim=1, keepdim=True) \n                    + torch.sum(self._embedding.weight**2, dim=1)\n                    - 2 * torch.matmul(flat_input, self._embedding.weight.t()))\n            \n        # Encoding\n        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)\n        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings).to(device)\n        encodings.scatter_(1, encoding_indices, 1)\n        \n        # Use EMA to update the embedding vectors\n        if self.training:\n            self._ema_cluster_size = self._ema_cluster_size * self._decay + \\\n                                     (1 - self._decay) * torch.sum(encodings, 0)\n            \n            # Laplace smoothing of the cluster size\n            n = torch.sum(self._ema_cluster_size.data)\n            self._ema_cluster_size = (\n                (self._ema_cluster_size + self._epsilon)\n                / (n + self._num_embeddings * self._epsilon) * n)\n            \n            dw = torch.matmul(encodings.t(), flat_input)\n            self._ema_w = nn.Parameter(self._ema_w * self._decay + (1 - self._decay) * dw)\n            \n            self._embedding.weight = nn.Parameter(self._ema_w / self._ema_cluster_size.unsqueeze(1))\n        \n        # Quantize and unflatten\n        quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)\n        \n        # Loss\n        e_latent_loss = torch.mean((quantized.detach() - inputs)**2)\n        loss = self._commitment_cost * e_latent_loss\n        \n        quantized = inputs + (quantized - inputs).detach()\n        avg_probs = torch.mean(encodings, dim=0)\n        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))\n        \n        # convert quantized from BHWC -> BCHW\n        return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test input and output\nnum_embeddings, embedding_dim, commitment_cost = 64, 2048, 0.25\ntestVectorQuantizer = VectorQuantizerEMA(num_embeddings, embedding_dim, commitment_cost, 0.25)\ntestVectorQuantizer.cuda()\ninput_tensor = torch.Tensor(np.random.normal(size = [32, 3, 8, 8]))\nprint(input_tensor.shape)\n_, output_tensor, perplexity, encodings = testVectorQuantizer(input_tensor.cuda())\nprint(output_tensor.shape)\nprint(encodings.shape)\nprint(perplexity.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"encodings, perplexity","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Encoder & Decoder Architecture\n\nThe encoder and decoder architecture is based on a ResNet and is implemented below:"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Residual(nn.Module):\n    def __init__(self, in_channels, num_hiddens, num_residual_hiddens):\n        super(Residual, self).__init__()\n        self._block = nn.Sequential(\n            nn.ReLU(True),\n            nn.Conv2d(in_channels=in_channels,\n                      out_channels=num_residual_hiddens,\n                      kernel_size=3, stride=1, padding=1, bias=False),\n            nn.ReLU(True),\n            nn.Conv2d(in_channels=num_residual_hiddens,\n                      out_channels=num_hiddens,\n                      kernel_size=1, stride=1, bias=False)\n        )\n    \n    def forward(self, x):\n        return x + self._block(x)\n\n\nclass ResidualStack(nn.Module):\n    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):\n        super(ResidualStack, self).__init__()\n        self._num_residual_layers = num_residual_layers\n        self._layers = nn.ModuleList([Residual(in_channels, num_hiddens, num_residual_hiddens)\n                             for _ in range(self._num_residual_layers)])\n\n    def forward(self, x):\n        for i in range(self._num_residual_layers):\n            x = self._layers[i](x)\n        return F.relu(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test input and output\nnum_hiddens = 64\nnum_residual_hiddens = 32\nnum_residual_layers = 2\n\ntestResidual = Residual(64, num_hiddens, num_residual_hiddens)\ntestResidual.cuda()\ninput_tensor = torch.Tensor(np.random.normal(size = [128, 64, 64, 64]))\nprint(input_tensor.shape)\noutput_tensor = testResidual(input_tensor.cuda())\nprint(output_tensor.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Encoder(nn.Module):\n    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):\n        super(Encoder, self).__init__()\n\n        self._conv_1 = nn.Conv2d(in_channels=in_channels,\n                                 out_channels=num_hiddens//4,\n                                 kernel_size=4,\n                                 stride=2, padding=1)\n        self._conv_2 = nn.Conv2d(in_channels=num_hiddens//4,\n                                 out_channels=num_hiddens//2,\n                                 kernel_size=4,\n                                 stride=2, padding=1)\n        self._conv_2_1 = nn.Conv2d(in_channels=num_hiddens//2,\n                                   out_channels=num_hiddens,\n                                   kernel_size=4,\n                                   stride=2, padding=1)\n        self._conv_3 = nn.Conv2d(in_channels=num_hiddens,\n                                 out_channels=num_hiddens,\n                                 kernel_size=3,\n                                 stride=1, padding=1)\n        self._residual_stack = ResidualStack(in_channels=num_hiddens,\n                                             num_hiddens=num_hiddens,\n                                             num_residual_layers=num_residual_layers,\n                                             num_residual_hiddens=num_residual_hiddens)\n\n    def forward(self, inputs):\n        x = self._conv_1(inputs)\n        x = F.relu(x)\n        \n        x = self._conv_2(x)\n        x = F.relu(x)\n        \n        x = self._conv_2_1(x)\n        x = F.relu(x)\n        \n        x = self._conv_3(x)\n        return self._residual_stack(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test input and output\nnum_hiddens = 64\nnum_residual_hiddens = 32\nnum_residual_layers = 2\n\ntestEncoder = Encoder(3, num_hiddens, 2, num_residual_hiddens)\ntestEncoder.cuda()\ninput_tensor = torch.Tensor(np.random.normal(size = [32, 3, 64, 64]))\nprint(input_tensor.shape)\noutput_tensor = testEncoder(input_tensor.cuda())\nprint(output_tensor.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Decoder(nn.Module):\n    def __init__(self, in_channels, num_hiddens, num_residual_layers, num_residual_hiddens):\n        super(Decoder, self).__init__()\n        \n        self._conv_1 = nn.Conv2d(in_channels=in_channels,\n                                 out_channels=num_hiddens,\n                                 kernel_size=3, \n                                 stride=1, padding=1)\n        \n        self._residual_stack = ResidualStack(in_channels=num_hiddens,\n                                             num_hiddens=num_hiddens,\n                                             num_residual_layers=num_residual_layers,\n                                             num_residual_hiddens=num_residual_hiddens)\n        \n        self._conv_trans_1 = nn.ConvTranspose2d(in_channels=num_hiddens, \n                                                out_channels=num_hiddens//2,\n                                                kernel_size=4, \n                                                stride=2, padding=1)\n        \n        self._conv_trans_2 = nn.ConvTranspose2d(in_channels=num_hiddens//2, \n                                                out_channels=num_hiddens//4,\n                                                kernel_size=4, \n                                                stride=2, padding=1)\n        \n        self._conv_trans_2_1 = nn.ConvTranspose2d(in_channels=num_hiddens//4, \n                                                out_channels=3,\n                                                kernel_size=4, \n                                                stride=2, padding=1)\n\n    def forward(self, inputs):\n        x = self._conv_1(inputs)\n        \n        x = self._residual_stack(x)\n        \n        x = self._conv_trans_1(x)\n        x = F.relu(x)\n\n        x = self._conv_trans_2(x)\n        x = F.relu(x)\n        \n        return self._conv_trans_2_1(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test input and output\nnum_hiddens = 64\nnum_residual_hiddens = 32\nnum_residual_layers = 2\n\ntestDecoder = Decoder(64, num_hiddens, 2, num_residual_hiddens)\ntestDecoder.cuda()\ninput_tensor = torch.Tensor(np.random.normal(size = [32, 64, 8, 8]))\nprint(input_tensor.shape)\noutput_tensor = testDecoder(input_tensor.cuda())\nprint(output_tensor.shape)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Train\n\nWe use the hyperparameters from the author's code:"},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 16\nnum_training_updates = 600000\n\nnum_hiddens = 64\nnum_residual_hiddens = 32\nnum_residual_layers = 2\n\nembedding_dim = 64\nnum_embeddings = 1024\n\ncommitment_cost = 0.25\n\ndecay = 0.99\n\nlearning_rate = 3e-4","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"training_loader = DataLoader(training_data, \n                             batch_size=batch_size, \n                             shuffle=True,\n                             pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"validation_loader = DataLoader(validation_data,\n                               batch_size=batch_size,\n                               shuffle=True,\n                               pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Model(nn.Module):\n    def __init__(self, num_hiddens, num_residual_layers, num_residual_hiddens, \n                 num_embeddings, embedding_dim, commitment_cost, decay=0):\n        super(Model, self).__init__()\n        \n        self._encoder = Encoder(3, num_hiddens,\n                                num_residual_layers, \n                                num_residual_hiddens)\n        self._pre_vq_conv = nn.Conv2d(in_channels=num_hiddens, \n                                      out_channels=embedding_dim,\n                                      kernel_size=1, \n                                      stride=1)\n        if decay > 0.0:\n            self._vq_vae = VectorQuantizerEMA(num_embeddings, embedding_dim, \n                                              commitment_cost, decay)\n        else:\n            self._vq_vae = VectorQuantizer(num_embeddings, embedding_dim,\n                                           commitment_cost)\n        self._decoder = Decoder(embedding_dim,\n                                num_hiddens, \n                                num_residual_layers, \n                                num_residual_hiddens)\n\n    def forward(self, x):\n        z = self._encoder(x)\n        z = self._pre_vq_conv(z)\n        loss, quantized, perplexity, _ = self._vq_vae(z)\n        x_recon = self._decoder(quantized)\n\n        return loss, x_recon, perplexity","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Model(num_hiddens, num_residual_layers, num_residual_hiddens,\n              num_embeddings, embedding_dim, \n              commitment_cost, decay).to(device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Test input and output\ninput_tensor = torch.Tensor(np.random.normal(size = [32, 3, 64, 64]))\nprint(input_tensor.shape)\n_, output_tensor, _ = model(input_tensor.cuda())\nprint(output_tensor.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"optimizer = optim.Adam(model.parameters(), lr=learning_rate, amsgrad=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.train()\ntrain_res_recon_error = []\ntrain_res_perplexity = []\n\nfor i in tqdm_notebook(xrange(num_training_updates)):\n    (data, _) = next(iter(training_loader))\n    data = data.to(device)\n    optimizer.zero_grad()\n\n    vq_loss, data_recon, perplexity = model(data)\n    recon_error = torch.mean((data_recon - data)**2) / data_variance\n    loss = recon_error + vq_loss\n    loss.backward()\n\n    optimizer.step()\n    \n    train_res_recon_error.append(recon_error.item())\n    train_res_perplexity.append(perplexity.item())\n\n    if (i+1) % 1000 == 0:\n        print('%d iterations' % (i+1))\n        print('recon_error: %.3f' % np.mean(train_res_recon_error[-100:]))\n        print('perplexity: %.3f' % np.mean(train_res_perplexity[-100:]))\n        print()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Plot Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_res_recon_error_smooth = savgol_filter(train_res_recon_error, 201, 7)\ntrain_res_perplexity_smooth = savgol_filter(train_res_perplexity, 201, 7)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"f = plt.figure(figsize=(16,8))\nax = f.add_subplot(1,2,1)\nax.plot(train_res_recon_error_smooth)\nax.set_yscale('log')\nax.set_title('Smoothed NMSE.')\nax.set_xlabel('iteration')\n\nax = f.add_subplot(1,2,2)\nax.plot(train_res_perplexity_smooth)\nax.set_title('Smoothed Average codebook usage (perplexity).')\nax.set_xlabel('iteration')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## View Reconstructions"},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\n(valid_originals, _) = next(iter(validation_loader))\nvalid_originals = valid_originals.to(device)\n\nvq_output_eval = model._pre_vq_conv(model._encoder(valid_originals))\n_, valid_quantize, _, _ = model._vq_vae(vq_output_eval)\nvalid_reconstructions = model._decoder(valid_quantize)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"(train_originals, _) = next(iter(training_loader))\ntrain_originals = train_originals.to(device)\n_, train_reconstructions, _, _ = model._vq_vae(train_originals)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def show(img):\n    npimg = img.numpy()\n    fig = plt.imshow(np.transpose(npimg, (1,2,0)), interpolation='nearest')\n    fig.axes.get_xaxis().set_visible(False)\n    fig.axes.get_yaxis().set_visible(False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from PIL import Image\n\na = (valid_reconstructions.cpu().detach().numpy() * 255)[0, :, :, :].transpose(2, 1, 0)\na.shape\nimg = Image.fromarray(a.astype('uint8'))\nimg","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show(make_grid(valid_reconstructions.cpu().data)+0.5)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show(make_grid(valid_originals.cpu()+0.5))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## View Embedding"},{"metadata":{"trusted":true},"cell_type":"code","source":"proj = umap.UMAP(n_neighbors=5,\n                 min_dist=0.2,\n                 metric='cosine').fit_transform(model._vq_vae._embedding.weight.data.cpu())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.scatter(proj[:,0], proj[:,1], alpha=0.3)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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}