{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"markdown","source":"### Check TPU is available\n\nThe cell below makes sure you have access to a TPU on Kaggle.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\n\ntry:\n    tpu = tf.distribute.cluster_resolver.TPUClusterResolver()\n    print(f'Running on TPU {tpu.master()}')\nexcept ValueError:\n    tpu = None\n\nif tpu:\n    tf.config.experimental_connect_to_cluster(tpu)\n    tf.tpu.experimental.initialize_tpu_system(tpu)\n    strategy = tf.distribute.experimental.TPUStrategy(tpu)\nelse:\n    strategy = tf.distribute.get_strategy()\n\n\nREPLICAS = strategy.num_replicas_in_sync\nprint(f'REPLICAS: {REPLICAS}')\nAUTO = tf.data.experimental.AUTOTUNE","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Generative Adersarial Networks (GANs)\nIn the landmark paper [Goodfellow et al.](https://arxiv.org/abs/1406.2661), published in 2014, authors introduced this novel paradigm for generative models. The fundamental idea proposed in the work is to train a Generator Network in adversarial setup, where a discriminator network downstream critiques the generated samples.\n\nSimply put, generator network generates a sample and discriminator network classifies it as a real or fake. Discriminator is also provided with real samples. The objective functions takes the following form:\n\n$$\\underset{G}{\\text{minimize}}\\; \\underset{D}{\\text{maximize}}\\; \\mathbb{E}_{x \\sim p_\\text{data}}\\left[\\log D(x)\\right] + \\mathbb{E}_{z \\sim p(z)}\\left[\\log \\left(1-D(G(z))\\right)\\right]$$\nwhere:\n$x \\sim p_\\text{data}$ are samples from the input data. $z \\sim p(z)$ are the random noise samples. $G(z)$ are the generated images using the neural network generator $G$, and $D$ is the output of the discriminator, specifying the probability of an input being real.\n\n### Training Setup\n\nThis example is **part 1 of 3 series** for **training a GAN on TPU using Torch XLA package**. This notebook illustrates distributed (data parallel) training of DC-GAN model using MNIST dataset on a TPU device. **The notebook will lay the foundations for our Cycle GAN training on TPU and more**. A TPU device consists of 4 chips (8 cores; 2 cores/chip). Both the discriminator and generator replica are created on each of 8 cores. The dataset is splitted across the 8 cores.\nAt every training step, each of the cores perfoms the forward (loss computation) and backward (gradient computation) on the given minibatch and then all_reduce is performed across TPU cores to update the parameters. Notice xm.optimizer_step call in the discriminator and optimizer train steps.\n\nGeneral GAN training looks like:\n\n* update the **generator** ($G$) to minimize the probability of the **discriminator making the correct choice**.\n* update the **discriminator** ($D$) to maximize the probability of the **discriminator making the correct choice**.\n\nWe will use a different objective when we update the generator: maximize the probability of the **discriminator making the incorrect choice.** This small change helps to alleviate problems with the generator gradient vanishing when the discriminator is confident. This is the standard update used in most GAN papers, and was used in the original paper from [Goodfellow et al..](https://arxiv.org/abs/1406.2661)\n\nTherefore the training loop in this notebook will entail:\n\n * Update the generator ($G$) to maximize the probability of the discriminator making the incorrect choice on generated data:$$\\underset{G}{\\text{maximize}}\\;  \\mathbb{E}_{z \\sim p(z)}\\left[\\log D(G(z))\\right]$$\n \n* Update the discriminator ($D$), to maximize the probability of the discriminator making the correct choice on real and generated data:$$\\underset{D}{\\text{maximize}}\\; \\mathbb{E}_{x \\sim p_\\text{data}}\\left[\\log D(x)\\right] + \\mathbb{E}_{z \\sim p(z)}\\left[\\log \\left(1-D(G(z))\\right)\\right]$$\n\n### Deep Convolutional Generative Adersarial Networks\n\nDCGAN is one of the popular and successful network design for GAN. It mainly composes of convolution layers without max pooling or fully connected layers. It uses convolutional stride and transposed convolution for the downsampling and the upsampling. The figure below is the network design for the generator and discriminator,\n\n![DC-GAN](https://gluon.mxnet.io/_images/dcgan.png)\n\nHere is the summary of DCGAN:\n   * Replace all `max pooling with convolutional stride`\n   * Use `transposed convolution` for upsampling.\n   * `Eliminate fully connected` layers.\n   * Use `Batch normalization` except the output layer for the generator and the input layer of the discriminator.\n   * Use `ReLU in the generator` except for the output which uses tanh.\n   * Use `LeakyReLU in the discriminator`.\n   \nHere are the tuning tips quote directly from the paper.\n> All models were trained with mini-batch stochastic gradient descent (SGD) with a mini-batch size of 128. All weights were initialized from a zero-centered Normal distribution with standard deviation 0.02. In the LeakyReLU, the slope of the leak was set to 0.2 in all models. While previous GAN work has used momentum to accelerate training, we used the Adam optimizer with tuned hyperparameters. We found the suggested learning rate of 0.001, to be too high, using 0.0002 instead. Additionally, we found leaving the momentum term β1 at the suggested value of 0.9 resulted in training oscillation and instability while reducing it to 0.5 helped stabilize training.\n\nThe simplicity of DCGAN contributes to its success. We reach certain bottleneck that increasing the complexity of the generator does not necessarily improve the image quality. Until we identify the bottleneck and know how to train GANs more effective, DCGAN remains a good start point for a new project.","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"### Setup Dependencies","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"#### Download Torch XLA nightly release","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n!python pytorch-xla-env-setup.py --version nightly  --apt-packages libomp5 libopenblas-dev","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Training Setup","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom torch import nn, optim\nfrom torchvision import transforms, datasets\nfrom torch.optim import Adam\nimport torch.nn.functional as F\n\nimport torch_xla\nimport torch_xla.core.xla_model as xm\nimport torch_xla.debug.metrics as met\nimport torch_xla.distributed.parallel_loader as pl\nimport torch_xla.distributed.xla_multiprocessing as xmp\nimport torch_xla.utils.utils as xu","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Setting up the Global Flags\n\nIn the current setup, Discriminator network was chosen to be a smaller capacity than generator. Even with similar capacity networks, generator update path is deeper than discriminator. Therefore uneven learning rates chosen here seems to yield a better convergence.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# Define Parameters\nFLAGS = {}\nFLAGS['datadir'] = \"/tmp/mnist\"\nFLAGS['batch_size'] = 128\nFLAGS['num_workers'] = 4\nFLAGS['gen_learning_rate'] = 0.005\nFLAGS['disc_learning_rate'] = 0.001\nFLAGS['num_epochs'] = 30\nFLAGS['num_cores'] = 8  ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Data & Image Utilities","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"from matplotlib.pyplot import imshow\nfrom matplotlib import pyplot as plt\nfrom IPython import display \nimport cv2\n    \nRESULT_IMG_PATH = '/tmp/test_result.png'\n\ndef plot_results(*images):\n    num_images = len(images)\n    n_rows = 4\n    n_columns =len(images) // n_rows\n    fig, axes = plt.subplots(n_rows, n_columns, figsize=(30, 18))\n\n    for i, ax in enumerate(fig.axes):\n        ax.axis('off') \n        if i >= num_images:\n          continue\n        img = images[i]\n        img = img.squeeze() # [1,Y,X] -> [Y,X]\n        ax.imshow(img)\n    plt.savefig(RESULT_IMG_PATH, transparent=True)\n\ndef display_results():\n    img = cv2.imread(RESULT_IMG_PATH, cv2.IMREAD_UNCHANGED)\n    plt.figure(figsize=(30,18))\n    plt.imshow(img)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def mnist_data():\n    compose = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])\n    out_dir = '{}/dataset'.format(FLAGS['datadir'])\n    return datasets.MNIST(root=out_dir, train=True, transform=compose, download=True)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Discriminator Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class DiscriminativeNet(torch.nn.Module):\n    \n    def __init__(self):\n        super(DiscriminativeNet, self).__init__()\n        self.conv1 = nn.Conv2d(1, 32, kernel_size=5)\n        self.bn1 = nn.BatchNorm2d(32)\n        self.conv2 = nn.Conv2d(32, 64, kernel_size=5)\n        self.bn2 = nn.BatchNorm2d(64)\n        self.fc1 = nn.Linear(4*4*64, 1)\n\n    def forward(self, x):\n        x = F.leaky_relu(F.max_pool2d(self.conv1(x), 2), 0.01)\n        x = self.bn1(x)\n        x = F.leaky_relu(F.max_pool2d(self.conv2(x), 2), 0.01)\n        x = self.bn2(x)\n        x = torch.flatten(x, 1)\n        x = F.leaky_relu(self.fc1(x), 0.01)\n        return torch.sigmoid(x)            \n        ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Generator Model","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"class GenerativeNet(torch.nn.Module):\n    \n    def __init__(self):\n        super(GenerativeNet, self).__init__()\n        self.input_size = 100\n        self.linear1 = nn.Linear(self.input_size, 1024)\n        self.bn1 = nn.BatchNorm1d(1024)\n        self.linear2 = nn.Linear(1024, 7*7*128)\n        self.bn2 = nn.BatchNorm1d(7*7*128)\n        self.conv1 = nn.ConvTranspose2d(\n            in_channels=128, \n            out_channels=64, \n            kernel_size=4,\n            stride=2, \n            padding=1, \n            bias=False\n        )\n        self.bn3 = nn.BatchNorm2d(64)\n        self.conv2 = nn.ConvTranspose2d(\n            in_channels=64, \n            out_channels=1, \n            kernel_size=4,\n            stride=2, \n            padding=1, \n            bias=False\n        )\n\n    # Noise\n    def generate_noise(self, size):\n        n = torch.randn(size, self.input_size)\n        return n \n              \n    def forward(self, x):\n        x = self.linear1(x)\n        x = F.relu(x)\n        x = self.bn1(x)\n        x = self.linear2(x)\n        x = F.relu(x)\n        x = self.bn2(x)\n        x = x.view(x.shape[0], 128, 7, 7)\n        x = self.conv1(x)\n        x = F.relu(x)\n        x = self.bn3(x)\n        x = self.conv2(x)\n        x = torch.tanh(x)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def init_weights(m):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1 or classname.find('BatchNorm') != -1:\n        m.weight.data.normal_(0.00, 0.02)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def real_data_target(size, device):\n    data = torch.ones(size, 1)\n    return data.to(device)\n\ndef fake_data_target(size, device):\n    data = torch.zeros(size, 1)\n    return data.to(device)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Note on the use of .detach() function\n\nYou will notice in the following code snippet that when the generator is used to create the fake_data, **.detach() for the discriminator training step, the .detach call is used to create a new view of the fake_data tensor for which the operations will not be recorded for gradient computation.**\n\n* Since fake_data is an output of an nn.module, by default, pytorch will record all the operations performed on this tensor during the forward pass as DAG. And after the backward pass these DAG and corresponding operations are cleared (unless retain_graph=True). Therefore such a tensor can be part of only one cone of logic where the forward and backward pass is done. If there are two loss function where this tensor is used and backward pass is performed on these two function (or even sum of the functions) for the second backward pass the operations DAG will not be found, leading to an error.\n\n* The second place, where detach() is used is when a numpy() call is to be made to tensor (for plotting purposes). Pytorch also requires that requires_grad should not be true on these tensor. (Ref: RuntimeError: Can't call numpy() on Variable that requires grad. Use var.detach().numpy() instead.)\n\n\n### Training ","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"SERIAL_EXEC = xmp.MpSerialExecutor()\n# Only instantiate model weights once in memory.\ngenerator = GenerativeNet()\ngenerator.apply(init_weights)\ndescriminator = DiscriminativeNet()\ndescriminator.apply(init_weights)\nWRAPPED_GENERATOR = xmp.MpModelWrapper(generator)\nWRAPPED_DISCRIMINATOR = xmp.MpModelWrapper(descriminator)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_gan(rank):\n    torch.manual_seed(1) \n    data = SERIAL_EXEC.run(lambda: mnist_data())\n    train_sampler = torch.utils.data.distributed.DistributedSampler(\n        data,\n        num_replicas=xm.xrt_world_size(),\n        rank=xm.get_ordinal(),\n        shuffle=True)\n    \n\n    # Create loader with data, so that we can iterate over it\n    train_loader = torch.utils.data.DataLoader(\n      data,\n      batch_size=FLAGS['batch_size'],\n      sampler=train_sampler,\n      num_workers=FLAGS['num_workers'],\n      drop_last=True)\n\n    # Num batches\n    num_batches = len(train_loader)\n    \n    device = xm.xla_device()\n    \n    generator = WRAPPED_GENERATOR.to(device)\n    discriminator = WRAPPED_DISCRIMINATOR.to(device)\n   \n    \n    # Optimizers\n    d_optimizer = Adam(discriminator.parameters(), lr=FLAGS['disc_learning_rate'], betas=(0.5, 0.999))\n    g_optimizer = Adam(generator.parameters(), lr=FLAGS['gen_learning_rate'], betas=(0.5, 0.999))\n\n    # Number of epochs\n    num_epochs = FLAGS['num_epochs'] \n    # Loss function\n    loss = nn.BCELoss()\n    \n\n    def train_step_discriminator(optimizer, real_data, fake_data, device):         \n        # Reset gradients\n        optimizer.zero_grad()\n\n        # 1. Train on Real Data\n        prediction_real = discriminator(real_data)\n        # Calculate error and backpropagate\n        error_real = loss(prediction_real, real_data_target(real_data.size(0), device))\n        \n\n        # 2. Train on Fake Data\n        prediction_fake = discriminator(fake_data)\n        # Calculate error and backpropagate\n        error_fake = loss(prediction_fake, fake_data_target(real_data.size(0), device))\n        \n        total_error = error_real + error_fake\n        total_error.backward()\n\n        # Update weights with gradients\n        xm.optimizer_step(optimizer)\n\n        return total_error, prediction_real, prediction_fake\n\n    def train_step_generator(optimizer, fake_data, device):\n        # Reset gradients\n        optimizer.zero_grad()\n        prediction = discriminator(fake_data)\n        # Calculate error and backpropagate\n        error = loss(prediction, real_data_target(prediction.size(0), device))\n        error.backward()\n        # Update weights with gradients\n        xm.optimizer_step(optimizer)\n\n        # Return error\n        return error\n\n    # Notice the use of .detach() when fake_data is to passed into discriminator\n    def train_loop_fn(loader):\n        tracker = xm.RateTracker()\n        for n_batch, (real_batch,_) in enumerate(loader):\n            # Train Step Descriminator\n            real_data = real_batch.to(device)\n            # sample noise and generate fake data\n            noise = generator.generate_noise(real_data.size(0)).to(device)\n            fake_data = generator(noise)\n            d_error, d_pred_real, d_pred_fake = train_step_discriminator(\n                d_optimizer, real_data, fake_data.detach(), device)\n            \n            #Train Step Generator\n            noise = generator.generate_noise(real_data.size(0)).to(device)\n            fake_data = generator(noise)\n            g_error = train_step_generator(g_optimizer, fake_data, device)\n        return d_error.item(), g_error.item()\n\n\n    for epoch in range(1, FLAGS['num_epochs'] + 1):\n        d_error, g_error = train_loop_fn (pl.MpDeviceLoader(train_loader, device))\n        xm.master_print(\"Finished training epoch {}: D_error:{}, G_error: {}\".format(epoch, d_error, g_error))\n        \n        if epoch == FLAGS['num_epochs']:\n            xm.master_print('Saving Model ..')\n            xm.save(generator.state_dict(), \"generator.bin\")\n            xm.save(discriminator.state_dict(), \"discriminator.bin\")\n            xm.master_print('Model Saved.')\n          \n    num_test_samples = 100\n    test_noise = generator.generate_noise(num_test_samples).to(device)\n    xm.do_on_ordinals(plot_results, generator(test_noise).detach(), (0,))","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"# Start training processes\ndef _mp_fn(rank, flags):\n    global FLAGS\n    FLAGS = flags\n    torch.set_default_tensor_type('torch.FloatTensor')\n    train_gan(rank)\n\nxmp.spawn(_mp_fn, args=(FLAGS,), nprocs=FLAGS['num_cores'],\n          start_method='fork')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"display_results()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"### References\n\n* This notebook is inspired from [Training DC-GAN using Colab Cloud TPU](https://github.com/pytorch/xla/blob/master/contrib/colab/DC-GAN.ipynb).\n\n* The [Unsupervised representation learning with Deep convolutional generative adversarial networks](https://arxiv.org/pdf/1511.06434.pdf) paper.\n\n* [[TPU Training] PyTorch nlp XLMRoberta](https://www.kaggle.com/rhtsingh/tpu-training-pytorch-nlp-xlmroberta)\n\n##### More to come, stay tuned.","execution_count":null}],"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}