{"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},"cell_type":"code","source":"import os\nimport xml.etree.ElementTree as ET\n\nimport torch\nimport torchvision\n\n# for testing only\nfrom tqdm import tqdm_notebook as tqdm\nimport matplotlib.pyplot as plt\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# This loader will use the underlying loader plus crop the image based on the annotation\ndef doggo_loader(path):\n    img = torchvision.datasets.folder.default_loader(path) # default loader\n    \n    # Get bounding box\n    annotation_basename = os.path.splitext(os.path.basename(path))[0]\n    annotation_dirname = next(dirname for dirname in os.listdir('../input/annotation/Annotation/') if dirname.startswith(annotation_basename.split('_')[0]))\n    annotation_filename = os.path.join('../input/annotation/Annotation', annotation_dirname, annotation_basename)\n    tree = ET.parse(annotation_filename)\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    bbox = (xmin, ymin, xmax, ymax)\n    \n    # return cropped image\n    return img.crop(bbox)\n\n\n# The dataset (example)\ndataset = torchvision.datasets.ImageFolder(\n    '../input/all-dogs/',\n    loader=doggo_loader, # THE CUSTOM LOADER\n    transform=torchvision.transforms.Compose([\n        torchvision.transforms.Resize(70),\n        torchvision.transforms.CenterCrop(64),\n        torchvision.transforms.ToTensor(),\n        torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n    ]) # some transformations, add your data preprocessing here\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check that it all loads without a bug\nfor i in tqdm(range(len(dataset))):\n    _ = dataset[i]\nprint('Ok.')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Check that we get only the CUTE DOGS OH YES WHOS THE GOOD DOGGO ITS YOU\nn = 10\n_, axes = plt.subplots(figsize=(4*n, 4*n), ncols=n, nrows=n)\nfor i, ax in enumerate(axes.flatten()):\n    ax.imshow(dataset[i][0].permute(1, 2, 0).detach().numpy())\nplt.show()\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"len(dataset)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_loader = torch.utils.data.DataLoader(dataset=dataset,\n                                          batch_size=128,\n                                          shuffle=True,\n                                          num_workers=0)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# obtain one batch of training images\ndataiter = iter(train_loader)\nimages, labels = dataiter.next()\n\n# plot the images in the batch, along with the corresponding labels\nfig = plt.figure(figsize=(25, 4))\nplot_size=20\nfor idx in np.arange(plot_size):\n    ax = fig.add_subplot(2, plot_size/2, idx+1, xticks=[], yticks=[])\n    ax.imshow(np.transpose(images[idx], (1, 2, 0)))\n    # print out the correct label for each image\n    # .item() gets the value contained in a Tensor\n    ax.set_title(str(labels[idx].item()))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# current range\nimg = images[0]\n\nprint('Min: ', img.min())\nprint('Max: ', img.max())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # helper scale function\n# def scale(x, feature_range=(-1, 1)):\n#     ''' Scale takes in an image x and returns that image, scaled\n#        with a feature_range of pixel values from -1 to 1. \n#        This function assumes that the input x is already scaled from 0-1.'''\n#     # assume x is scaled to (0, 1)\n#     # scale to feature_range and return scaled x\n#     min, max = feature_range\n#     x = x * (max - min) + min\n#     return x\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# # scaled range\n# scaled_img = scale(img)\n\n# print('Scaled min: ', scaled_img.min())\n# print('Scaled max: ', scaled_img.max())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.nn as nn\nimport torch.nn.functional as F\n\n# helper conv function\ndef conv(in_channels, out_channels, kernel_size, stride=2, padding=1, batch_norm=True):\n    \"\"\"Creates a convolutional layer, with optional batch normalization.\n    \"\"\"\n    layers = []\n    conv_layer = nn.Conv2d(in_channels, out_channels, \n                           kernel_size, stride, padding, bias=False)\n    \n    # append conv layer\n    layers.append(conv_layer)\n\n    if batch_norm:\n        # append batchnorm layer\n        layers.append(nn.BatchNorm2d(out_channels))\n     \n    # using Sequential container\n    return nn.Sequential(*layers)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Discriminator(nn.Module):\n\n    def __init__(self, conv_dim=64):\n        super(Discriminator, self).__init__()\n\n        # complete init function\n        self.conv_dim = conv_dim\n\n        # 64x64 input\n        self.conv1 = conv(3, conv_dim, 4, batch_norm=False) # first layer, no batch_norm\n        # 32x32 out\n        self.conv2 = conv(conv_dim, conv_dim*2, 4)\n        # 16x16 out\n        self.conv3 = conv(conv_dim*2, conv_dim*4, 4)\n        # 8x8 out\n        self.conv4 = conv(conv_dim*4, conv_dim*8, 4)\n        \n        # final, fully-connected layer\n        self.fc = nn.Linear(conv_dim*8*4*4, 1)\n\n    def forward(self, x):\n        # all hidden layers + leaky relu activation\n        out = F.leaky_relu(self.conv1(x), 0.2)\n        out = F.leaky_relu(self.conv2(out), 0.2)\n        out = F.leaky_relu(self.conv3(out), 0.2)\n        out = F.leaky_relu(self.conv4(out), 0.2)\n        \n        # flatten\n        out = out.view(-1, self.conv_dim*8*4*4)\n        out = self.fc(out)        \n        return out   \n                # final output layer\n        \n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# helper deconv function\ndef deconv(in_channels, out_channels, kernel_size, stride=2, padding=1, batch_norm=True):\n    \"\"\"Creates a transposed-convolutional layer, with optional batch normalization.\n    \"\"\"\n    # create a sequence of transpose + optional batch norm layers\n    layers = []\n    transpose_conv_layer = nn.ConvTranspose2d(in_channels, out_channels, \n                                              kernel_size, stride, padding, bias=False)\n    # append transpose convolutional layer\n    layers.append(transpose_conv_layer)\n    \n    if batch_norm:\n        # append batchnorm layer\n        layers.append(nn.BatchNorm2d(out_channels))\n        \n    return nn.Sequential(*layers)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Generator(nn.Module):\n    \n    def __init__(self, z_size, conv_dim=64):\n        super(Generator, self).__init__()\n\n        # complete init function\n        \n        self.conv_dim = conv_dim\n        \n        # first, fully-connected layer\n        self.fc = nn.Linear(z_size, conv_dim*8*4*4)\n\n        # transpose conv layers\n        self.t_conv1 = deconv(conv_dim*8, conv_dim*4, 4)\n        self.t_conv2 = deconv(conv_dim*4, conv_dim*2, 4)\n        self.t_conv3 = deconv(conv_dim*2, conv_dim, 4)\n        self.t_conv4 = deconv(conv_dim, 3, 4, batch_norm=False)\n        \n\n    def forward(self, x):\n        # fully-connected + reshape \n        out = self.fc(x)\n        out = out.view(-1, self.conv_dim*8, 4, 4) # (batch_size, depth, 4, 4)\n        \n        # hidden transpose conv layers + relu\n        out = F.relu(self.t_conv1(out))\n        out = F.relu(self.t_conv2(out))\n        out = F.relu(self.t_conv3(out))\n        \n        # last layer + tanh activation\n        out = self.t_conv4(out)\n        out = torch.tanh(out)\n        \n        return out\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# define hyperparams\nconv_dim = 64\nz_size = 100\n\n# define discriminator and generator\nD = Discriminator(conv_dim)\nG = Generator(z_size=z_size, conv_dim=conv_dim)\n\nprint(D)\nprint()\nprint(G)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_on_gpu = torch.cuda.is_available()\n\nif train_on_gpu:\n    # move models to GPU\n    G.cuda()\n    D.cuda()\n    print('GPU available for training. Models moved to GPU')\nelse:\n    print('Training on CPU.')\n    ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def real_loss(D_out, smooth=False):\n    batch_size = D_out.size(0)\n    # label smoothing\n    if smooth:\n        # smooth, real labels = 0.9\n        labels = torch.ones(batch_size)*0.9\n    else:\n        labels = torch.ones(batch_size) # real labels = 1\n    # move labels to GPU if available     \n    if train_on_gpu:\n        labels = labels.cuda()\n    # binary cross entropy with logits loss\n    criterion = nn.BCEWithLogitsLoss()\n    # calculate loss\n    loss = criterion(D_out.squeeze(), labels)\n    return loss\n\ndef fake_loss(D_out):\n    batch_size = D_out.size(0)\n    labels = torch.zeros(batch_size) # fake labels = 0\n    if train_on_gpu:\n        labels = labels.cuda()\n    criterion = nn.BCEWithLogitsLoss()\n    # calculate loss\n    loss = criterion(D_out.squeeze(), labels)\n    return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch.optim as optim\n\n# params\nlr = 0.0002\nbeta1=0.5\nbeta2=0.999 # default value\n\n# Create optimizers for the discriminator and generator\nd_optimizer = optim.Adam(D.parameters(), lr, [beta1, beta2])\ng_optimizer = optim.Adam(G.parameters(), lr, [beta1, beta2])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import pickle as pkl\n\n# training hyperparams\nnum_epochs = 6\n\n# keep track of loss and generated, \"fake\" samples\nsamples = []\nlosses = []\n\nprint_every = 300\n\n# Get some fixed data for sampling. These are images that are held\n# constant throughout training, and allow us to inspect the model's performance\nsample_size=16\nfixed_z = np.random.uniform(-1, 1, size=(sample_size, z_size))\nfixed_z = torch.from_numpy(fixed_z).float()\n\n# train the network\nfor epoch in range(num_epochs):\n    \n    for batch_i, (real_images, _) in enumerate(train_loader):\n                \n        batch_size = real_images.size(0)\n        \n        # important rescaling step\n        real_images = (real_images)\n        \n        # ============================================\n        #            TRAIN THE DISCRIMINATOR\n        # ============================================\n        \n        d_optimizer.zero_grad()\n        \n        # 1. Train with real images\n\n        # Compute the discriminator losses on real images \n        if train_on_gpu:\n            real_images = real_images.cuda()\n        D_real = D(real_images)\n        d_real_loss = real_loss(D_real)\n        \n        # 2. Train with fake images\n        \n        # Generate fake images\n        z = np.random.uniform(-1, 1, size=(batch_size, z_size))\n        z = torch.from_numpy(z).float()\n        # move x to GPU, if available\n        if train_on_gpu:\n            z = z.cuda()\n        fake_images = G(z)\n        \n        # Compute the discriminator losses on fake images            \n        D_fake = D(fake_images)\n        d_fake_loss = fake_loss(D_fake)\n        \n        # add up loss and perform backprop\n        d_loss = d_real_loss + d_fake_loss\n        d_loss.backward()\n        d_optimizer.step()\n        \n        \n        # =========================================\n        #            TRAIN THE GENERATOR\n        # =========================================\n        g_optimizer.zero_grad()\n        \n        # 1. Train with fake images and flipped labels\n        \n        # Generate fake images\n        z = np.random.uniform(-1, 1, size=(batch_size, z_size))\n        z = torch.from_numpy(z).float()\n        if train_on_gpu:\n            z = z.cuda()\n        fake_images = G(z)\n        \n        # Compute the discriminator losses on fake images \n        # using flipped labels!\n        D_fake = D(fake_images)\n        g_loss = real_loss(D_fake) # use real loss to flip labels\n        \n        # perform backprop\n        g_loss.backward()\n        g_optimizer.step()\n\n        # Print some loss stats\n        if batch_i % print_every == 0:\n            # append discriminator loss and generator loss\n            losses.append((d_loss.item(), g_loss.item()))\n            # print discriminator and generator loss\n            print('Epoch [{:5d}/{:5d}] | d_loss: {:6.4f} | g_loss: {:6.4f}'.format(\n                    epoch+1, num_epochs, d_loss.item(), g_loss.item()))\n\n    \n    ## AFTER EACH EPOCH##    \n    # generate and save sample, fake images\n    G.eval() # for generating samples\n    if train_on_gpu:\n        fixed_z = fixed_z.cuda()\n    samples_z = G(fixed_z)\n    samples.append(samples_z)\n    G.train() # back to training mode\n\n\n# Save training generator samples\nwith open('train_samples.pkl', 'wb') as f:\n    pkl.dump(samples, f)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"fig, ax = plt.subplots()\nlosses = np.array(losses)\nplt.plot(losses.T[0], label='Discriminator', alpha=0.5)\nplt.plot(losses.T[1], label='Generator', alpha=0.5)\nplt.title(\"Training Losses\")\nplt.legend()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from torchvision.utils import save_image\nif not os.path.exists('../output_images'):\n    os.mkdir('../output_images')\nsample_size=50\nfixed_z = np.random.uniform(-1, 1, size=(sample_size, z_size))\nfixed_z = torch.from_numpy(fixed_z).float()\nim_batch_size = 50\nn_images=10000\nfor i_batch in range(0, n_images, im_batch_size):\n    G.eval() # for generating samples\n    if train_on_gpu:\n        fixed_z = fixed_z.cuda()\n    gen_images = G(fixed_z)\n    images = gen_images.to(\"cpu\").clone().detach()\n    images = images.numpy().transpose(0, 2, 3, 1)\n    for i_image in range(gen_images.size(0)):\n        save_image(gen_images[i_image, :, :, :], os.path.join('../output_images', f'image_{i_batch+i_image:05d}.png'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nshutil.make_archive('images', 'zip', '../output_images')","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}