{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list the files in the input directory\n\nimport os\nprint(os.listdir(\"../input\"))\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true,"_kg_hide-input":false,"_kg_hide-output":false},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nimport torch.optim as optim\nfrom tqdm import tqdm_notebook\nimport torchvision.datasets as datasets\nimport torchvision.transforms as transforms\nfrom torchvision.utils import make_grid\nfrom scipy.signal import savgol_filter\nfrom math import cos, pi, floor, sin\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## VQ VAE Implementation"},{"metadata":{},"cell_type":"markdown","source":"code mainly from https://github.com/zalandoresearch/pytorch-vq-vae/blob/master/vq-vae.ipynb"},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.utils.data import DataLoader\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\nComputeLB = False\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 \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 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 = np.min((xmax - xmin, ymax - ymin))\n img2 = img.crop((xmin, ymin, xmin+w, ymin+w))\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 IMAGES = np.sort(IMAGES)\n np.random.seed(810)\n x = np.random.choice(np.arange(20579),10000)\n np.random.seed(None)\n for k in range(len(x)):\n img = Image.open(ROOT + 'all-dogs/all-dogs/' + IMAGES[x[k]])\n w = img.size[0]; h = img.size[1];\n if (k%2==0)|(k%3==0):\n w2 = 100; h2 = int(h/(w/100))\n a = 18; b = 0 \n else:\n a=0; b=0\n if w 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, encoding_indices","execution_count":null,"outputs":[]},{"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":"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, encoding_indices","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//2,\n kernel_size=4,\n stride=2, padding=1)\n self._conv_2 = 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_3(x)\n return self._residual_stack(x)","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=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 return self._conv_trans_2(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_size = 32\nnum_training_updates = 25000\n\nnum_hiddens = 128\nnum_residual_hiddens = 32\nnum_residual_layers = 2\n\nembedding_dim = 64\nnum_embeddings = 8\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":"np.moveaxis(imagesIn,3,-3).shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"imagesIn = imagesIn.astype(\"float32\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### shuffle data"},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.randperm(len(imagesIn))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"imagesIn = imagesIn[torch.randperm(len(imagesIn))]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(imagesIn)-len(imagesIn)//10","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = imagesIn[:len(imagesIn)-len(imagesIn)//10]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val = imagesIn[len(imagesIn)-len(imagesIn)//10:]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"training_loader = DataLoader(((np.moveaxis(train,3,-3)/255)), \n batch_size=batch_size, \n shuffle=True,\n pin_memory=True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_loader = DataLoader((((np.moveaxis(val,3,-3)/255))), \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 loss, quantized, perplexity, emb_codes = self.encode(x)\n x_recon = self._decoder(quantized)\n\n return loss, x_recon, perplexity\n \n def encode(self,x):\n z = self._encoder(x)\n z = self._pre_vq_conv(z)\n loss, quantized, perplexity, emb_codes = self._vq_vae(z)\n return loss, quantized, perplexity, emb_codes\n \n ","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":"optimizer = optim.Adam(model.parameters(), lr=learning_rate, amsgrad=False)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data= next(iter(training_loader))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data.std()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data.mean()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.train()\ntrain_res_recon_error = []\ntrain_res_perplexity = []\n\nfor i in range(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) % 100 == 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":{"trusted":true},"cell_type":"code","source":"from scipy.signal import savgol_filter\n","execution_count":null,"outputs":[]},{"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":{"trusted":true},"cell_type":"code","source":"model.eval()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_originals= next(iter(val_loader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"valid_originals = valid_originals.to(device)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"vq_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":"valid_quantize.shape","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":{},"cell_type":"markdown","source":"> ## Creating reconstractions from unseed examples"},{"metadata":{"trusted":true},"cell_type":"code","source":"show(make_grid(valid_reconstructions.cpu().data))\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show(make_grid(valid_originals.cpu()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(),\"./vqvae.pth\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## extract codes from all samples"},{"metadata":{"trusted":true},"cell_type":"code","source":"codes = []\nval_codes =[]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"iterator = iter(training_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_iterator = iter(val_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with torch.no_grad():\n for batch in iterator:\n _, quantize, _, emb_code = model.encode(batch.cuda())\n codes.append(emb_code)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"with torch.no_grad():\n for batch in val_iterator:\n _, quantize, _, emb_code = model.encode(batch.cuda())\n val_codes.append(emb_code)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"codes = torch.cat(codes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_codes = torch.cat(val_codes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"codes = codes.view(-1,16,16)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_codes = val_codes.view(-1,16,16)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## pixelsnail code "},{"metadata":{},"cell_type":"markdown","source":"code from https://github.com/rosinality/vq-vae-2-pytorch/blob/master/pixelsnail.py"},{"metadata":{"trusted":true},"cell_type":"code","source":"# Copyright (c) Xi Chen\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\n# Borrowed from https://github.com/neocxi/pixelsnail-public and ported it to PyTorch\n\nfrom math import sqrt\nfrom functools import partial, lru_cache\n\nimport numpy as np\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\ndef wn_linear(in_dim, out_dim):\n return nn.utils.weight_norm(nn.Linear(in_dim, out_dim))\n\n\nclass WNConv2d(nn.Module):\n def __init__(\n self,\n in_channel,\n out_channel,\n kernel_size,\n stride=1,\n padding=0,\n bias=True,\n activation=None,\n ):\n super().__init__()\n\n self.conv = nn.utils.weight_norm(\n nn.Conv2d(\n in_channel,\n out_channel,\n kernel_size,\n stride=stride,\n padding=padding,\n bias=bias,\n )\n )\n\n self.out_channel = out_channel\n\n if isinstance(kernel_size, int):\n kernel_size = [kernel_size, kernel_size]\n\n self.kernel_size = kernel_size\n\n self.activation = activation\n\n def forward(self, input):\n out = self.conv(input)\n\n if self.activation is not None:\n out = self.activation(out)\n\n return out\n\n\ndef shift_down(input, size=1):\n return F.pad(input, [0, 0, size, 0])[:, :, : input.shape[2], :]\n\n\ndef shift_right(input, size=1):\n return F.pad(input, [size, 0, 0, 0])[:, :, :, : input.shape[3]]\n\n\nclass CausalConv2d(nn.Module):\n def __init__(\n self,\n in_channel,\n out_channel,\n kernel_size,\n stride=1,\n padding='downright',\n activation=None,\n ):\n super().__init__()\n\n if isinstance(kernel_size, int):\n kernel_size = [kernel_size] * 2\n\n self.kernel_size = kernel_size\n\n if padding == 'downright':\n pad = [kernel_size[1] - 1, 0, kernel_size[0] - 1, 0]\n\n elif padding == 'down' or padding == 'causal':\n pad = kernel_size[1] // 2\n\n pad = [pad, pad, kernel_size[0] - 1, 0]\n\n self.causal = 0\n if padding == 'causal':\n self.causal = kernel_size[1] // 2\n\n self.pad = nn.ZeroPad2d(pad)\n\n self.conv = WNConv2d(\n in_channel,\n out_channel,\n kernel_size,\n stride=stride,\n padding=0,\n activation=activation,\n )\n\n def forward(self, input):\n out = self.pad(input)\n\n if self.causal > 0:\n self.conv.conv.weight_v.data[:, :, -1, self.causal :].zero_()\n\n out = self.conv(out)\n\n return out\n\n\nclass GatedResBlock(nn.Module):\n def __init__(\n self,\n in_channel,\n channel,\n kernel_size,\n conv='wnconv2d',\n activation=nn.ELU,\n dropout=0.1,\n auxiliary_channel=0,\n condition_dim=0,\n ):\n super().__init__()\n\n if conv == 'wnconv2d':\n conv_module = partial(WNConv2d, padding=kernel_size // 2)\n\n elif conv == 'causal_downright':\n conv_module = partial(CausalConv2d, padding='downright')\n\n elif conv == 'causal':\n conv_module = partial(CausalConv2d, padding='causal')\n\n self.activation = activation(inplace=True)\n self.conv1 = conv_module(in_channel, channel, kernel_size)\n\n if auxiliary_channel > 0:\n self.aux_conv = WNConv2d(auxiliary_channel, channel, 1)\n\n self.dropout = nn.Dropout(dropout)\n\n self.conv2 = conv_module(channel, in_channel * 2, kernel_size)\n\n if condition_dim > 0:\n # self.condition = nn.Linear(condition_dim, in_channel * 2, bias=False)\n self.condition = WNConv2d(condition_dim, in_channel * 2, 1, bias=False)\n\n self.gate = nn.GLU(1)\n\n def forward(self, input, aux_input=None, condition=None):\n out = self.conv1(self.activation(input))\n\n if aux_input is not None:\n out = out + self.aux_conv(self.activation(aux_input))\n\n out = self.activation(out)\n out = self.dropout(out)\n out = self.conv2(out)\n\n if condition is not None:\n condition = self.condition(condition)\n out += condition\n # out = out + condition.view(condition.shape[0], 1, 1, condition.shape[1])\n\n out = self.gate(out)\n out += input\n\n return out\n\n\n@lru_cache(maxsize=64)\ndef causal_mask(size):\n shape = [size, size]\n mask = np.triu(np.ones(shape), k=1).astype(np.uint8).T\n start_mask = np.ones(size).astype(np.float32)\n start_mask[0] = 0\n\n return (\n torch.from_numpy(mask).unsqueeze(0),\n torch.from_numpy(start_mask).unsqueeze(1),\n )\n\n\nclass CausalAttention(nn.Module):\n def __init__(self, query_channel, key_channel, channel, n_head=8, dropout=0.1):\n super().__init__()\n\n self.query = wn_linear(query_channel, channel)\n self.key = wn_linear(key_channel, channel)\n self.value = wn_linear(key_channel, channel)\n\n self.dim_head = channel // n_head\n self.n_head = n_head\n\n self.dropout = nn.Dropout(dropout)\n\n def forward(self, query, key):\n batch, _, height, width = key.shape\n\n def reshape(input):\n return input.view(batch, -1, self.n_head, self.dim_head).transpose(1, 2)\n\n query_flat = query.view(batch, query.shape[1], -1).transpose(1, 2)\n key_flat = key.view(batch, key.shape[1], -1).transpose(1, 2)\n query = reshape(self.query(query_flat))\n key = reshape(self.key(key_flat)).transpose(2, 3)\n value = reshape(self.value(key_flat))\n\n attn = torch.matmul(query, key) / sqrt(self.dim_head)\n mask, start_mask = causal_mask(height * width)\n mask = mask.type_as(query)\n start_mask = start_mask.type_as(query)\n attn = attn.masked_fill(mask == 0, -1e4)\n attn = torch.softmax(attn, 3) * start_mask\n attn = self.dropout(attn)\n\n out = attn @ value\n out = out.transpose(1, 2).reshape(\n batch, height, width, self.dim_head * self.n_head\n )\n out = out.permute(0, 3, 1, 2)\n\n return out\n\n\nclass PixelBlock(nn.Module):\n def __init__(\n self,\n in_channel,\n channel,\n kernel_size,\n n_res_block,\n attention=True,\n dropout=0.1,\n condition_dim=0,\n ):\n super().__init__()\n\n resblocks = []\n for i in range(n_res_block):\n resblocks.append(\n GatedResBlock(\n in_channel,\n channel,\n kernel_size,\n conv='causal',\n dropout=dropout,\n condition_dim=condition_dim,\n )\n )\n\n self.resblocks = nn.ModuleList(resblocks)\n\n self.attention = attention\n\n if attention:\n self.key_resblock = GatedResBlock(\n in_channel * 2 + 2, in_channel, 1, dropout=dropout\n )\n self.query_resblock = GatedResBlock(\n in_channel + 2, in_channel, 1, dropout=dropout\n )\n\n self.causal_attention = CausalAttention(\n in_channel + 2, in_channel * 2 + 2, in_channel // 2, dropout=dropout\n )\n\n self.out_resblock = GatedResBlock(\n in_channel,\n in_channel,\n 1,\n auxiliary_channel=in_channel // 2,\n dropout=dropout,\n )\n\n else:\n self.out = WNConv2d(in_channel + 2, in_channel, 1)\n\n def forward(self, input, background, condition=None):\n out = input\n\n for resblock in self.resblocks:\n out = resblock(out, condition=condition)\n\n if self.attention:\n key_cat = torch.cat([input, out, background], 1)\n key = self.key_resblock(key_cat)\n query_cat = torch.cat([out, background], 1)\n query = self.query_resblock(query_cat)\n attn_out = self.causal_attention(query, key)\n out = self.out_resblock(out, attn_out)\n\n else:\n bg_cat = torch.cat([out, background], 1)\n out = self.out(bg_cat)\n\n return out\n\n\nclass CondResNet(nn.Module):\n def __init__(self, in_channel, channel, kernel_size, n_res_block):\n super().__init__()\n\n blocks = [WNConv2d(in_channel, channel, kernel_size, padding=kernel_size // 2)]\n\n for i in range(n_res_block):\n blocks.append(GatedResBlock(channel, channel, kernel_size))\n\n self.blocks = nn.Sequential(*blocks)\n\n def forward(self, input):\n return self.blocks(input)\n\n\nclass PixelSNAIL(nn.Module):\n def __init__(\n self,\n shape,\n n_class,\n channel,\n kernel_size,\n n_block,\n n_res_block,\n res_channel,\n attention=True,\n dropout=0.1,\n n_cond_res_block=0,\n cond_res_channel=0,\n cond_res_kernel=3,\n n_out_res_block=0,\n ):\n super().__init__()\n\n height, width = shape\n\n self.n_class = n_class\n\n if kernel_size % 2 == 0:\n kernel = kernel_size + 1\n\n else:\n kernel = kernel_size\n\n self.horizontal = CausalConv2d(\n n_class, channel, [kernel // 2, kernel], padding='down'\n )\n self.vertical = CausalConv2d(\n n_class, channel, [(kernel + 1) // 2, kernel // 2], padding='downright'\n )\n\n coord_x = (torch.arange(height).float() - height / 2) / height\n coord_x = coord_x.view(1, 1, height, 1).expand(1, 1, height, width)\n coord_y = (torch.arange(width).float() - width / 2) / width\n coord_y = coord_y.view(1, 1, 1, width).expand(1, 1, height, width)\n self.register_buffer('background', torch.cat([coord_x, coord_y], 1))\n\n self.blocks = nn.ModuleList()\n\n for i in range(n_block):\n self.blocks.append(\n PixelBlock(\n channel,\n res_channel,\n kernel_size,\n n_res_block,\n attention=attention,\n dropout=dropout,\n condition_dim=cond_res_channel,\n )\n )\n\n if n_cond_res_block > 0:\n self.cond_resnet = CondResNet(\n n_class, cond_res_channel, cond_res_kernel, n_cond_res_block\n )\n\n out = []\n\n for i in range(n_out_res_block):\n out.append(GatedResBlock(channel, res_channel, 1))\n\n out.extend([nn.ELU(inplace=True), WNConv2d(channel, n_class, 1)])\n\n self.out = nn.Sequential(*out)\n\n def forward(self, input, condition=None, cache=None):\n if cache is None:\n cache = {}\n batch, height, width = input.shape\n input = (\n F.one_hot(input, self.n_class).permute(0, 3, 1, 2).type_as(self.background)\n )\n horizontal = shift_down(self.horizontal(input))\n vertical = shift_right(self.vertical(input))\n out = horizontal + vertical\n\n background = self.background[:, :, :height, :].expand(batch, 2, height, width)\n\n if condition is not None:\n if 'condition' in cache:\n condition = cache['condition']\n condition = condition[:, :, :height, :]\n\n else:\n condition = (\n F.one_hot(condition, self.n_class)\n .permute(0, 3, 1, 2)\n .type_as(self.background)\n )\n condition = self.cond_resnet(condition)\n condition = F.interpolate(condition, scale_factor=2)\n cache['condition'] = condition.detach().clone()\n condition = condition[:, :, :height, :]\n\n for block in self.blocks:\n out = block(out, background, condition=condition)\n\n out = self.out(out)\n\n return out, cache","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# pxModel = PixelSNAIL(\n# list(codes.shape[1:3]),\n# 512,\n# 128,\n# 5,\n# 4,\n# 4,\n# 256,\n# dropout=0.1,\n# n_out_res_block=0)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"PixelSNAIL??","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pxModel = PixelSNAIL(\n list(codes.shape[1:3]),\n 8,\n 128,\n 5,\n 4,\n 4,\n 256,\n dropout=0.1,\n n_out_res_block=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"pxModel = pxModel.to(device)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"optimizer = optim.Adam(pxModel.parameters(), lr=1e-3)\nscheduler = None","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loader = DataLoader(codes,128,True)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"val_loader = DataLoader(val_codes,32,False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torch.optim import lr_scheduler\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train( epoch, loader, model, optimizer, scheduler, device):\n model.train()\n loader = tqdm_notebook(loader)\n criterion = nn.CrossEntropyLoss()\n losses = []\n for i, batch in enumerate(loader):\n model.zero_grad()\n\n batch = batch.to(device)\n\n target = batch\n out, _ = model(batch)\n\n\n\n loss = criterion(out, target)\n loss.backward()\n\n if scheduler is not None:\n scheduler.step()\n optimizer.step()\n\n _, pred = out.max(1)\n correct = (pred == target).float()\n accuracy = correct.sum() / target.numel()\n losses.append(loss.item())\n lr = optimizer.param_groups[0]['lr']\n loader.set_description(\n (\n f'epoch: {epoch + 1}; loss: {loss.item():.5f}; '\n f'acc: {accuracy:.5f}; lr: {lr:.5f}'\n )\n )\n \n epoch_loss = np.array(losses).mean()\n loader.set_description(\n (\n f'epoch: {epoch + 1}; loss: {epoch_loss:.5f}; '\n f'acc: {accuracy:.5f}; lr: {lr:.5f}'\n )\n )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def validate(epoch,loader, pxModel, device):\n loader = tqdm_notebook(loader)\n losses = 0\n currect = 0\n total = 0\n batches = len(loader)\n \n with torch.no_grad():\n pxModel.eval()\n criterion = nn.CrossEntropyLoss()\n for batch in loader:\n total+=batch.numel()\n val_out, _ = pxModel(batch)\n val_loss = criterion(val_out, batch)\n _, val_pred = val_out.max(1)\n batch_currect = (val_pred == batch).float().sum()\n currect += batch_currect\n losses += val_loss\n batch_accuracy = batch_currect/batch.numel()\n loader.set_description(\n (\n f'epoch: {epoch + 1}; loss: {val_loss.item():.5f}; '\n f'acc: {batch_accuracy:.5f};'\n )\n )\n print(\"val_accuracy: {}\".format(currect/total))\n print(\"val_loss: {}\".format(losses/batches))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"epoch = 260","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def anneal_linear(start, end, proportion):\n return start + proportion * (end - start)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def anneal_cos(start, end, proportion):\n cos_val = cos(pi * proportion) + 1\n\n return end + (start - end) / 2 * cos_val","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n@torch.no_grad()\ndef sample_model(model, device, batch, size, temperature, condition=None):\n row = torch.zeros(batch, *size, dtype=torch.int32).to(device)\n cache = {}\n\n for i in tqdm_notebook(range(size[0])):\n for j in range(size[1]):\n out, cache = model(row[:, : i + 1, :].long(), condition=condition, cache=cache)\n prob = torch.softmax(out[:, :, i, j] / temperature, 1)\n sample = torch.multinomial(prob, 1).squeeze(-1)\n row[:, i, j] = sample\n\n return row","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def sample_model_full(pxModel,vqmodel,device,batch,size,temperature,condition=None):\n pxModel.eval()\n vqmodel.eval()\n sample_codes = sample_model(pxModel, device, batch, size, temperature)\n quantized = vqmodel._vq_vae._embedding(sample_codes.long())\n quantized = quantized.permute(0,3,1,2)\n samples = vqmodel._decoder(quantized)\n return samples","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef show_samples(pxModel,vqmodel,device,batch,size,temperature,condition=None):\n samples = sample_model_full(pxModel,vqmodel,device,batch,size,temperature,condition)\n show(make_grid(samples.detach().cpu()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef get_samples(pxmodel,vqmodel,device):\n samples = []\n for i in range(50):\n sample_batch = sample_model_full(pxmodel,vqmodel, device, 200, [16,16], 0.7)\n samples.append(sample_batch)\n return torch.cat(samples)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"sched = CycleScheduler(optimizer,1e-2,epoch*len(loader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.cuda.empty_cache()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## recontraction form codes generated by untrained pixelSanil "},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(epoch//5):\n train(i,loader, pxModel, optimizer, sched, device)\n if i%5==0:\n print(\"saving model\")\n torch.save(pxModel.state_dict(),\"./pxSnail.pth\")\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(epoch//5,2*(epoch//5)):\n train(i,loader, pxModel, optimizer, sched, device)\n if i%5==0:\n print(\"saving model\")\n torch.save(pxModel.state_dict(),\"./pxSnail.pth\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(2*(epoch//5),3*(epoch//5)):\n train(i,loader, pxModel, optimizer, sched, device)\n if i%5==0:\n print(\"saving model\")\n torch.save(pxModel.state_dict(),\"./pxSnail.pth\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(3*(epoch//5),4*(epoch//5),):\n train(i,loader, pxModel, optimizer, sched, device)\n if i%5==0:\n print(\"saving model\")\n torch.save(pxModel.state_dict(),\"./pxSnail.pth\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for i in range(4*(epoch//5),epoch):\n train(i,loader, pxModel, optimizer, sched, device)\n if i%5==0:\n print(\"saving model\")\n torch.save(pxModel.state_dict(),\"./pxSnail.pth\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## reconstractions after pixleSnail training"},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"show_samples(pxModel,model,device,16,[16,16],1.0)\n","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":{"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}