{"cells":[{"metadata":{},"cell_type":"markdown","source":"# 1 VQ-VAE \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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"ComputeLB = True\nDogsOnly = True\n\nimport gc\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":"%%time\nfrom 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(imagesIn[i]/255.0).transpose(2, 0, 1) for i in range(idxIn)])\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]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"imagesIntorch.shape","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# 3 Model-VQ VAE 2\n- VQVAE2\n- PixelCNN"},{"metadata":{},"cell_type":"markdown","source":"## 3.1 General function"},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"\"\"\"\nVector-Quantization for the VQ-VAE itself.\n\"\"\"\n\nimport random\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef vq_loss(inputs, embedded, commitment=0.25):\n    \"\"\"\n    Compute the codebook and commitment losses for an\n    input-output pair from a VQ layer.\n    \"\"\"\n    return (torch.mean(torch.pow(inputs.detach() - embedded, 2)) +\n            commitment * torch.mean(torch.pow(inputs - embedded.detach(), 2)))\n\n\nclass VQ(nn.Module):\n    \"\"\"\n    A vector quantization layer.\n    This layer takes continuous inputs and produces a few\n    different types of outputs, including a discretized\n    output, a commitment loss, a codebook loss, etc.\n    Args:\n        num_channels: the depth of the input Tensors.\n        num_latents: the number of latent values in the\n          dictionary to choose from.\n        dead_rate: the number of forward passes after\n          which a dictionary entry is considered dead if\n          it has not been used.\n    \"\"\"\n\n    def __init__(self, num_channels, num_latents, dead_rate=100):\n        super().__init__()\n        self.num_channels = num_channels\n        self.num_latents = num_latents\n        self.dead_rate = dead_rate\n\n        self.dictionary = nn.Parameter(torch.randn(num_latents, num_channels))\n        self.usage_count = nn.Parameter(dead_rate * torch.ones(num_latents).long(),\n                                        requires_grad=False)\n        self._last_batch = None\n\n    def embed(self, idxs):\n        \"\"\"\n        Convert encoded indices into embeddings.\n        Args:\n            idxs: an [N x H x W] or [N] Tensor.\n        Returns:\n            An [N x H x W x C] or [N x C] Tensor.\n        \"\"\"\n        embedded = F.embedding(idxs, self.dictionary)\n        if len(embedded.shape) == 4:\n            # NHWC to NCHW\n            embedded = embedded.permute(0, 3, 1, 2).contiguous()\n        return embedded\n\n    def forward(self, inputs):\n        \"\"\"\n        Apply vector quantization.\n        If the module is in training mode, this will also\n        update the usage tracker and re-initialize dead\n        dictionary entries.\n        Args:\n            inputs: the input Tensor. Either [N x C] or\n              [N x C x H x W].\n        Returns:\n            A tuple (embedded, embedded_pt, idxs):\n              embedded: the new [N x C x H x W] Tensor\n                which passes gradients to the dictionary.\n              embedded_pt: like embedded, but with a\n                passthrough gradient estimator. Gradients\n                through this pass directly to the inputs.\n              idxs: a [N x H x W] Tensor of Longs\n                indicating the chosen dictionary entries.\n        \"\"\"\n        channels_last = inputs\n        if len(inputs.shape) == 4:\n            # NCHW to NHWC\n            channels_last = inputs.permute(0, 2, 3, 1).contiguous()\n\n        diffs = embedding_distances(self.dictionary, channels_last)\n        idxs = torch.argmin(diffs, dim=-1)\n        embedded = self.embed(idxs)\n        embedded_pt = embedded.detach() + (inputs - inputs.detach())\n\n        if self.training:\n            self._update_tracker(idxs)\n            self._last_batch = channels_last.detach()\n\n        return embedded, embedded_pt, idxs\n\n    def revive_dead_entries(self, inputs=None):\n        \"\"\"\n        Use the dictionary usage tracker to re-initialize\n        entries that aren't being used often.\n        Args:\n          inputs: a batch of inputs from which random\n            values are sampled for new entries. If None,\n            the previous input to forward() is used.\n        \"\"\"\n        if inputs is None:\n            assert self._last_batch is not None, ('cannot revive dead entries until a batch has ' +\n                                                  'been run')\n            inputs = self._last_batch\n        counts = self.usage_count.detach().cpu().numpy()\n        new_dictionary = None\n        inputs_numpy = None\n        for i, count in enumerate(counts):\n            if count:\n                continue\n            if new_dictionary is None:\n                new_dictionary = self.dictionary.detach().cpu().numpy()\n            if inputs_numpy is None:\n                inputs_numpy = inputs.detach().cpu().numpy().reshape([-1, inputs.shape[-1]])\n            new_dictionary[i] = random.choice(inputs_numpy)\n            counts[i] = self.dead_rate\n        if new_dictionary is not None:\n            dict_tensor = torch.from_numpy(new_dictionary).to(self.dictionary.device)\n            counts_tensor = torch.from_numpy(counts).to(self.usage_count.device)\n            self.dictionary.data.copy_(dict_tensor)\n            self.usage_count.data.copy_(counts_tensor)\n\n    def _update_tracker(self, idxs):\n        raw_idxs = set(idxs.detach().cpu().numpy().flatten())\n        update = -np.ones([self.num_latents], dtype=np.int)\n        for idx in raw_idxs:\n            update[idx] = self.dead_rate\n        self.usage_count.data.add_(torch.from_numpy(update).to(self.usage_count.device).long())\n        self.usage_count.data.clamp_(0, self.dead_rate)\n\n\ndef embedding_distances(dictionary, tensor):\n    \"\"\"\n    Compute distances between every embedding in a\n    dictionary and every vector in a Tensor.\n    This will not generate a huge intermediate Tensor,\n    unlike the naive implementation.\n    Args:\n        dictionary: a [D x C] Tensor.\n        tensor: a [... x C] Tensor.\n    Returns:\n        A [... x D] Tensor of distances.\n    \"\"\"\n    dict_norms = torch.sum(torch.pow(dictionary, 2), dim=-1)\n    tensor_norms = torch.sum(torch.pow(tensor, 2), dim=-1)\n\n    # Work-around for https://github.com/pytorch/pytorch/issues/18862.\n    exp_tensor = tensor[..., None].view(-1, tensor.shape[-1], 1)\n    exp_dict = dictionary[None].expand(exp_tensor.shape[0], *dictionary.shape)\n    dots = torch.bmm(exp_dict, exp_tensor)[..., 0]\n    dots = dots.view(*tensor.shape[:-1], dots.shape[-1])\n\n    return -2 * dots + dict_norms + tensor_norms[..., None]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nAn implementation of the hierarchical VQ-VAE.\nSee https://arxiv.org/abs/1906.00446.\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n# from .vq import VQ, vq_loss\n\n\nclass Encoder(nn.Module):\n    \"\"\"\n    An abstract VQ-VAE encoder, which takes input Tensors,\n    shrinks them, and quantizes the result.\n    Sub-classes should overload the encode() method.\n    Args:\n        num_channels: the number of channels in the latent\n          codebook.\n        num_latents: the number of entries in the latent\n          codebook.\n        kwargs: arguments to pass to the VQ layer.\n    \"\"\"\n\n    def __init__(self, num_channels, num_latents, **kwargs):\n        super().__init__()\n        self.vq = VQ(num_channels, num_latents, **kwargs)\n\n    def encode(self, x):\n        \"\"\"\n        Encode a Tensor before the VQ layer.\n        Args:\n            x: the input Tensor.\n        Returns:\n            A Tensor with the correct number of output\n              channels (according to self.vq).\n        \"\"\"\n        raise NotImplementedError\n\n    def forward(self, x):\n        \"\"\"\n        Apply the encoder.\n        See VQ.forward() for return values.\n        \"\"\"\n        return self.vq(self.encode(x))\n\n\nclass QuarterEncoder(Encoder):\n    \"\"\"\n    The encoder from the original VQ-VAE paper that cuts\n    the dimensions down by a factor of 4 in both\n    directions.\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, num_latents, **kwargs):\n        super().__init__(out_channels, num_latents, **kwargs)\n        self.conv1 = nn.Conv2d(in_channels, out_channels, 4, stride=2)\n        self.conv2 = nn.Conv2d(out_channels, out_channels, 4, stride=2)\n        self.residual1 = _make_residual(out_channels)\n        self.residual2 = _make_residual(out_channels)\n\n    def encode(self, x):\n        # Padding is uneven, so we make the right and\n        # bottom more padded arbitrarily.\n        x = F.pad(x, (1, 2, 1, 2))\n        x = self.conv1(x)\n        x = F.relu(x)\n        x = F.pad(x, (1, 2, 1, 2))\n        x = self.conv2(x)\n        x = x + self.residual1(x)\n        x = x + self.residual2(x)\n        return x\n\n\nclass HalfEncoder(Encoder):\n    \"\"\"\n    An encoder that cuts the input size in half in both\n    dimensions.\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels, num_latents, **kwargs):\n        super().__init__(out_channels, num_latents, **kwargs)\n        self.conv = nn.Conv2d(in_channels, out_channels, 3, stride=2, padding=1)\n        self.residual1 = _make_residual(out_channels)\n        self.residual2 = _make_residual(out_channels)\n\n    def encode(self, x):\n        x = self.conv(x)\n        x = x + self.residual1(x)\n        x = x + self.residual2(x)\n        return x\n\n\nclass Decoder(nn.Module):\n    \"\"\"\n    An abstract VQ-VAE decoder, which takes a stack of\n    (differently-sized) input Tensors and produces a\n    predicted output Tensor.\n    Sub-classes should overload the forward() method.\n    \"\"\"\n\n    def forward(self, inputs):\n        \"\"\"\n        Apply the decoder to a list of inputs.\n        Args:\n            inputs: a sequence of input Tensors. There may\n              be more than one in the case of a hierarchy,\n              in which case the top levels come first.\n        Returns:\n            A decoded Tensor.\n        \"\"\"\n        raise NotImplementedError\n\n\nclass QuarterDecoder(Decoder):\n    \"\"\"\n    The decoder from the original VQ-VAE paper that\n    upsamples the dimensions by a factor of 4 in both\n    directions.\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.residual1 = _make_residual(in_channels)\n        self.residual2 = _make_residual(in_channels)\n        self.conv1 = nn.ConvTranspose2d(in_channels, in_channels, 4, stride=2, padding=1)\n        self.conv2 = nn.ConvTranspose2d(in_channels, out_channels, 4, stride=2, padding=1)\n\n    def forward(self, inputs):\n        assert len(inputs) == 1\n        x = inputs[0]\n        x = x + self.residual1(x)\n        x = x + self.residual2(x)\n        x = F.relu(x)\n        x = self.conv1(x)\n        x = F.relu(x)\n        x = self.conv2(x)\n        return x\n\n\nclass HalfDecoder(Decoder):\n    \"\"\"\n    A decoder that upsamples by a factor of 2 in both\n    dimensions.\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.residual1 = _make_residual(in_channels)\n        self.residual2 = _make_residual(in_channels)\n        self.conv = nn.ConvTranspose2d(in_channels, out_channels, 4, stride=2, padding=1)\n\n    def forward(self, inputs):\n        assert len(inputs) == 1\n        x = inputs[0]\n        x = x + self.residual1(x)\n        x = x + self.residual2(x)\n        x = F.relu(x)\n        x = self.conv(x)\n        return x\n\n\nclass HalfQuarterDecoder(Decoder):\n    \"\"\"\n    A decoder that takes two inputs. The first one is\n    upsampled by a factor of two, and then combined with\n    the second input which is further upsampled by a\n    factor of four.\n    \"\"\"\n\n    def __init__(self, in_channels, out_channels):\n        super().__init__()\n        self.residual1 = _make_residual(in_channels)\n        self.residual2 = _make_residual(in_channels)\n        self.conv1 = nn.ConvTranspose2d(in_channels, in_channels, 4, stride=2, padding=1)\n        self.conv2 = nn.Conv2d(in_channels * 2, in_channels, 3, padding=1)\n        self.residual3 = _make_residual(in_channels)\n        self.residual4 = _make_residual(in_channels)\n        self.conv3 = nn.ConvTranspose2d(in_channels, in_channels, 4, stride=2, padding=1)\n        self.conv4 = nn.ConvTranspose2d(in_channels, out_channels, 4, stride=2, padding=1)\n\n    def forward(self, inputs):\n        assert len(inputs) == 2\n\n        # Upsample the top input to match the shape of the\n        # bottom input.\n        x = inputs[0]\n        x = x + self.residual1(x)\n        x = x + self.residual2(x)\n        x = F.relu(x)\n        x = self.conv1(x)\n        x = F.relu(x)\n\n        # Mix together the bottom and top inputs.\n        x = torch.cat([x, inputs[1]], dim=1)\n        x = self.conv2(x)\n\n        x = x + self.residual3(x)\n        x = x + self.residual4(x)\n        x = F.relu(x)\n        x = self.conv3(x)\n        x = F.relu(x)\n        x = self.conv4(x)\n        return x\n\n\nclass VQVAE(nn.Module):\n    \"\"\"\n    A complete VQ-VAE hierarchy.\n    There are N encoders, stored from the bottom level to\n    the top level, and N decoders stored from top to\n    bottom.\n    \"\"\"\n\n    def __init__(self, encoders, decoders):\n        super().__init__()\n        assert len(encoders) == len(decoders)\n        self.encoders = encoders\n        self.decoders = decoders\n        for i, enc in enumerate(encoders):\n            self.add_module('encoder_%d' % i, enc)\n        for i, dec in enumerate(decoders):\n            self.add_module('decoder_%d' % i, dec)\n\n    def forward(self, inputs, commitment=0.25):\n        \"\"\"\n        Compute training losses for a batch of inputs.\n        Args:\n            inputs: the input Tensor. If this is a Tensor\n              of integers, then cross-entropy loss will be\n              used for the final decoder. Otherwise, MSE\n              will be used.\n            commitment: the commitment loss coefficient.\n        Returns:\n            A dict of Tensors, containing at least:\n              loss: the total training loss.\n              losses: the MSE/log-loss from each decoder.\n              reconstructions: a reconstruction Tensor\n                from each decoder.\n              embedded: outputs from every encoder, passed\n                through the vector-quantization table.\n                Ordered from bottom to top level.\n        \"\"\"\n        all_encoded = [inputs]\n        all_vq_outs = []\n        total_vq_loss = 0.0\n        total_recon_loss = 0.0\n        for encoder in self.encoders:\n            encoded = encoder.encode(all_encoded[-1])\n            embedded, embedded_pt, _ = encoder.vq(encoded)\n            all_encoded.append(encoded)\n            all_vq_outs.append(embedded_pt)\n            total_vq_loss = total_vq_loss + vq_loss(encoded, embedded, commitment=commitment)\n        losses = []\n        reconstructions = []\n        for i, decoder in enumerate(self.decoders):\n            dec_inputs = all_vq_outs[::-1][:i + 1]\n            target = all_encoded[::-1][i + 1]\n            recon = decoder(dec_inputs)\n            reconstructions.append(recon)\n            if target.dtype.is_floating_point:\n                recon_loss = torch.mean(torch.pow(recon - target.detach(), 2))\n            else:\n                recon_loss = F.cross_entropy(recon.view(-1, recon.shape[-1]), target.view(-1))\n            total_recon_loss = total_recon_loss + recon_loss\n            losses.append(recon_loss)\n        return {\n            'loss': total_vq_loss + total_recon_loss,\n            'losses': losses,\n            'reconstructions': reconstructions,\n            'embedded': all_vq_outs,\n        }\n\n    def revive_dead_entries(self):\n        \"\"\"\n        Revive dead entries from all of the VQ layers.\n        Only call this once the encoders have all been\n        through a forward pass in training mode.\n        \"\"\"\n        for enc in self.encoders:\n            enc.vq.revive_dead_entries()\n\n    def full_reconstructions(self, inputs):\n        \"\"\"\n        Compute reconstructions of the inputs using all\n        the different layers of the hierarchy.\n        The first reconstruction uses only information\n        from the top-level codes, the second uses only\n        information from the top-level and second-to-top\n        level codes, etc.\n        This is not forward(inputs)['reconstructions'],\n        since said reconstructions are simply each level's\n        reconstruction of the next level's features.\n        Instead, full_reconstructions reconstructs the\n        original inputs.\n        \"\"\"\n        terms = self(inputs)\n        layer_recons = []\n        for encoder, recon in zip(self.encoders[:-1][::-1], terms['reconstructions'][:-1]):\n            _, embedded_pt, _ = encoder.vq(recon)\n            layer_recons.append(embedded_pt)\n        hierarchy_size = len(self.decoders)\n        results = []\n        for i in range(hierarchy_size - 1):\n            num_actual = i + 1\n            dec_in = terms['embedded'][-num_actual:][::-1] + layer_recons[num_actual - 1:]\n            results.append(self.decoders[-1](dec_in))\n        results.append(terms['reconstructions'][-1])\n        return results\n\n\ndef _make_residual(channels):\n    return nn.Sequential(\n        nn.ReLU(),\n        nn.Conv2d(channels, channels, 3, padding=1),\n        nn.ReLU(),\n        nn.Conv2d(channels, channels, 1),\n    )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport torch\n\n# from .vq import embedding_distances\n\ndef test_embedding_distances():\n    dictionary = torch.randn(15, 7)\n    tensor = torch.randn(3, 3, 7)\n    with torch.no_grad():\n        actual = embedding_distances(dictionary, tensor).numpy()\n        print(actual.shape)\n        expected = naive_embedding_distances(dictionary, tensor).numpy()\n        print(expected.shape)\n        assert np.allclose(actual, expected, atol=1e-4)\n\n\ndef naive_embedding_distances(dictionary, tensor):\n    return torch.sum(torch.pow(tensor[..., None, :] - dictionary, 2), dim=-1)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"test_embedding_distances()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nAn implementation of multi-head attention, based off of\nhttps://github.com/unixpickle/xformer\n\"\"\"\n\nimport math\n\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass PixelAttention(nn.Module):\n    \"\"\"\n    An attention layer that operates on images.\n    Args:\n        num_channels: the input image depth.\n        num_heads: the number of attention heads.\n    \"\"\"\n\n    def __init__(self, num_channels, num_heads=8):\n        super().__init__()\n        self.attention = MaskedAttention(num_channels, num_heads=num_heads)\n\n    def forward(self, *images, conds=None):\n        \"\"\"\n        Apply masked attention to a batch of images.\n        Args:\n            images: one or more [N x C x H x W] Tensors.\n            conds: ignored. Here for compatibility with\n              the PixelCNN aggregator.\n        Returns:\n            A new list of [N x C x H x W] Tensors.\n        \"\"\"\n        results = []\n        for image in images:\n            batch, num_channels, height, width = image.shape\n            result = image.permute(0, 2, 3, 1)\n            result = result.view(batch, height * width, num_channels)\n            result = self.attention(result)\n            result = result.view(batch, height, width, num_channels)\n            result = result.permute(0, 3, 1, 2)\n            results.append(result + image)\n        if len(results) == 1:\n            return results[0]\n        return tuple(results)\n\n\nclass MaskedAttention(nn.Module):\n    \"\"\"\n    An attention layer that operates on sequences of the\n    shape [N x T x C], where N is the batch size, T is the\n    number of timesteps, and C is the number of channels.\n    Args:\n        num_channels: the number of channels in the input\n          sequences.\n        num_heads: the number of attention heads to use.\n    \"\"\"\n\n    def __init__(self, num_channels, num_heads=8):\n        super().__init__()\n\n        assert not num_channels % num_heads, 'heads must evenly divide channels'\n        self.num_channels = num_channels\n        self.num_heads = num_heads\n\n        self.kqv_projection = nn.Linear(num_channels, num_channels * 3)\n        self.mix_heads = nn.Linear(num_channels, num_channels)\n\n    def forward(self, sequence):\n        \"\"\"\n        Apply masked multi-head attention.\n        Args:\n            sequence: an [N x T x C] Tensor.\n        Returns:\n            A new [N x T x C] Tensor.\n        \"\"\"\n        projected = self.kqv_projection(sequence)\n        kqv = torch.split(projected, self.num_channels, dim=-1)\n        keys, queries, values = [self._split_heads(x) for x in kqv]\n        logits = torch.bmm(queries, keys.permute(0, 2, 1))\n        logits /= math.sqrt(self.num_channels / self.num_heads)\n        logits += self._logit_mask(sequence.shape[1])\n        weights = F.softmax(logits, dim=-1)\n        weighted_sum = torch.bmm(weights, values)\n        combined = self._combine_heads(weighted_sum)\n        return self.mix_heads(combined)\n\n    def _split_heads(self, batch):\n        \"\"\"\n        Split up the channels in a batch into groups, one\n        per head.\n        Args:\n            batch: an [N x T x C] Tensor.\n        Returns:\n            An [N*H x T x C/H] Tensor.\n        \"\"\"\n        batch_size = batch.shape[0]\n        num_steps = batch.shape[1]\n        split_channels = self.num_channels // self.num_heads\n        batch = batch.view(batch_size, num_steps, self.num_heads, split_channels)\n        batch = batch.permute(0, 2, 1, 3).contiguous()\n        batch = batch.view(batch_size * self.num_heads, num_steps, split_channels)\n        return batch\n\n    def _combine_heads(self, batch):\n        \"\"\"\n        Perform the inverse of _split_heads().\n        Args:\n            batch: an [N*H x T x C/H] Tensor.\n        Returns:\n            An [N x T x C] Tensor.\n        \"\"\"\n        batch_size = batch.shape[0] // self.num_heads\n        num_steps = batch.shape[1]\n        split_channels = self.num_channels // self.num_heads\n        batch = batch.view(batch_size, self.num_heads, num_steps, split_channels)\n        batch = batch.permute(0, 2, 1, 3).contiguous()\n        batch = batch.view(batch_size, num_steps, self.num_channels)\n        return batch\n\n    def _logit_mask(self, num_steps):\n        row_indices = np.arange(num_steps)[:, None]\n        col_indices = np.arange(num_steps)[None]\n        upper = (row_indices >= col_indices)\n        mask = np.where(upper, 0, -np.inf).astype(np.float32)\n        return torch.from_numpy(mask).to(next(self.parameters()).device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nAn implementation of the Gated PixelCNN from\nhttps://arxiv.org/abs/1606.05328.\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass PixelCNN(nn.Module):\n    \"\"\"\n    A PixelCNN is a stack of PixelConv layers.\n    \"\"\"\n\n    def __init__(self, *layers):\n        super().__init__()\n        for i, layer in enumerate(layers):\n            self.add_module('layer_%d' % i, layer)\n        self.layers = layers\n\n    def forward(self, images, conds=None):\n        \"\"\"\n        Apply the stack of PixelConv layers.\n        It is assumed that the first layer is a\n        PixelConvA, and the rest are PixelConvB's.\n        This way, the first layer takes one input and the\n        rest take two.\n        Returns:\n            A tuple (vertical, horizontal), one for each\n              of the two directional stacks.\n        \"\"\"\n        outputs = self.layers[0](images, conds=conds)\n        for layer in self.layers[1:]:\n            outputs = layer(*outputs, conds=conds)\n        return outputs\n\n\nclass PixelConv(nn.Module):\n    \"\"\"\n    An abstract base class for PixelCNN layers.\n    \"\"\"\n\n    def __init__(self, depth_in, depth_out, cond_depth=None, horizontal=2, vertical=2):\n        super().__init__()\n        self.depth_in = depth_in\n        self.depth_out = depth_out\n        self.horizontal = horizontal\n        self.vertical = vertical\n\n        self._init_directional_convs()\n        self.vert_to_horiz = nn.Conv2d(depth_out * 2, depth_out * 2, 1)\n        self.cond_layer = None\n        if cond_depth is not None:\n            self.cond_layer = nn.Linear(cond_depth, depth_out * 4)\n\n    def _init_directional_convs(self):\n        raise NotImplementedError\n\n    def _run_stacks(self, vert_in, horiz_in, conds):\n        vert_out = self._run_padded_vertical(vert_in)\n        horiz_out = self._run_padded_horizontal(horiz_in)\n        horiz_out = horiz_out + self.vert_to_horiz(vert_out)\n\n        if conds is not None:\n            cond_bias = self._compute_cond_bias(conds)\n            vert_out = vert_out + cond_bias[:, :self.depth_out*2]\n            horiz_out = horiz_out + cond_bias[:, self.depth_out*2:]\n\n        vert_out = gated_activation(vert_out)\n        horiz_out = gated_activation(horiz_out)\n        return vert_out, horiz_out\n\n    def _run_padded_vertical(self, vert_in):\n        raise NotImplementedError\n\n    def _run_padded_horizontal(self, horiz_in):\n        raise NotImplementedError\n\n    def _compute_cond_bias(self, conds):\n        if len(conds.shape) == 2:\n            outputs = self.cond_layer(conds)\n            return outputs.view(-1, outputs.shape[1], 1, 1)\n        assert len(conds.shape) == 4\n        conds_perm = conds.permute(0, 2, 3, 1)\n        outputs = self.cond_layer(conds_perm)\n        return outputs.permute(0, 3, 1, 2)\n\n\nclass PixelConvA(PixelConv):\n    \"\"\"\n    The first layer in a PixelCNN. This layer is unlike\n    the other layers, in that it does not allow the stack\n    to see the current pixel.\n    Args:\n        depth_in: the number of input filters.\n        depth_out: the number of output filters.\n        cond_depth: the number of conditioning channels.\n          If None, this is an unconditional model.\n        horizontal: the receptive field of the horizontal\n          stack.\n        vertical: the receptive field of the vertical\n          stack.\n    \"\"\"\n\n    def __init__(self, depth_in, depth_out, cond_depth=None, horizontal=2, vertical=2):\n        super().__init__(depth_in, depth_out, cond_depth=cond_depth, horizontal=2, vertical=2)\n\n    def forward(self, images, conds=None):\n        \"\"\"\n        Apply the layer to some images, producing latents.\n        Args:\n            images: an NCHW batch of images.\n            conds: an optional conditioning value. If set,\n              either an NCHW Tensor or an NxM Tensor.\n        Returns:\n            A tuple (vertical, horizontal), one for each\n              of the two directional stacks.\n        \"\"\"\n        return self._run_stacks(images, images, conds)\n\n    def _init_directional_convs(self):\n        self.vertical_conv = nn.Conv2d(self.depth_in, self.depth_out * 2,\n                                       (self.vertical, self.horizontal*2 + 1))\n        self.horizontal_conv = nn.Conv2d(self.depth_in, self.depth_out * 2, (1, self.horizontal))\n\n    def _run_padded_vertical(self, vert_in):\n        vert_pad = (self.horizontal, self.horizontal, self.vertical, 0)\n        return self.vertical_conv(F.pad(vert_in, vert_pad))[:, :, :-1, :]\n\n    def _run_padded_horizontal(self, horiz_in):\n        return self.horizontal_conv(F.pad(horiz_in, (self.horizontal, 0, 0, 0)))[:, :, :, :-1]\n\n\nclass PixelConvB(PixelConv):\n    \"\"\"\n    Any layer except the first in a PixelCNN.\n    Args:\n        depth_in: the number of input filters.\n        cond_depth: the number of conditioning channels.\n          If None, this is an unconditional model.\n        horizontal: the receptive field of the horizontal\n          stack.\n        vertical: the receptive field of the vertical\n          stack.\n    \"\"\"\n\n    def __init__(self, depth_in, cond_depth=None, norm=False, horizontal=2, vertical=2):\n        super().__init__(depth_in, depth_in, cond_depth=cond_depth, horizontal=horizontal,\n                         vertical=vertical)\n        self.horiz_residual = nn.Conv2d(depth_in, depth_in, 1)\n        self.vert_norm = lambda x: x\n        self.horiz_norm = lambda x: x\n        if norm:\n            self.vert_norm = ChannelNorm(depth_in)\n            self.horiz_norm = ChannelNorm(depth_in)\n\n    def forward(self, vert_in, horiz_in, conds=None):\n        \"\"\"\n        Apply the layer to the outputs of previous\n        vertical and horizontal stacks.\n        Args:\n            vert_in: an NCHW Tensor.\n            horiz_in: an NCHW Tensor.\n            conds: an optional conditioning value. If set,\n              either an NCHW Tensor or an NxM Tensor.\n        Returns:\n            A tuple (vertical, horizontal), one for each\n              of the two directional stacks.\n        \"\"\"\n        vert_out, horiz_out = self._run_stacks(vert_in, horiz_in, conds)\n        horiz_out = horiz_in + self.horiz_norm(self.horiz_residual(horiz_out))\n        return self.vert_norm(vert_out), horiz_out\n\n    def _init_directional_convs(self):\n        self.vertical_conv = nn.Conv2d(self.depth_in, self.depth_out * 2,\n                                       (self.vertical + 1, self.horizontal*2 + 1))\n        self.horizontal_conv = nn.Conv2d(self.depth_in, self.depth_out * 2,\n                                         (1, self.horizontal + 1))\n\n    def _run_padded_vertical(self, vert_in):\n        vert_pad = (self.horizontal, self.horizontal, self.vertical, 0)\n        return self.vertical_conv(F.pad(vert_in, vert_pad))\n\n    def _run_padded_horizontal(self, horiz_in):\n        return self.horizontal_conv(F.pad(horiz_in, (self.horizontal, 0, 0, 0)))\n\n\nclass ChannelNorm(nn.Module):\n    \"\"\"\n    A layer which applies layer normalization to the\n    channels at each spacial location separately.\n    \"\"\"\n\n    def __init__(self, num_channels):\n        super().__init__()\n        self.norm = nn.LayerNorm((num_channels,))\n\n    def forward(self, x):\n        x = x.permute(0, 2, 3, 1).contiguous()\n        x = self.norm(x)\n        x = x.permute(0, 3, 1, 2).contiguous()\n        return x\n\n\ndef gated_activation(outputs):\n    depth = outputs.shape[1] // 2\n    tanh = torch.tanh(outputs[:, :depth])\n    sigmoid = torch.sigmoid(outputs[:, depth:])\n    return tanh * sigmoid","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from math import cos, pi, floor, sin\n\nfrom torch.optim import lr_scheduler\n\nclass CosineLR(lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, lr_min, lr_max, step_size):\n        self.lr_min = lr_min\n        self.lr_max = lr_max\n        self.step_size = step_size\n        self.iteration = 0\n\n        super().__init__(optimizer, -1)\n\n    def get_lr(self):\n        lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (\n            1 + cos(self.iteration / self.step_size * pi)\n        )\n        self.iteration += 1\n\n        if self.iteration == self.step_size:\n            self.iteration = 0\n\n        return [lr for base_lr in self.base_lrs]\n\n\nclass PowerLR(lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, lr_min, lr_max, warmup):\n        self.lr_min = lr_min\n        self.lr_max = lr_max\n        self.warmup = warmup\n        self.iteration = 0\n\n        super().__init__(optimizer, -1)\n\n    def get_lr(self):\n        if self.iteration < self.warmup:\n            lr = (\n                self.lr_min + (self.lr_max - self.lr_min) / self.warmup * self.iteration\n            )\n\n        else:\n            lr = self.lr_max * (self.iteration - self.warmup + 1) ** -0.5\n\n        self.iteration += 1\n\n        return [lr for base_lr in self.base_lrs]\n\n\nclass SineLR(lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, lr_min, lr_max, step_size):\n        self.lr_min = lr_min\n        self.lr_max = lr_max\n        self.step_size = step_size\n        self.iteration = 0\n\n        super().__init__(optimizer, -1)\n\n    def get_lr(self):\n        lr = self.lr_min + (self.lr_max - self.lr_min) * sin(\n            self.iteration / self.step_size * pi\n        )\n        self.iteration += 1\n\n        if self.iteration == self.step_size:\n            self.iteration = 0\n\n        return [lr for base_lr in self.base_lrs]\n\n\nclass LinearLR(lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, lr_min, lr_max, warmup, step_size):\n        self.lr_min = lr_min\n        self.lr_max = lr_max\n        self.step_size = step_size\n        self.warmup = warmup\n        self.iteration = 0\n\n        super().__init__(optimizer, -1)\n\n    def get_lr(self):\n        if self.iteration < self.warmup:\n            lr = self.lr_max\n\n        else:\n            lr = self.lr_max + (self.iteration - self.warmup) * (\n                self.lr_min - self.lr_max\n            ) / (self.step_size - self.warmup)\n        self.iteration += 1\n\n        if self.iteration == self.step_size:\n            self.iteration = 0\n\n        return [lr for base_lr in self.base_lrs]\n\n\nclass CLR(lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, lr_min, lr_max, step_size):\n        self.epoch = 0\n        self.lr_min = lr_min\n        self.lr_max = lr_max\n        self.current_lr = lr_min\n        self.step_size = step_size\n\n        super().__init__(optimizer, -1)\n\n    def get_lr(self):\n        cycle = floor(1 + self.epoch / (2 * self.step_size))\n        x = abs(self.epoch / self.step_size - 2 * cycle + 1)\n        lr = self.lr_min + (self.lr_max - self.lr_min) * max(0, 1 - x)\n        self.current_lr = lr\n\n        self.epoch += 1\n\n        return [lr for base_lr in self.base_lrs]\n\n\nclass Warmup(lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, model_dim, factor=1, warmup=16000):\n        self.optimizer = optimizer\n        self.model_dim = model_dim\n        self.factor = factor\n        self.warmup = warmup\n        self.iteration = 0\n\n        super().__init__(optimizer, -1)\n\n    def get_lr(self):\n        self.iteration += 1\n        lr = (\n            self.factor\n            * self.model_dim ** (-0.5)\n            * min(self.iteration ** (-0.5), self.iteration * self.warmup ** (-1.5))\n        )\n\n        return [lr for base_lr in self.base_lrs]\n\n\n# Copyright 2019 fastai\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n# Borrowed from https://github.com/fastai/fastai and changed to make it runs like PyTorch lr scheduler\n\n\nclass CycleAnnealScheduler:\n    def __init__(\n        self, optimizer, lr_max, lr_divider, cut_point, step_size, momentum=None\n    ):\n        self.lr_max = lr_max\n        self.lr_divider = lr_divider\n        self.cut_point = step_size // cut_point\n        self.step_size = step_size\n        self.iteration = 0\n        self.cycle_step = int(step_size * (1 - cut_point / 100) / 2)\n        self.momentum = momentum\n        self.optimizer = optimizer\n\n    def get_lr(self):\n        if self.iteration > 2 * self.cycle_step:\n            cut = (self.iteration - 2 * self.cycle_step) / (\n                self.step_size - 2 * self.cycle_step\n            )\n            lr = self.lr_max * (1 + (cut * (1 - 100) / 100)) / self.lr_divider\n\n        elif self.iteration > self.cycle_step:\n            cut = 1 - (self.iteration - self.cycle_step) / self.cycle_step\n            lr = self.lr_max * (1 + cut * (self.lr_divider - 1)) / self.lr_divider\n\n        else:\n            cut = self.iteration / self.cycle_step\n            lr = self.lr_max * (1 + cut * (self.lr_divider - 1)) / self.lr_divider\n\n        return lr\n\n    def get_momentum(self):\n        if self.iteration > 2 * self.cycle_step:\n            momentum = self.momentum[0]\n\n        elif self.iteration > self.cycle_step:\n            cut = 1 - (self.iteration - self.cycle_step) / self.cycle_step\n            momentum = self.momentum[0] + cut * (self.momentum[1] - self.momentum[0])\n\n        else:\n            cut = self.iteration / self.cycle_step\n            momentum = self.momentum[0] + cut * (self.momentum[1] - self.momentum[0])\n\n        return momentum\n\n    def step(self):\n        lr = self.get_lr()\n\n        if self.momentum is not None:\n            momentum = self.get_momentum()\n\n        self.iteration += 1\n\n        if self.iteration == self.step_size:\n            self.iteration = 0\n\n        for group in self.optimizer.param_groups:\n            group['lr'] = lr\n\n            if self.momentum is not None:\n                group['betas'] = (momentum, group['betas'][1])\n\n        return lr\n\n\ndef anneal_linear(start, end, proportion):\n    return start + proportion * (end - start)\n\n\ndef anneal_cos(start, end, proportion):\n    cos_val = cos(pi * proportion) + 1\n\n    return end + (start - end) / 2 * cos_val\n\n\nclass Phase:\n    def __init__(self, start, end, n_iter, anneal_fn):\n        self.start, self.end = start, end\n        self.n_iter = n_iter\n        self.anneal_fn = anneal_fn\n        self.n = 0\n\n    def step(self):\n        self.n += 1\n\n        return self.anneal_fn(self.start, self.end, self.n / self.n_iter)\n\n    def reset(self):\n        self.n = 0\n\n    @property\n    def is_done(self):\n        return self.n >= self.n_iter\n\n\nclass CycleScheduler:\n    def __init__(\n        self,\n        optimizer,\n        lr_max,\n        n_iter,\n        momentum=(0.95, 0.85),\n        divider=25,\n        warmup_proportion=0.3,\n        phase=('linear', 'cos'),\n    ):\n        self.optimizer = optimizer\n\n        phase1 = int(n_iter * warmup_proportion)\n        phase2 = n_iter - phase1\n        lr_min = lr_max / divider\n\n        phase_map = {'linear': anneal_linear, 'cos': anneal_cos}\n\n        self.lr_phase = [\n            Phase(lr_min, lr_max, phase1, phase_map[phase[0]]),\n            Phase(lr_max, lr_min / 1e4, phase2, phase_map[phase[1]]),\n        ]\n\n        self.momentum = momentum\n\n        if momentum is not None:\n            mom1, mom2 = momentum\n            self.momentum_phase = [\n                Phase(mom1, mom2, phase1, phase_map[phase[0]]),\n                Phase(mom2, mom1, phase2, phase_map[phase[1]]),\n            ]\n\n        else:\n            self.momentum_phase = []\n\n        self.phase = 0\n\n    def step(self):\n        lr = self.lr_phase[self.phase].step()\n\n        if self.momentum is not None:\n            momentum = self.momentum_phase[self.phase].step()\n\n        else:\n            momentum = None\n\n        for group in self.optimizer.param_groups:\n            group['lr'] = lr\n\n            if self.momentum is not None:\n                if 'betas' in group:\n                    group['betas'] = (momentum, group['betas'][1])\n\n                else:\n                    group['momentum'] = momentum\n\n        if self.lr_phase[self.phase].is_done:\n            self.phase += 1\n\n        if self.phase >= len(self.lr_phase):\n            for phase in self.lr_phase:\n                phase.reset()\n\n            for phase in self.momentum_phase:\n                phase.reset()\n\n            self.phase = 0\n\n        return lr, momentum\n\n\nclass LRFinder(lr_scheduler._LRScheduler):\n    def __init__(self, optimizer, lr_min, lr_max, step_size, linear=False):\n        ratio = lr_max / lr_min\n        self.linear = linear\n        self.lr_min = lr_min\n        self.lr_mult = (ratio / step_size) if linear else ratio ** (1 / step_size)\n        self.iteration = 0\n        self.lrs = []\n        self.losses = []\n\n        super().__init__(optimizer, -1)\n\n    def get_lr(self):\n        lr = (\n            self.lr_mult * self.iteration\n            if self.linear\n            else self.lr_mult ** self.iteration\n        )\n        lr = self.lr_min + lr if self.linear else self.lr_min * lr\n\n        self.iteration += 1\n        self.lrs.append(lr)\n\n        return [lr for base_lr in self.base_lrs]\n\n    def record(self, loss):\n        self.losses.append(loss)\n\n    def save(self, filename):\n        with open(filename, 'w') as f:\n            for lr, loss in zip(self.lrs, self.losses):\n                f.write('{},{}\\n'.format(lr, loss))\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 3.2 Using Data to Generate Images\n- a example with mnist\n- for generating dog images"},{"metadata":{},"cell_type":"markdown","source":"### 3.2.1 Dog Image\n- model.py\n- train_vae.py\n- train_top.py\n- train_bottom.py\n- sample_top.py"},{"metadata":{"trusted":true},"cell_type":"code","source":"def load_images(train=True, batch_size = 16):\n    while True:\n        for data, _ in create_data_loader(batch_size):\n            yield data\n\n\ndef create_data_loader(batch_size):\n#     mnist = torchvision.datasets.MNIST('./data', train=train, download=True,\n#                                        transform=torchvision.transforms.ToTensor())\n    ds = torch.utils.data.TensorDataset(torch.Tensor(imagesIntorch), torch.Tensor(np.zeros(22125)))\n    return torch.utils.data.DataLoader(ds, batch_size=batch_size, shuffle=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nModels for hierarchical image generation.\n\"\"\"\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\ndef make_vae():\n#     encoders = [QuarterEncoder(3, 128, 512), HalfEncoder(128, 128, 512)]\n#     decoders = [HalfDecoder(128, 128), HalfQuarterDecoder(128, 3)]\n    encoders = [QuarterEncoder(3, 64, 128), HalfEncoder(64, 64, 128)]\n    decoders = [HalfDecoder(64, 64), HalfQuarterDecoder(64, 3)]\n    return VQVAE(encoders, decoders)\n\n\nclass TopPrior(nn.Module):\n    def __init__(self, depth=64, num_heads=2):\n        super().__init__()\n        self.embed = nn.Embedding(128, depth)\n        self.pixel_cnn = PixelCNN(\n            PixelConvA(depth, depth),\n\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelAttention(depth, num_heads=num_heads),\n\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelAttention(depth, num_heads=num_heads),\n\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelAttention(depth, num_heads=num_heads),\n\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelConvB(depth, norm=True),\n            PixelAttention(depth, num_heads=num_heads),\n        )\n        self.out_stack = nn.Sequential(\n            nn.Conv2d(depth * 2, depth, 1),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            Residual1x1(depth),\n            nn.Conv2d(depth, 256, 1),\n        )\n\n    def forward(self, x):\n        x = self.embed(x)\n        x = x.permute(0, 3, 1, 2).contiguous()\n        out1, out2 = self.pixel_cnn(x)\n        return self.out_stack(torch.cat([out1, out2], dim=1))\n\n\nclass BottomPrior(nn.Module):\n    def __init__(self, depth=64, num_heads=2):\n        super().__init__()\n        self.embed_top = nn.Embedding(128, depth)\n        self.embed_bottom = nn.Embedding(128, depth)\n        self.cond_stack = nn.Sequential(\n            Residual3x3(depth),\n            Residual3x3(depth),\n            Residual3x3(depth),\n            Residual3x3(depth),\n            Residual3x3(depth),\n            Residual3x3(depth),\n            Residual3x3(depth),\n            Residual3x3(depth),\n            Residual3x3(depth),\n            nn.ConvTranspose2d(depth, depth, 4, stride=2, padding=1),\n        )\n        self.pixel_cnn = PixelCNN(\n            PixelConvA(depth, depth, cond_depth=depth),\n\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n            PixelConvB(depth, cond_depth=depth, norm=True),\n        )\n        self.out_stack = nn.Sequential(\n            nn.Conv2d(depth * 2, depth, 1),\n            nn.Conv2d(depth, 128, 1),\n        )\n\n    def forward(self, bottom, top):\n        conds = self.embed_top(top)\n        conds = conds.permute(0, 3, 1, 2).contiguous()\n        conds = self.cond_stack(conds)\n\n        out = self.embed_bottom(bottom)\n        out = out.permute(0, 3, 1, 2).contiguous()\n        out1, out2 = self.pixel_cnn(out, conds=conds)\n        return self.out_stack(torch.cat([out1, out2], dim=1))\n\n\nclass Residual1x1(nn.Module):\n    def __init__(self, num_channels):\n        super().__init__()\n        self.conv1 = nn.Conv2d(num_channels, num_channels, 1)\n        self.conv2 = nn.Conv2d(num_channels, num_channels, 1)\n        self.norm = ChannelNorm(num_channels)\n\n    def forward(self, x):\n        inputs = x\n        x = F.relu(x)\n        x = self.conv1(x)\n        x = F.relu(x)\n        x = self.conv2(x)\n        return inputs + self.norm(x)\n\n\nclass Residual3x3(nn.Module):\n    def __init__(self, num_channels):\n        super().__init__()\n        self.conv1 = nn.Conv2d(num_channels, num_channels, 3, padding=1)\n        self.conv2 = nn.Conv2d(num_channels, num_channels, 3, padding=1)\n        self.norm = ChannelNorm(num_channels)\n\n    def forward(self, x):\n        inputs = x\n        x = F.relu(x)\n        x = self.conv1(x)\n        x = F.relu(x)\n        x = self.conv2(x)\n        return inputs + self.norm(x)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !rm vae.pt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nTrain a hierarchical VQ-VAE on 256x256 images.\n\"\"\"\n\nimport argparse\nimport itertools\nimport os\n\nfrom PIL import Image\nimport numpy as np\nimport torch\nimport torch.optim as optim\n\nVAE_PATH = 'vae.pt'\n\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n\ndef main_vae():\n    epochs = 5000\n    batch_size = 16\n    device = torch.device(DEVICE)\n    model = make_vae()\n    if os.path.exists(VAE_PATH):\n        model.load_state_dict(torch.load(VAE_PATH, map_location=DEVICE))\n    model.to(device)\n    optimizer = optim.Adam(model.parameters())\n    schedule = CycleScheduler(optimizer, 0.001, batch_size * epochs, momentum=(0.95, 0.85),)\n    data = load_images(batch_size = batch_size)\n    for i in itertools.count():\n        images = next(data).to(device)\n        terms = model(images)\n        mse, mse_top = terms['losses'][-1].item(), terms['losses'][0].item()\n        print('step %d: mse=%f mse_top=%f' %\n              (i, mse, mse_top))\n        optimizer.zero_grad()\n        terms['loss'].backward()\n        schedule.step()\n        optimizer.step()\n        model.revive_dead_entries()\n        \n        if i == epochs: \n            torch.save(model.state_dict(), VAE_PATH)\n            save_reconstructions(model, images)\n            break\n\ndef save_reconstructions(vae, images):\n    vae.eval()\n    with torch.no_grad():\n        recons = [torch.clamp(x, 0, 1).permute(0, 2, 3, 1).detach().cpu().numpy()\n                  for x in vae.full_reconstructions(images)]\n    vae.train()\n    top_recons, real_recons = recons\n    images = images.permute(0, 2, 3, 1).detach().cpu().numpy()\n\n    columns = np.concatenate([top_recons, real_recons, images], axis=-2)\n    columns = np.concatenate(columns, axis=0)\n    Image.fromarray((columns * 255).astype('uint8')).save('reconstructions.png')\n\nmain_vae()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from PIL import Image\nvae = Image.open('reconstructions.png')\nvae","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !rm top.pt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nTrain the top-level prior.\n\"\"\"\n\nimport itertools\nimport os\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\n\nTOP_PRIOR_PATH = 'top.pt'\n\ndef main_top():\n    epochs = 60000\n    batch_size = 64\n    lr = 3e-4\n    device = torch.device(DEVICE)\n\n    vae = make_vae()\n    vae.load_state_dict(torch.load(VAE_PATH, map_location='cpu'))\n    vae.to(device)\n    vae.eval()\n\n    top_prior = TopPrior()\n    if os.path.exists(TOP_PRIOR_PATH):\n        top_prior.load_state_dict(torch.load(TOP_PRIOR_PATH, map_location='cuda'))\n    top_prior.to(device)\n\n    optimizer = optim.Adam(top_prior.parameters(), lr=lr)\n    loss_fn = nn.CrossEntropyLoss()\n    \n    schedule = CycleScheduler(optimizer, lr, batch_size * epochs, momentum=(0.95, 0.85),)\n    data = load_images(batch_size = batch_size)\n    for i in itertools.count():\n        images = next(data).to(device)\n        _, _, encoded = vae.encoders[1](vae.encoders[0].encode(images))\n        logits = top_prior(encoded)\n        logits = logits.permute(0, 2, 3, 1).contiguous()\n        logits = logits.view(-1, logits.shape[-1])\n        loss = loss_fn(logits, encoded.view(-1))\n        print('step %d: loss=%f' % (i, loss.item()))\n        optimizer.zero_grad()\n        loss.backward()\n        schedule.step()\n        optimizer.step()\n        if not i % 30:\n            torch.save(top_prior.state_dict(), TOP_PRIOR_PATH)\n            \n        if i == epochs: break\n\nmain_top()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# !rm bottom.pt","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nTrain the bottom-level prior.\n\"\"\"\n\nimport itertools\nimport os\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\n\nBOTTOM_PRIOR_PATH = 'bottom.pt'\n\ndef main_bottom():\n    epochs = 60000\n    batch_size = 64\n    lr = 3e-4\n    device = torch.device(DEVICE)\n\n    vae = make_vae()\n    vae.load_state_dict(torch.load(VAE_PATH, map_location='cuda'))\n    vae.to(device)\n    vae.eval()\n\n    bottom_prior = BottomPrior()\n    if os.path.exists(BOTTOM_PRIOR_PATH):\n        bottom_prior.load_state_dict(torch.load(BOTTOM_PRIOR_PATH, map_location='cuda'))\n    bottom_prior.to(device)\n\n    optimizer = optim.Adam(bottom_prior.parameters(), lr=lr)\n    loss_fn = nn.CrossEntropyLoss()\n\n    schedule = CycleScheduler(optimizer, lr, batch_size * epochs, momentum=(0.95, 0.85),)\n    data = load_images(batch_size = batch_size)\n    for i in itertools.count():\n        images = next(data).to(device)\n        bottom_enc = vae.encoders[0].encode(images)\n        _, _, bottom_idxs = vae.encoders[0].vq(bottom_enc)\n        _, _, top_idxs = vae.encoders[1](bottom_enc)\n        logits = bottom_prior(bottom_idxs, top_idxs)\n        logits = logits.permute(0, 2, 3, 1).contiguous()\n        logits = logits.view(-1, logits.shape[-1])\n        loss = loss_fn(logits, bottom_idxs.view(-1))\n        print('step %d: loss=%f' % (i, loss.item()))\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n        if not i % 30:\n            torch.save(bottom_prior.state_dict(), BOTTOM_PRIOR_PATH)\n        if i == epochs: break\n\nmain_bottom()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\"\"\"\nGenerate samples using the top-level prior.\n\"\"\"\n\nimport argparse\nimport random\n\nfrom PIL import Image\nimport numpy as np\nimport torch\n\nNUM_SAMPLES = 4\n\ndef main_sample():\n    device = torch.device(DEVICE)\n\n    vae = make_vae()\n    vae.load_state_dict(torch.load(VAE_PATH))\n    vae.to(device)\n    vae.eval()\n\n    top_prior = TopPrior()\n    top_prior.load_state_dict(torch.load(TOP_PRIOR_PATH))\n    top_prior.to(device)\n\n    results = np.zeros([NUM_SAMPLES, 8, 8], dtype=np.long)\n    for row in range(results.shape[1]):\n        for col in range(results.shape[2]):\n            partial_in = torch.from_numpy(results[:, :row + 1]).to(device)\n            with torch.no_grad():\n                outputs = torch.softmax(top_prior(partial_in), dim=1).cpu().numpy()\n            for i, out in enumerate(outputs):\n                probs = out[:, row, col]\n                results[i, row, col] = sample_softmax(probs)\n        print('done row', row)\n    with torch.no_grad():\n        full_latents = torch.from_numpy(results).to(device)\n        top_embedded = vae.encoders[1].vq.embed(full_latents)\n        bottom_encoded = vae.decoders[0]([top_embedded])\n        bottom_embedded, _, _ = vae.encoders[0].vq(bottom_encoded)\n        decoded = torch.clamp(vae.decoders[1]([top_embedded, bottom_embedded]), 0, 1)\n    decoded = decoded.permute(0, 2, 3, 1).cpu().numpy()\n    decoded = np.concatenate(decoded, axis=1)\n    Image.fromarray((decoded * 255).astype(np.uint8)).save('top_samples.png')\n\n\ndef sample_softmax(probs):\n    number = random.random()\n    for i, x in enumerate(probs):\n        number -= x\n        if number <= 0:\n            return i\n    return len(probs) - 1\n\n\ndef arg_parser():\n    parser = argparse.ArgumentParser()\n    parser.add_argument('--device', help='torch device', default='cuda')\n    return parser\n\nmain_sample()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"samples = Image.open('top_samples.png')\nsamples","execution_count":null,"outputs":[]}],"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}