{"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\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 read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\n# import os\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n#     for filename in filenames:\n#         print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from __future__ import print_function\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom scipy.signal import savgol_filter\nimport umap\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import DataLoader\nimport torch.optim as optim\nimport torchvision.datasets as datasets\nimport torchvision.transforms as transforms\nfrom torchvision.utils import make_grid\nfrom six.moves import xrange\nimport torchvision\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.manual_seed(42)\ntorch.cuda.manual_seed(42)\n \n## Finds images\ntrain_data_path = '../input/gan-getting-started/'\n \n### Rescaling incoming image to 28 by 28 pixels\n### After Rescaling, convert the image to a tensor\ntransform = transforms.Compose([transforms.Resize((32,32)),transforms.ToTensor(),transforms.Normalize((0.5,0.5,0.5), (1.0,1.0,1.0))])\ntrain_data = torchvision.datasets.ImageFolder(root=train_data_path,transform=transform)\ntest_data = torchvision.datasets.ImageFolder(root=train_data_path,transform=transform)\nbatch_size = 256\ntraining_loader = torch.utils.data.DataLoader(train_data,batch_size,shuffle=True)\nvalidation_loader = torch.utils.data.DataLoader(train_data,32,shuffle=True)\n##### Declare the model architecture","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#data_variance = np.var(train_loader.dataset/255.0)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class VectorQuantizer(nn.Module):\n    def __init__(self, num_embeddings, embedding_dim, commitment_cost):\n        super(VectorQuantizer, self).__init__()\n        \n        ### Create an embedding matrix with size number of embedding X embedding dimension\n        self._embedding_dim = embedding_dim\n        self._num_embeddings = num_embeddings\n        \n        self._embedding = nn.Embedding(self._num_embeddings, self._embedding_dim)\n        self._embedding.weight.data.uniform_(-1/self._num_embeddings, 1/self._num_embeddings)\n        self._commitment_cost = commitment_cost\n\n    def forward(self, inputs):\n        # convert inputs from BCHW -> BHWC\n        inputs = inputs.permute(0, 2, 3, 1).contiguous()\n        input_shape = inputs.shape\n        \n        # Flatten input\n        flat_input = inputs.view(-1, self._embedding_dim)\n        \n        # Calculate distances between flattened input and embedding vector\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            \n        # Choose indices that are min in each row\n        encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)\n        ## Create a matrix of dimensions B*H*W into number of embeddings\n        encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)\n        ### Convert index to on hot encoding \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 = F.mse_loss(quantized.detach(), inputs)\n        q_latent_loss = F.mse_loss(quantized, inputs.detach())\n        loss = q_latent_loss + self._commitment_cost * e_latent_loss\n        \n        quantized = inputs + (quantized - inputs).detach()\n        \n        avg_probs = torch.mean(encodings, dim=0)\n        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))\n        \n        # convert quantized from BHWC -> BCHW\n        return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"### Create Residual connections\nclass 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)\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) for _ in range(self._num_residual_layers)])\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)\n\n","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        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                                )\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        x = self._conv_2(x)\n        x = F.relu(x)\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        self._conv_1 = nn.Conv2d(in_channels=in_channels,\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        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        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    def forward(self,inputs):\n        x = self._conv_1(inputs)\n        x = self._residual_stack(x)\n        x = self._conv_trans_1(x)\n        x = F.relu(x)\n        return self._conv_trans_2(x)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"num_training_updates = 11000\nnum_hiddens = 128\nnum_residual_hiddens = 32\nnum_residual_layers = 2\nembedding_dim= 512\nnum_embeddings = 512\ncommitment_cost = 0.25\nlearning_rate = 1e-3\n","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,num_embeddings,embedding_dim,commitment_cost,decay=0):\n        super(Model,self).__init__()\n        self._encoder_= Encoder(3,num_hiddens,num_residual_layers,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        self._vq_vae = VectorQuantizer(num_embeddings,embedding_dim,commitment_cost)\n        self._decoder = Decoder(embedding_dim,\n                              num_hiddens,\n                              num_residual_layers,\n                              num_residual_hiddens)\n    def forward(self,x):\n        z = self._encoder_(x)\n        z = self._pre_vq_conv(z)\n        loss,quantized,perplexity,_ = self._vq_vae(z)\n        x_recon = self._decoder(quantized)\n        return loss,x_recon,perplexity\n            \n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = Model(num_hiddens,num_residual_layers,num_residual_hiddens,num_embeddings,embedding_dim,commitment_cost,decay=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"optimizer = optim.Adam(model.parameters(), lr=learning_rate, amsgrad=False)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.train()\ntrain_res_recon_error = []\ntrain_res_perplexity = []\n\nfor i in xrange(num_training_updates):\n    (data, _) = next(iter(training_loader))\n    data = data.to(device)\n    optimizer.zero_grad()\n\n    vq_loss, data_recon, perplexity = model(data)\n    recon_error = F.mse_loss(data_recon, data)\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":"model.eval()\n\n(valid_originals, _) = next(iter(validation_loader))\nvalid_originals = valid_originals.to(device)\n\nvq_output_eval = model._pre_vq_conv(model._encoder_(valid_originals))\n_, valid_quantize, _, _ = model._vq_vae(vq_output_eval)\nvalid_reconstructions = model._decoder(valid_quantize)\ndef 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)\n    \nshow(make_grid(valid_reconstructions.cpu().data)+0.5, )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# train_loss = []\n# test_loss = []\n\n\n# for epoch in range(0,num_training_updates):\n#     total_train_loss = 0\n#     total_test_loss = 0\n#     model.train()\n#     for x,_ in train_loader:\n        \n#         vq_loss ,data_recon,perplexity = model(x)\n#         recon_error = F.mse_loss(x,data_recon)\n#         loss = recon_error + vq_loss\n#         optimizer.zero_grad()\n#         loss.backward()\n#         optimizer.step()\n#         total_train_loss+=loss.item()\n#     train_loss.append(total_train_loss)\n#     print('--------------------------------')\n#     print('Epoch :' , epoch)\n#     print('Total Train Loss :', total_train_loss)\n        \n#     with torch.no_grad():\n#         model.eval()\n#         for x,_ in testing_loader:\n#             vq_loss,data_recon,perplexity = model(x)\n#             recon_error = F.mse_loss(data_recon,x)\n#             loss = recon_error+vq_loss\n#             total_test_loss+=loss.item()\n#     test_loss.append(total_test_loss)\n#     print('Epoch :' , epoch)\n#     print('Total Test Loss :', total_test_loss)\n    \n            \n            \n        ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# device = torch.device('cuda' if torch.cuda.is_available() else \"cpu\")\n\n# testing_loader = torch.utils.data.DataLoader(train_data,28,shuffle=True)\n\n\n# 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)\n\n\n# model.eval()\n\n# (valid_originals, _) = next(iter(testing_loader))\n# valid_originals = valid_originals.to(device)\n\n# vq_output_eval = model._pre_vq_conv(model._encoder_(valid_originals))\n# _, valid_quantize, _, _ = model._vq_vae(vq_output_eval)\n# valid_reconstructions = model._decoder(valid_quantize)\n# show(make_grid(valid_reconstructions.cpu().data), )","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"checkpoint = {'model': model,\n              'state_dict': model.state_dict(),\n              'optimizer' : optimizer.state_dict()}\n \ntorch.save(checkpoint, 'checkpoint.pth')","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}