{"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\nimport os\nfor 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 5GB 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":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.transforms as transforms\nimport torchvision\nimport torch.utils.data as data\nimport numpy as np","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# 2-d latent space, parameter count in same order of magnitude\n# as in the original VAE paper (VAE paper has about 3x as many)\nlatent_dims = 2\nnum_epochs = 20 # Originally at 100\nbatch_size = 128\ncapacity = 64\nlearning_rate = 1e-3\nvariational_beta = 1\nuse_gpu = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Make the dataset\n\n# Configure directory to look appropriate for making dataset later\ntry:\n  !mkdir ../working/actual_monet_jpg\n  !cp -R ../input/gan-getting-started/monet_jpg ../working/actual_monet_jpg/monet_jpg_inner\nexcept:\n  pass\n\n# Define paramaters and transformation for data augmentation\nBATCH_SIZE = 10\nDATA_PATH = '../working/actual_monet_jpg'\nmyTransforms = transforms.Compose([\n    transforms.ColorJitter(hue=.05, saturation=.05),\n    transforms.RandomHorizontalFlip(),\n    transforms.ToTensor()\n])\n\n# make loader and data loader\ntrain_data = torchvision.datasets.ImageFolder(root=DATA_PATH, transform = myTransforms)\ntrain_data = torch.utils.data.ConcatDataset([train_data] * 32)\ntrain_data_loader = data.DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True)\nprint(\"Done\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Encoder(nn.Module):\n    def __init__(self):\n        super(Encoder, self).__init__()\n        c = capacity\n        self.conv00 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=4, stride=3, padding=0) # out: (c/4) x 85 x 85\n        self.conv01 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=4, stride=3, padding=0) # out: (c/2) x 85 x 85\n        self.conv1 = nn.Conv2d(in_channels=32, out_channels=c, kernel_size=4, stride=2, padding=1) # out: c x 14 x 14 (CHANGE in_channels=1 for monet)\n        self.conv2 = nn.Conv2d(in_channels=c, out_channels=c*2, kernel_size=4, stride=2, padding=1) # out: c x 7 x 7\n        self.fc_mu = nn.Linear(in_features=c*2*7*7, out_features=latent_dims)\n        self.fc_logvar = nn.Linear(in_features=c*2*7*7, out_features=latent_dims)\n            \n    def forward(self, x):\n        x = F.relu(self.conv00(x)) # for monet\n        x = F.relu(self.conv01(x)) # for monet\n        x = F.relu(self.conv1(x))\n        x = F.relu(self.conv2(x))\n        x = x.view(x.size(0), -1) # flatten batch of multi-channel feature maps to a batch of feature vectors\n        x_mu = self.fc_mu(x)\n        x_logvar = self.fc_logvar(x)\n        return x_mu, x_logvar\n\nclass Decoder(nn.Module):\n    def __init__(self):\n        super(Decoder, self).__init__()\n        c = capacity\n        self.fc = nn.Linear(in_features=latent_dims, out_features=c*2*7*7)\n        self.conv2 = nn.ConvTranspose2d(in_channels=c*2, out_channels=c, kernel_size=4, stride=2, padding=1)\n        self.conv1 = nn.ConvTranspose2d(in_channels=c, out_channels=32, kernel_size=4, stride=2, padding=1) # change to out_channels = 32 for monet\n        self.conv01 = nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=4, stride=3, padding=0)\n        self.conv00 = nn.ConvTranspose2d(in_channels=16, out_channels=3, kernel_size=4, stride=3, padding=0)\n            \n    def forward(self, x):\n        x = self.fc(x)\n        x = x.view(x.size(0), capacity*2, 7, 7) # unflatten batch of feature vectors to a batch of multi-channel feature maps\n        x = F.relu(self.conv2(x))\n        x = F.relu(self.conv1(x))\n        x = F.relu(self.conv01(x))\n        x = torch.sigmoid(self.conv00(x)) # last layer before output is sigmoid, since we are using BCE as reconstruction loss\n        return x\n    \nclass VariationalAutoencoder(nn.Module):\n    def __init__(self):\n        super(VariationalAutoencoder, self).__init__()\n        self.encoder = Encoder()\n        self.decoder = Decoder()\n    \n    def forward(self, x):\n        latent_mu, latent_logvar = self.encoder(x)\n        latent = self.latent_sample(latent_mu, latent_logvar)\n        x_recon = self.decoder(latent)\n        return x_recon, latent_mu, latent_logvar\n    \n    def latent_sample(self, mu, logvar):\n        if self.training:\n            # the reparameterization trick\n            std = logvar.mul(0.5).exp_()\n            eps = torch.empty_like(std).normal_()\n            return eps.mul(std).add_(mu)\n        else:\n            return mu\n    \ndef vae_loss(recon_x, x, mu, logvar):\n    # recon_x is the probability of a multivariate Bernoulli distribution p.\n    # -log(p(x)) is then the pixel-wise binary cross-entropy.\n    # Averaging or not averaging the binary cross-entropy over all pixels here\n    # is a subtle detail with big effect on training, since it changes the weight\n    # we need to pick for the other loss term by several orders of magnitude.\n    # Not averaging is the direct implementation of the negative log likelihood,\n    # but averaging makes the weight of the other loss term independent of the image resolution.\n    size = x.size()[0] * x.size()[1] * x.size()[2] * x.size()[3]\n    #print(size)\n    recon_loss = F.binary_cross_entropy(recon_x.view(-1, size), x.view(-1, size), reduction='sum') # change 784 to 6291456 for monet\n    \n    # KL-divergence between the prior distribution over latent vectors\n    # (the one we are going to sample from when generating new images)\n    # and the distribution estimated by the generator for the given image.\n    kldivergence = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())\n    \n    return recon_loss + variational_beta * kldivergence\n    \n    \nvae = VariationalAutoencoder()\n\ndevice = torch.device(\"cuda:0\" if use_gpu and torch.cuda.is_available() else \"cpu\")\nvae = vae.to(device)\n\nnum_params = sum(p.numel() for p in vae.parameters() if p.requires_grad)\nprint('Number of parameters: %d' % num_params)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"optimizer = torch.optim.Adam(params=vae.parameters(), lr=learning_rate, weight_decay=1e-5)\n\n# set to training mode\nvae.train()\n\ntrain_loss_avg = []\n\nprint('Training ...')\nfor epoch in range(num_epochs):\n    train_loss_avg.append(0)\n    num_batches = 0\n    \n    for image_batch, _ in train_data_loader:\n        #print(image_batch.size())\n        image_batch = image_batch.to(device)\n        #print(image_batch.size())\n\n        # vae reconstruction\n        image_batch_recon, latent_mu, latent_logvar = vae(image_batch)\n        # RuntimeError: Given groups=1, weight of size [64, 1, 4, 4], expected input[32, 3, 256, 256] to have 1 channels, but got 3 channels instead\n        \n        # reconstruction error\n        loss = vae_loss(image_batch_recon, image_batch, latent_mu, latent_logvar)\n        \n        # backpropagation\n        optimizer.zero_grad()\n        loss.backward()\n        \n        # one step of the optmizer (using the gradients from backpropagation)\n        optimizer.step()\n        \n        train_loss_avg[-1] += loss.item()\n        num_batches += 1\n        \n    train_loss_avg[-1] /= num_batches\n    print('Epoch [%d / %d] average reconstruction error: %f' % (epoch+1, num_epochs, train_loss_avg[-1]))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Plot training loss across epochs\n\nimport matplotlib.pyplot as plt\nplt.ion()\n\nfig = plt.figure()\nplt.plot(train_loss_avg)\nplt.xlabel('Epochs')\nplt.ylabel('Loss')\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def to_img(x):\n    x = x.clamp(0, 1)\n    return x\n\ndef show_image(img):\n    img = to_img(img)\n    npimg = img.numpy()\n    plt.imshow(np.transpose(npimg, (1, 2, 0)))\n\nvae.eval()\n\nwith torch.no_grad():\n\n    # sample latent vectors from the normal distribution\n    latent = torch.randn(7000, latent_dims, device=device)\n\n    # reconstruct images from the latent vectors\n    img_recon = vae.decoder(latent)\n    img_recon = img_recon.cpu()\n\n    fig, ax = plt.subplots(figsize=(5, 5))\n    show_image(torchvision.utils.make_grid(img_recon.data[:9],3))\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from zipfile import ZipFile\nzipObj = ZipFile('images.zip', 'w')\nfor i in range(7000):\n    img = img_recon[i]\n    fp = 'img' + str(i) + '.jpg'\n    torchvision.utils.save_image(img, fp)\n    zipObj.write(fp)\nzipObj.close()\n!rm *.jpg\nprint(\"done\")","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}