{"cells":[{"metadata":{},"cell_type":"markdown","source":"#All credits goes to Andrej Karpathy\n#I **do not** own any of the code below, this is just a notebook for you to play around on Kaggle\n#github: https://github.com/karpathy/minGPT\n#Train GPT on images\n\nEffectively re-implements OpenAI's Image GPT model, getting GPT to model images instead of text, but using a near identical model. It's truly quite remarkable that a single model can agnostically do a great job modeling whatever data you give it: text, images, or whatever else. At the end of the day it is just a sequence of integers. Notice that unlike models like PixelCNN++ etc, this model knows nothing at all about the spatial layout of the pixels and has to learn the appropriate positional embeddings that reflect the spatial topology of the data.\n"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import numpy as np\nimport torchvision\nimport torch\nimport matplotlib.pyplot as plt\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# set up logging\nimport logging\nlogging.basicConfig(\n        format=\"%(asctime)s - %(levelname)s - %(name)s -   %(message)s\",\n        datefmt=\"%m/%d/%Y %H:%M:%S\",\n        level=logging.INFO,\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#utils\nimport random\nimport numpy as np\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\n\ndef set_seed(seed):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n\ndef top_k_logits(logits, k):\n    v, ix = torch.topk(logits, k)\n    out = logits.clone()\n    out[out < v[:, [-1]]] = -float('Inf')\n    return out\n\n@torch.no_grad()\ndef sample(model, x, steps, temperature=1.0, sample=False, top_k=None):\n    \"\"\"\n    take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in\n    the sequence, feeding the predictions back into the model each time. Clearly the sampling\n    has quadratic complexity unlike an RNN that is only linear, and has a finite context window\n    of block_size, unlike an RNN that has an infinite context window.\n    \"\"\"\n    block_size = model.get_block_size()\n    model.eval()\n    for k in range(steps):\n        x_cond = x if x.size(1) <= block_size else x[:, -block_size:] # crop context if needed\n        logits, _ = model(x_cond)\n        # pluck the logits at the final step and scale by temperature\n        logits = logits[:, -1, :] / temperature\n        # optionally crop probabilities to only the top k options\n        if top_k is not None:\n            logits = top_k_logits(logits, top_k)\n        # apply softmax to convert to probabilities\n        probs = F.softmax(logits, dim=-1)\n        # sample from the distribution or take the most likely\n        if sample:\n            ix = torch.multinomial(probs, num_samples=1)\n        else:\n            _, ix = torch.topk(probs, k=1, dim=-1)\n        # append to the sequence and continue\n        x = torch.cat((x, ix), dim=1)\n\n    return x\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#trainer\n\"\"\"\nSimple training loop; Boilerplate that could apply to any arbitrary neural network,\nso nothing in this file really has anything to do with GPT specifically.\n\"\"\"\n\nimport math\nimport logging\n\nfrom tqdm import tqdm\nimport numpy as np\n\nimport torch\nimport torch.optim as optim\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom torch.utils.data.dataloader import DataLoader\n\nlogger = logging.getLogger(__name__)\n\nclass TrainerConfig:\n    # optimization parameters\n    max_epochs = 10\n    batch_size = 64\n    learning_rate = 3e-4\n    betas = (0.9, 0.95)\n    grad_norm_clip = 1.0\n    weight_decay = 0.1 # only applied on matmul weights\n    # learning rate decay params: linear warmup followed by cosine decay to 10% of original\n    lr_decay = False\n    warmup_tokens = 375e6 # these two numbers come from the GPT-3 paper, but may not be good defaults elsewhere\n    final_tokens = 260e9 # (at what point we reach 10% of original LR)\n    # checkpoint settings\n    ckpt_path = None\n    num_workers = 0 # for DataLoader\n\n    def __init__(self, **kwargs):\n        for k,v in kwargs.items():\n            setattr(self, k, v)\n\nclass Trainer:\n\n    def __init__(self, model, train_dataset, test_dataset, config):\n        self.model = model\n        self.train_dataset = train_dataset\n        self.test_dataset = test_dataset\n        self.config = config\n\n        # take over whatever gpus are on the system\n        self.device = 'cpu'\n        if torch.cuda.is_available():\n            self.device = torch.cuda.current_device()\n            self.model = torch.nn.DataParallel(self.model).to(self.device)\n\n    def save_checkpoint(self):\n        # DataParallel wrappers keep raw model object in .module attribute\n        raw_model = self.model.module if hasattr(self.model, \"module\") else self.model\n        logger.info(\"saving %s\", self.config.ckpt_path)\n        torch.save(raw_model.state_dict(), self.config.ckpt_path)\n\n    def train(self):\n        model, config = self.model, self.config\n        raw_model = model.module if hasattr(self.model, \"module\") else model\n        optimizer = raw_model.configure_optimizers(config)\n\n        def run_epoch(split):\n            is_train = split == 'train'\n            model.train(is_train)\n            data = self.train_dataset if is_train else self.test_dataset\n            loader = DataLoader(data, shuffle=True, pin_memory=True,\n                                batch_size=config.batch_size,\n                                num_workers=config.num_workers)\n\n            losses = []\n            pbar = tqdm(enumerate(loader), total=len(loader)) if is_train else enumerate(loader)\n            for it, (x, y) in pbar:\n\n                # place data on the correct device\n                x = x.to(self.device)\n                y = y.to(self.device)\n\n                # forward the model\n                with torch.set_grad_enabled(is_train):\n                    logits, loss = model(x, y)\n                    loss = loss.mean() # collapse all losses if they are scattered on multiple gpus\n                    losses.append(loss.item())\n\n                if is_train:\n\n                    # backprop and update the parameters\n                    model.zero_grad()\n                    loss.backward()\n                    torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_norm_clip)\n                    optimizer.step()\n\n                    # decay the learning rate based on our progress\n                    if config.lr_decay:\n                        self.tokens += (y >= 0).sum() # number of tokens processed this step (i.e. label is not -100)\n                        if self.tokens < config.warmup_tokens:\n                            # linear warmup\n                            lr_mult = float(self.tokens) / float(max(1, config.warmup_tokens))\n                        else:\n                            # cosine learning rate decay\n                            progress = float(self.tokens - config.warmup_tokens) / float(max(1, config.final_tokens - config.warmup_tokens))\n                            lr_mult = max(0.1, 0.5 * (1.0 + math.cos(math.pi * progress)))\n                        lr = config.learning_rate * lr_mult\n                        for param_group in optimizer.param_groups:\n                            param_group['lr'] = lr\n                    else:\n                        lr = config.learning_rate\n\n                    # report progress\n                    pbar.set_description(f\"epoch {epoch+1} iter {it}: train loss {loss.item():.5f}. lr {lr:e}\")\n\n            if not is_train:\n                test_loss = float(np.mean(losses))\n                logger.info(\"test loss: %f\", test_loss)\n                return test_loss\n\n        best_loss = float('inf')\n        self.tokens = 0 # counter used for learning rate decay\n        for epoch in range(config.max_epochs):\n\n            run_epoch('train')\n            if self.test_dataset is not None:\n                test_loss = run_epoch('test')\n\n            # supports early stopping based on the test loss, or just save always if no test set is provided\n            good_model = self.test_dataset is None or test_loss < best_loss\n            if self.config.ckpt_path is not None and good_model:\n                best_loss = test_loss\n                self.save_checkpoint()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#model\n\"\"\"\nGPT model:\n- the initial stem consists of a combination of token encoding and a positional encoding\n- the meat of it is a uniform sequence of Transformer blocks\n    - each Transformer is a sequential combination of a 1-hidden-layer MLP block and a self-attention block\n    - all blocks feed into a central residual pathway similar to resnets\n- the final decoder is a linear projection into a vanilla Softmax classifier\n\"\"\"\n\nimport math\nimport logging\n\nimport torch\nimport torch.nn as nn\nfrom torch.nn import functional as F\n\nlogger = logging.getLogger(__name__)\n\nclass GPTConfig:\n    \"\"\" base GPT config, params common to all GPT versions \"\"\"\n    embd_pdrop = 0.1\n    resid_pdrop = 0.1\n    attn_pdrop = 0.1\n\n    def __init__(self, vocab_size, block_size, **kwargs):\n        self.vocab_size = vocab_size\n        self.block_size = block_size\n        for k,v in kwargs.items():\n            setattr(self, k, v)\n\nclass GPT1Config(GPTConfig):\n    \"\"\" GPT-1 like network roughly 125M params \"\"\"\n    n_layer = 12\n    n_head = 12\n    n_embd = 768\n\nclass CausalSelfAttention(nn.Module):\n    \"\"\"\n    A vanilla multi-head masked self-attention layer with a projection at the end.\n    It is possible to use torch.nn.MultiheadAttention here but I am including an\n    explicit implementation here to show that there is nothing too scary here.\n    \"\"\"\n\n    def __init__(self, config):\n        super().__init__()\n        assert config.n_embd % config.n_head == 0\n        # key, query, value projections for all heads\n        self.key = nn.Linear(config.n_embd, config.n_embd)\n        self.query = nn.Linear(config.n_embd, config.n_embd)\n        self.value = nn.Linear(config.n_embd, config.n_embd)\n        # regularization\n        self.attn_drop = nn.Dropout(config.attn_pdrop)\n        self.resid_drop = nn.Dropout(config.resid_pdrop)\n        # output projection\n        self.proj = nn.Linear(config.n_embd, config.n_embd)\n        # causal mask to ensure that attention is only applied to the left in the input sequence\n        self.register_buffer(\"mask\", torch.tril(torch.ones(config.block_size, config.block_size))\n                                     .view(1, 1, config.block_size, config.block_size))\n        self.n_head = config.n_head\n\n    def forward(self, x, layer_past=None):\n        B, T, C = x.size()\n\n        # calculate query, key, values for all heads in batch and move head forward to be the batch dim\n        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)\n\n        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)\n        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))\n        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))\n        att = F.softmax(att, dim=-1)\n        att = self.attn_drop(att)\n        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)\n        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side\n\n        # output projection\n        y = self.resid_drop(self.proj(y))\n        return y\n\nclass Block(nn.Module):\n    \"\"\" an unassuming Transformer block \"\"\"\n\n    def __init__(self, config):\n        super().__init__()\n        self.ln1 = nn.LayerNorm(config.n_embd)\n        self.ln2 = nn.LayerNorm(config.n_embd)\n        self.attn = CausalSelfAttention(config)\n        self.mlp = nn.Sequential(\n            nn.Linear(config.n_embd, 4 * config.n_embd),\n            nn.GELU(),\n            nn.Linear(4 * config.n_embd, config.n_embd),\n            nn.Dropout(config.resid_pdrop),\n        )\n\n    def forward(self, x):\n        x = x + self.attn(self.ln1(x))\n        x = x + self.mlp(self.ln2(x))\n        return x\n\nclass GPT(nn.Module):\n    \"\"\"  the full GPT language model, with a context size of block_size \"\"\"\n\n    def __init__(self, config):\n        super().__init__()\n\n        # input embedding stem\n        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)\n        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))\n        self.drop = nn.Dropout(config.embd_pdrop)\n        # transformer\n        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])\n        # decoder head\n        self.ln_f = nn.LayerNorm(config.n_embd)\n        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)\n\n        self.block_size = config.block_size\n        self.apply(self._init_weights)\n\n        logger.info(\"number of parameters: %e\", sum(p.numel() for p in self.parameters()))\n\n    def get_block_size(self):\n        return self.block_size\n\n    def _init_weights(self, module):\n        if isinstance(module, (nn.Linear, nn.Embedding)):\n            module.weight.data.normal_(mean=0.0, std=0.02)\n            if isinstance(module, nn.Linear) and module.bias is not None:\n                module.bias.data.zero_()\n        elif isinstance(module, nn.LayerNorm):\n            module.bias.data.zero_()\n            module.weight.data.fill_(1.0)\n\n    def configure_optimizers(self, train_config):\n        \"\"\"\n        This long function is unfortunately doing something very simple and is being very defensive:\n        We are separating out all parameters of the model into two buckets: those that will experience\n        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).\n        We are then returning the PyTorch optimizer object.\n        \"\"\"\n\n        # separate out all parameters to those that will and won't experience regularizing weight decay\n        decay = set()\n        no_decay = set()\n        whitelist_weight_modules = (torch.nn.Linear, )\n        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n        for mn, m in self.named_modules():\n            for pn, p in m.named_parameters():\n                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name\n\n                if pn.endswith('bias'):\n                    # all biases will not be decayed\n                    no_decay.add(fpn)\n                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):\n                    # weights of whitelist modules will be weight decayed\n                    decay.add(fpn)\n                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):\n                    # weights of blacklist modules will NOT be weight decayed\n                    no_decay.add(fpn)\n\n        # special case the position embedding parameter in the root GPT module as not decayed\n        no_decay.add('pos_emb')\n\n        # validate that we considered every parameter\n        param_dict = {pn: p for pn, p in self.named_parameters()}\n        inter_params = decay & no_decay\n        union_params = decay | no_decay\n        assert len(inter_params) == 0, \"parameters %s made it into both decay/no_decay sets!\" % (str(inter_params), )\n        assert len(param_dict.keys() - union_params) == 0, \"parameters %s were not separated into either decay/no_decay set!\" \\\n                                                    % (str(param_dict.keys() - union_params), )\n\n        # create the pytorch optimizer object\n        optim_groups = [\n            {\"params\": [param_dict[pn] for pn in sorted(list(decay))], \"weight_decay\": train_config.weight_decay},\n            {\"params\": [param_dict[pn] for pn in sorted(list(no_decay))], \"weight_decay\": 0.0},\n        ]\n        optimizer = torch.optim.AdamW(optim_groups, lr=train_config.learning_rate, betas=train_config.betas)\n        return optimizer\n\n    def forward(self, idx, targets=None):\n        b, t = idx.size()\n        assert t <= self.block_size, \"Cannot forward, model block size is exhausted.\"\n\n        # forward the GPT model\n        token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector\n        position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector\n        x = self.drop(token_embeddings + position_embeddings)\n        x = self.blocks(x)\n        x = self.ln_f(x)\n        logits = self.head(x)\n\n        # if we are given some desired targets also calculate the loss\n        loss = None\n        if targets is not None:\n            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))\n\n        return logits, loss\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# make deterministic\nset_seed(42)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# pytorch helpfully makes it easy to download datasets, e.g. the common CIFAR-10 https://www.kaggle.com/c/cifar-10\nroot = './'\ntrain_data = torchvision.datasets.CIFAR10(root, train=True, transform=None, target_transform=None, download=True)\ntest_data  = torchvision.datasets.CIFAR10(root, train=False, transform=None, target_transform=None, download=True)\nprint(len(train_data), len(test_data))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# get random 5 pixels per image and stack them all up as rgb values to get half a million random pixels\npluck_rgb = lambda x: torch.from_numpy(np.array(x)).view(32*32, 3)[torch.randperm(32*32)[:5], :]\npx = torch.cat([pluck_rgb(x) for x, y in train_data], dim=0).float()\nprint(px.size())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# run kmeans to get our codebook\n\ndef kmeans(x, ncluster, niter=10):\n    N, D = x.size()\n    c = x[torch.randperm(N)[:ncluster]] # init clusters at random\n    for i in range(niter):\n        # assign all pixels to the closest codebook element\n        a = ((x[:, None, :] - c[None, :, :])**2).sum(-1).argmin(1)\n        # move each codebook element to be the mean of the pixels that assigned to it\n        c = torch.stack([x[a==k].mean(0) for k in range(ncluster)])\n        # re-assign any poorly positioned codebook elements\n        nanix = torch.any(torch.isnan(c), dim=1)\n        ndead = nanix.sum().item()\n        print('done step %d/%d, re-initialized %d dead clusters' % (i+1, niter, ndead))\n        c[nanix] = x[torch.randperm(N)[:ndead]] # re-init dead clusters\n    return c\n\nncluster = 512\nwith torch.no_grad():\n    C = kmeans(px, ncluster, niter=8)\n\nprint(C.size())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# encode the training examples with our codebook to visualize how much we've lost in the discretization\nn_samples = 16\nncol = 8\nnrow = n_samples // ncol + 1\nplt.figure(figsize=(20, 10))\nfor i in range(n_samples):\n    \n    # encode and decode random data\n    x, y = train_data[np.random.randint(0, len(train_data))]\n    xpt = torch.from_numpy(np.array(x)).float().view(32*32, 3)\n    ix = ((xpt[:, None, :] - C[None, :, :])**2).sum(-1).argmin(1) # cluster assignments for each pixel\n    \n    # these images should look normal ideally\n    plt.subplot(nrow, ncol, i+1)\n    plt.imshow(C[ix].view(32, 32, 3).numpy().astype(np.uint8))\n    plt.axis('off')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The images above look relatively reasonable, so our 512-sized codebook is enough to reasonably re-represent RGB values. Ok cool. So now every image is just a 1024-long sequence of numbers between 0..511. Time to train a GPT."},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import Dataset\n\nclass ImageDataset(Dataset):\n    \"\"\"\n    wrap up the pytorch CIFAR-10 dataset into our own, which will convert images into sequences of integers\n    \"\"\"\n    \n    def __init__(self, pt_dataset, clusters, perm=None):\n        self.pt_dataset = pt_dataset\n        self.clusters = clusters\n        self.perm = torch.arange(32*32) if perm is None else perm\n        \n        self.vocab_size = clusters.size(0)\n        self.block_size = 32*32 - 1\n        \n    def __len__(self):\n        return len(self.pt_dataset)\n\n    def __getitem__(self, idx):\n        x, y = self.pt_dataset[idx]\n        x = torch.from_numpy(np.array(x)).view(-1, 3) # flatten out all pixels\n        x = x[self.perm].float() # reshuffle pixels with any fixed permutation and -> float\n        a = ((x[:, None, :] - self.clusters[None, :, :])**2).sum(-1).argmin(1) # cluster assignments\n        return a[:-1], a[1:] # always just predict the next one in the sequence\n\ntrain_dataset = ImageDataset(train_data, C)\ntest_dataset = ImageDataset(test_data, C)\ntrain_dataset[0][0] # one example image flattened out into integers","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\nFor reference, **iGPT-S** from the paper is:\n- batch size of 128 and trained for 1M terations\n- Adam lr 0.003 with betas = (0.9, 0.95)\n- learning rate is warmed up for one epoch, then decays to 0\n- did not use weight decay or dropout\n- `n_layer=24, n_head=8, n_embd=512`\n\nWe will do something similar but smaller"},{"metadata":{"trusted":true},"cell_type":"code","source":"\n# we'll do something a bit smaller\nmconf = GPTConfig(train_dataset.vocab_size, train_dataset.block_size,\n                  embd_pdrop=0.0, resid_pdrop=0.0, attn_pdrop=0.0,\n                  n_layer=12, n_head=8, n_embd=256)\nmodel = GPT(mconf)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n\"\"\"\nNote that Karpathy running on an 8-GPU V100 machine so each GPU has 32GB.\nHe use a batch of 16*8 in the original implementation. I brought this down to 8\nso it could fit into Kaggle TPU.\nIf you don't have as many computational resources you have to bring down\nthe batch_size until the model fits into your memory, and then you may\nalso need to adjust the learning rate (e.g. decrease it a bit). Alternatively,\nyou can use an even smaller model up above, bringing down the number of layers,\nnumber of heads, and the embedding size.\n\"\"\"\n\ntokens_per_epoch = len(train_data) * train_dataset.block_size\ntrain_epochs = 20 # todo run a bigger model and longer, this is tiny\n\n# initialize a trainer instance and kick off training\ntconf = TrainerConfig(max_epochs=train_epochs, batch_size=8, learning_rate=3e-3,\n                      betas = (0.9, 0.95), weight_decay=0,\n                      lr_decay=True, warmup_tokens=tokens_per_epoch, final_tokens=train_epochs*tokens_per_epoch,\n                      ckpt_path='cifar10_model.pt',\n                      num_workers=8)\ntrainer = Trainer(model, train_dataset, test_dataset, tconf)\ntrainer.train()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# load the state of the best model we've seen based on early stopping\ncheckpoint = torch.load('cifar10_model.pt')\nmodel.load_state_dict(checkpoint)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# to sample we also have to technically \"train\" a separate model for the first token in the sequence\n# we are going to do so below simply by calculating and normalizing the histogram of the first token\ncounts = torch.ones(ncluster) # start counts as 1 not zero, this is called \"smoothing\"\nrp = torch.randperm(len(train_dataset))\nnest = 5000 # how many images to use for the estimation\nfor i in range(nest):\n    a, _ = train_dataset[int(rp[i])]\n    t = a[0].item() # index of first token in the sequence\n    counts[t] += 1\nprob = counts/counts.sum()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nfrom mingpt.utils import sample\n\nn_samples = 32\nstart_pixel = np.random.choice(np.arange(C.size(0)), size=(n_samples, 1), replace=True, p=prob)\nstart_pixel = torch.from_numpy(start_pixel).to(trainer.device)\npixels = sample(model, start_pixel, 32*32-1, temperature=1.0, sample=True, top_k=100)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# for visualization we have to invert the permutation used to produce the pixels\niperm = torch.argsort(train_dataset.perm)\n\nncol = 8\nnrow = n_samples // ncol\nplt.figure(figsize=(16, 8))\nfor i in range(n_samples):\n    pxi = pixels[i][iperm] # note: undo the encoding permutation\n    \n    plt.subplot(nrow, ncol, i+1)\n    plt.imshow(C[pxi].view(32, 32, 3).numpy().astype(np.uint8))\n    plt.axis('off')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# visualize some of the learned positional embeddings, maybe they contain structure\nplt.figure(figsize=(5, 5))\nnsee = 8*8\nncol = 8\nnrow = nsee // ncol\nfor i in range(nsee):\n    \n    ci = model.pos_emb.data[0, :, i].cpu()\n    zci = torch.cat((torch.tensor([0.0]), ci)) # pre-cat a zero\n    rzci = zci[iperm] # undo the permutation to recover the pixel space of the image\n    \n    plt.subplot(nrow, ncol, i+1)\n    plt.imshow(rzci.view(32, 32).numpy())\n    plt.axis('off')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}