BigGAN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kz7GMf9fruXG"
   },
   "source": [
    "# Components of BigGAN\n",
    "\n",
    "In this notebook, you'll learn about and implement the components of BigGAN, the first large-scale GAN architecture proposed in [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://arxiv.org/abs/1809.11096) (Brock et al. 2019). BigGAN performs a conditional generation task, so unlike StyleGAN, it conditions on a certain class to generate results. BigGAN is based mainly on empirical results and shows extremely good results when trained on ImageNet and its 1000 classes.\n",
    "\n",
    "The authors propose a several changes that improve state-of-the-art Inception Score (IS) and Frechet Inception Distance (FID), including:\n",
    " - **Increasing batch size by a factor of 8**, which improves IS by 46% and improves FID by 35%, but also induces complete mode collapse in training.\n",
    " - **Increasing the number of convolutional channels by 1.5x**, which improves IS by 21% and FID by 23%.\n",
    " - **Using shared class-conditional embeddings $c$ in BatchNorm layers**, which reduces the number of parameters and increases IS by 2% and FID by 4%.\n",
    " - **Adding skip connections from latent noise $z$** by concatenating chunks of $z$ to $c$. This improves IS by 1% and FID by 5%.\n",
    "\n",
    "> ![BigGAN Architecture](BigGAN.png)\n",
    "*BigGAN Architecture Components, taken from Figure 15 in [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://arxiv.org/abs/1809.11096) (Brock et al. 2019). (a) A typical architectural layout for BigGAN’s generator. See Appendix B for details. (b) A Residual Block (ResBlock up) in BigGAN’s generator. (c) A Residual Block (ResBlock down) in BigGAN’s discriminator.*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6vK-QqG_mefA"
   },
   "source": [
    "## The Truncation Trick and Orthogonal Regularization\n",
    "\n",
    "You should already be familiar with the truncation trick, which truncates the range of values of random noise $z$. Truncation to values close to 0 increases fidelity but decreases variety. Truncation to values further from 0 does the opposite.\n",
    "\n",
    "Truncation results in a different distribution of $z$ values from the one seen in training, which can cause saturation artifacts. The authors address this by making $G$ well-defined, or *smooth*, on the full distribution of $z$ values.\n",
    "\n",
    "To do this, they employ orthogonal regularization, first introduced in [Neural Photo Editing with Introspective Adversarial Networks](https://arxiv.org/abs/1609.07093) (Brock et al. 2017). The authors modify this regularization technique for BigGAN and formulate it as\n",
    "\n",
    "\\begin{align*}\n",
    "  R_\\beta(W) = \\beta\\big|\\big|W^\\top W \\odot (\\pmb{1} - I)\\big|\\big|^2_F,\n",
    "\\end{align*}\n",
    "where $\\pmb{1}$ denotes a matrix of 1's. This regularization term removes the diagonal terms from the regularization and aims to minimize the pairwise cosine similarity between filters without constraining their norm.\n",
    "\n",
    "> ![Truncation Trick](BigGAN-truncation-trick.png)\n",
    "*Generated images with different truncation thresholds, taken from Figure 2 in [Large Scale GAN Training for High Fidelity Natural Image Synthesis](https://arxiv.org/abs/1809.11096) (Brock et al. 2019). (a) The effects of increasing truncation. From left to right, the threshold is set to 2, 1, 0.5, 0.04. (b) Saturation artifacts from applying truncation to a poorly conditioned model.*\n",
    "\n",
    "Below is the implementation for orthogonal regularization. You can refer to the StyleGAN notebook for the truncation trick code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dPpIAAOjykrQ"
   },
   "outputs": [],
   "source": [
    "# Some setup\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "def orthogonal_regularization(weight):\n",
    "    '''\n",
    "    Function for computing the orthogonal regularization term for a given weight matrix.\n",
    "    '''\n",
    "    weight = weight.flatten(1)\n",
    "    return torch.norm(\n",
    "        torch.dot(weight, weight) * (torch.ones_like(weight) - torch.eye(weight.shape[0]))\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uAtedvsisf1j"
   },
   "source": [
    "## BigGAN Parts\n",
    "\n",
    "Before jumping into the full implementation, let's first take a look at some submodules that will be important in our BigGAN implementation later."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "h-F0HvK8RMzI"
   },
   "source": [
    "### Class-conditional Batch Normalization\n",
    "\n",
    "Recall that batch norm aims to normalize activation statistics to a standard gaussian distribution (via an exponential moving average of minibatch mean and variances) but also applies trainable parameters, $\\gamma$ and $\\beta$, to invert this operation if the model sees fit:\n",
    "\n",
    "\\begin{align*}\n",
    "    y &= \\dfrac{x - \\hat{\\mu}}{\\hat{\\sigma} + \\epsilon} * \\gamma + \\beta.\n",
    "\\end{align*}\n",
    "\n",
    "BigGAN injects class-conditional information by parameterizing $\\gamma$ and $\\beta$ as linear transformations of the class embedding, $c$. Recall that BigGAN also concatenates $c$ with $z$ skip connections (denoted $[c, z]$), so\n",
    "\n",
    "\\begin{align*}\n",
    "    \\gamma &:= W_\\gamma^\\top[c, z] \\\\\n",
    "    \\beta &:= W_\\beta^\\top[c, z]\n",
    "\\end{align*}\n",
    "\n",
    "The idea is actually very similar to the adaptive instance normalization (AdaIN) module that you implemented in the StyleGAN notebook, so we've copied that code in comments below for reference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xCGq1YSz3ezh"
   },
   "outputs": [],
   "source": [
    "class ClassConditionalBatchNorm2d(nn.Module):\n",
    "    '''\n",
    "    ClassConditionalBatchNorm2d Class\n",
    "    Values:\n",
    "    in_channels: the dimension of the class embedding (c) + noise vector (z), a scalar\n",
    "    out_channels: the dimension of the activation tensor to be normalized, a scalar\n",
    "    '''\n",
    "\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "        self.bn = torch.nn.BatchNorm2d(out_channels)\n",
    "        self.class_scale_transform = nn.utils.spectral_norm(nn.Linear(in_channels, out_channels, bias=False))\n",
    "        self.class_shift_transform = nn.utils.spectral_norm(nn.Linear(in_channels, out_channels, bias=False))\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        normalized_image = self.bn(x)\n",
    "        class_scale = (1 + self.class_scale_transform(y))[:, :, None, None]\n",
    "        class_shift = self.class_shift_transform(y)[:, :, None, None]\n",
    "        transformed_image = class_scale * normalized_image + class_shift\n",
    "        return transformed_image\n",
    "\n",
    "# class AdaIN(nn.Module):\n",
    "#     '''\n",
    "#     AdaIN Class, extends/subclass of nn.Module\n",
    "#     Values:\n",
    "#       channels: the number of channels the image has, a scalar\n",
    "#       w_dim: the dimension of the intermediate tensor, w, a scalar \n",
    "#     '''\n",
    "\n",
    "#     def __init__(self, channels, w_dim):\n",
    "#         super().__init__()\n",
    "#         self.instance_norm = nn.InstanceNorm2d(channels)\n",
    "#         self.style_scale_transform = nn.Linear(w_dim, channels)\n",
    "#         self.style_shift_transform = nn.Linear(w_dim, channels)\n",
    "\n",
    "#     def forward(self, image, w):\n",
    "#         normalized_image = self.instance_norm(image)\n",
    "#         style_scale = self.style_scale_transform(w)[:, :, None, None]\n",
    "#         style_shift = self.style_shift_transform(w)[:, :, None, None]\n",
    "#         transformed_image = style_scale * normalized_image + style_shift\n",
    "#         return transformed_image\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "m5MrMUuvW5ac"
   },
   "source": [
    "### Self-Attention Block\n",
    "\n",
    "As you may already know, self-attention has been a successful technique in helping models learn arbitrary, long-term dependencies. [Self-Attention Generative Adversarial Networks](https://arxiv.org/abs/1805.08318) (Zhang et al. 2018) first introduced the self-attention mechanism into the GAN architecture. BigGAN augments its residual blocks with these attention blocks.\n",
    "\n",
    "**A Quick Primer on Self-Attention**\n",
    "\n",
    "Self-attention is just **scaled dot product attention**. Given a sequence $S$ (with images, $S$ is just the image flattened across its height and width), the model learns mappings to query ($Q$), key ($K$), and value ($V$) matrices:\n",
    "\n",
    "\\begin{align*}\n",
    "    Q &:= W_q^\\top S \\\\\n",
    "    K &:= W_k^\\top S \\\\\n",
    "    V &:= W_v^\\top S\n",
    "\\end{align*}\n",
    "\n",
    "where $W_q$, $W_k$, and $W_v$ are learned parameters. The subsequent self-attention mechanism is then computed as\n",
    "\n",
    "\\begin{align*}\n",
    "    \\text{Attention}(Q, K, V) = \\text{softmax}\\left(\\dfrac{QK^\\top}{\\sqrt{d_k}}\\right)V\n",
    "\\end{align*}\n",
    "\n",
    "where $d_k$ is the dimensionality of the $Q, K$ matrices (SA-GAN and BigGAN both omit this term). Intuitively, you can think of the *query* matrix as containing the representations of each position with respect to itself and the *key* matrix as containing the representations of each position with respect to the others. How important two positions are to each other is measured by dot product as $QK^\\top$, hence **dot product attention**. A softmax is applied to convert these relative importances to a probability distribution over all positions.\n",
    "\n",
    "\n",
    "Intuitively, the *value* matrix provides the importance weighting of the attention at each position, hence **scaled dot product attention**. Relevant positions should be assigned larger weight and irrelevant ones should be assigned smaller weight.\n",
    "\n",
    "Don't worry if you don't understand this right away - it's a tough concept! For extra reading, you should check out [Attention Is All You Need](https://arxiv.org/abs/1706.03762) (Vaswani et al. 2017), which is the paper that first introduces this technique, and [The Illustrated Transformer](http://jalammar.github.io/illustrated-transformer/), which breaks down and explains the self-attention mechanism clearly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "75rBFkY6XQGR"
   },
   "outputs": [],
   "source": [
    "class AttentionBlock(nn.Module):\n",
    "    '''\n",
    "    AttentionBlock Class\n",
    "    Values:\n",
    "    channels: number of channels in input\n",
    "    '''\n",
    "    def __init__(self, channels):\n",
    "        super().__init__()\n",
    "\n",
    "        self.channels = channels\n",
    "\n",
    "        self.theta = nn.utils.spectral_norm(nn.Conv2d(channels, channels // 8, kernel_size=1, padding=0, bias=False))\n",
    "        self.phi = nn.utils.spectral_norm(nn.Conv2d(channels, channels // 8, kernel_size=1, padding=0, bias=False))\n",
    "        self.g = nn.utils.spectral_norm(nn.Conv2d(channels, channels // 2, kernel_size=1, padding=0, bias=False))\n",
    "        self.o = nn.utils.spectral_norm(nn.Conv2d(channels // 2, channels, kernel_size=1, padding=0, bias=False))\n",
    "\n",
    "        self.gamma = nn.Parameter(torch.tensor(0.), requires_grad=True)\n",
    "\n",
    "    def forward(self, x):\n",
    "        spatial_size = x.shape[2] * x.shape[3]\n",
    "\n",
    "        # Apply convolutions to get query (theta), key (phi), and value (g) transforms\n",
    "        theta = self.theta(x)\n",
    "        phi = F.max_pool2d(self.phi(x), kernel_size=2)\n",
    "        g = F.max_pool2d(self.g(x), kernel_size=2)\n",
    "\n",
    "        # Reshape spatial size for self-attention\n",
    "        theta = theta.view(-1, self.channels // 8, spatial_size)\n",
    "        phi = phi.view(-1, self.channels // 8, spatial_size // 4)\n",
    "        g = g.view(-1, self.channels // 2, spatial_size // 4)\n",
    "\n",
    "        # Compute dot product attention with query (theta) and key (phi) matrices\n",
    "        beta = F.softmax(torch.bmm(theta.transpose(1, 2), phi), dim=-1)\n",
    "\n",
    "        # Compute scaled dot product attention with value (g) and attention (beta) matrices\n",
    "        o = self.o(torch.bmm(g, beta.transpose(1, 2)).view(-1, self.channels // 2, x.shape[2], x.shape[3]))\n",
    "\n",
    "        # Apply gain and residual\n",
    "        return self.gamma * o + x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ChIw0phn3l2U"
   },
   "source": [
    "## BigGAN Generator\n",
    "\n",
    "Before implementing the generator in full, you first need to implement the generator residual block.\n",
    "\n",
    "### Generator Residual Block\n",
    "\n",
    "As with many state-of-the-art computer vision models, BigGAN employs skip connections in the form of residual blocks to map random noise to a fake image. You can think of BigGAN residual blocks as having 3 steps. Given input $x$ and class embedding $y$:\n",
    " 1. $h :=$ `bn-relu-upsample-conv`$(x, y)$\n",
    " 2. $h :=$ `bn-relu-conv`$(h, y)$\n",
    " 3. $x :=$ `upsample-conv`$(x)$,\n",
    "\n",
    "after which you can apply a residual connection and return $h + x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NZelofm5T-fy"
   },
   "outputs": [],
   "source": [
    "class GResidualBlock(nn.Module):\n",
    "    '''\n",
    "    GResidualBlock Class\n",
    "    Values:\n",
    "    c_dim: the dimension of conditional vector [c, z], a scalar\n",
    "    in_channels: the number of channels in the input, a scalar\n",
    "    out_channels: the number of channels in the output, a scalar\n",
    "    '''\n",
    "\n",
    "    def __init__(self, c_dim, in_channels, out_channels):\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv1 = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))\n",
    "        self.conv2 = nn.utils.spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))\n",
    "\n",
    "        self.bn1 = ClassConditionalBatchNorm2d(c_dim, in_channels)\n",
    "        self.bn2 = ClassConditionalBatchNorm2d(c_dim, out_channels)\n",
    "\n",
    "        self.activation = nn.ReLU()\n",
    "        self.upsample_fn = nn.Upsample(scale_factor=2)     # upsample occurs in every gblock\n",
    "\n",
    "        self.mixin = (in_channels != out_channels)\n",
    "        if self.mixin:\n",
    "            self.conv_mixin = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0))\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        # h := upsample(x, y)\n",
    "        h = self.bn1(x, y)\n",
    "        h = self.activation(h)\n",
    "        h = self.upsample_fn(h)\n",
    "        h = self.conv1(h)\n",
    "\n",
    "        # h := conv(h, y)\n",
    "        h = self.bn2(h, y)\n",
    "        h = self.activation(h)\n",
    "        h = self.conv2(h)\n",
    "\n",
    "        # x := upsample(x)\n",
    "        x = self.upsample_fn(x)\n",
    "        if self.mixin:\n",
    "            x = self.conv_mixin(x)\n",
    "\n",
    "        return h + x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VaXN6UeehAFm"
   },
   "source": [
    "You can now implement the BigGAN generator in full!! Below is an implementation of the base model (at 128x128 resolution) from the paper.\n",
    "\n",
    "> This implementation uses `nn.ModuleList` for convenience. If you're not familiar with this, you can think of it as simply a Pythonic list that registers your modules with the Pytorch backend. For more information, see the [torch.nn.ModuleList](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html) documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZoZ1WEbH3xLc"
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    '''\n",
    "    Generator Class\n",
    "    Values:\n",
    "    z_dim: the dimension of random noise sampled, a scalar\n",
    "    shared_dim: the dimension of shared class embeddings, a scalar\n",
    "    base_channels: the number of base channels, a scalar\n",
    "    bottom_width: the height/width of image before it gets upsampled, a scalar\n",
    "    n_classes: the number of image classes, a scalar\n",
    "    '''\n",
    " \n",
    "    def __init__(self, base_channels=96, bottom_width=4, z_dim=120, shared_dim=128, n_classes=1000):\n",
    "        super().__init__()\n",
    " \n",
    "        n_chunks = 6    # 5 (generator blocks) + 1 (generator input)\n",
    "        self.z_chunk_size = z_dim // n_chunks\n",
    "        self.z_dim = z_dim\n",
    "        self.shared_dim = shared_dim\n",
    "        self.bottom_width = bottom_width\n",
    " \n",
    "        # No spectral normalization on embeddings, which authors observe to cripple the generator\n",
    "        self.shared_emb = nn.Embedding(n_classes, shared_dim)\n",
    " \n",
    "        self.proj_z = nn.Linear(self.z_chunk_size, 16 * base_channels * bottom_width ** 2)\n",
    " \n",
    "        # Can't use one big nn.Sequential since we are adding class+noise at each block\n",
    "        self.g_blocks = nn.ModuleList([\n",
    "            nn.ModuleList([\n",
    "                GResidualBlock(shared_dim + self.z_chunk_size, 16 * base_channels, 16 * base_channels),\n",
    "                AttentionBlock(16 * base_channels),\n",
    "            ]),\n",
    "            nn.ModuleList([\n",
    "                GResidualBlock(shared_dim + self.z_chunk_size, 16 * base_channels, 8 * base_channels),\n",
    "                AttentionBlock(8 * base_channels),\n",
    "            ]),\n",
    "            nn.ModuleList([\n",
    "                GResidualBlock(shared_dim + self.z_chunk_size, 8 * base_channels, 4 * base_channels),\n",
    "                AttentionBlock(4 * base_channels),\n",
    "            ]),\n",
    "            nn.ModuleList([\n",
    "                GResidualBlock(shared_dim + self.z_chunk_size, 4 * base_channels, 2 * base_channels),\n",
    "                AttentionBlock(2 * base_channels),\n",
    "            ]),\n",
    "            nn.ModuleList([\n",
    "                GResidualBlock(shared_dim + self.z_chunk_size, 2 * base_channels, base_channels),\n",
    "                AttentionBlock(base_channels),\n",
    "            ]),\n",
    "        ])\n",
    "        self.proj_o = nn.Sequential(\n",
    "            nn.BatchNorm2d(base_channels),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.utils.spectral_norm(nn.Conv2d(base_channels, 3, kernel_size=1, padding=0)),\n",
    "            nn.Tanh(),\n",
    "        )\n",
    " \n",
    "    def forward(self, z, y):\n",
    "        '''\n",
    "        z: random noise with size self.z_dim\n",
    "        y: class embeddings with size self.shared_dim\n",
    "            = NOTE =\n",
    "            y should be class embeddings from self.shared_emb, not the raw class labels\n",
    "        '''\n",
    "        # Chunk z and concatenate to shared class embeddings\n",
    "        zs = torch.split(z, self.z_chunk_size, dim=1)\n",
    "        z = zs[0]\n",
    "        ys = [torch.cat([y, z], dim=1) for z in zs[1:]]\n",
    " \n",
    "        # Project noise and reshape to feed through generator blocks\n",
    "        h = self.proj_z(z)\n",
    "        h = h.view(h.size(0), -1, self.bottom_width, self.bottom_width)\n",
    " \n",
    "        # Feed through generator blocks\n",
    "        for idx, g_block in enumerate(self.g_blocks):\n",
    "            h = g_block[0](h, ys[idx])\n",
    "            h = g_block[1](h)\n",
    " \n",
    "        # Project to 3 RGB channels with tanh to map values to [-1, 1]\n",
    "        h = self.proj_o(h)\n",
    " \n",
    "        return h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EmWXtRgVfQMF"
   },
   "source": [
    "## BigGAN Discriminator\n",
    "\n",
    "Before implementing the discriminator in full, you need to implement a discriminator residual block, which is simpler than the generator's. Note that the last residual block does not apply downsampling.\n",
    " 1. $h :=$ `relu-conv-relu-downsample`$(x)$\n",
    " 2. $x :=$ `conv-downsample`$(x)$\n",
    "\n",
    "In the official BigGAN implementation, the architecture is slightly different for the first discriminator residual block, since it handles the raw image as input:\n",
    " 1. $h :=$ `conv-relu-downsample`$(x)$\n",
    " 2. $x :=$ `downsample-conv`$(x)$\n",
    "\n",
    "After these two steps, you can return the residual connection $h + x$. You might notice that there is no class information in these residual blocks. As you'll see later in the code, the authors inject class-conditional information after the final hidden layer (and before the output layer) via channel-wise dot product."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JLm65xx1hoz-"
   },
   "outputs": [],
   "source": [
    "class DResidualBlock(nn.Module):\n",
    "    '''\n",
    "    DResidualBlock Class\n",
    "    Values:\n",
    "    in_channels: the number of channels in the input, a scalar\n",
    "    out_channels: the number of channels in the output, a scalar\n",
    "    downsample: whether to apply downsampling\n",
    "    use_preactivation: whether to apply an activation function before the first convolution\n",
    "    '''\n",
    "\n",
    "    def __init__(self, in_channels, out_channels, downsample=True, use_preactivation=False):\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv1 = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1))\n",
    "        self.conv2 = nn.utils.spectral_norm(nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))\n",
    "\n",
    "        self.activation = nn.ReLU()\n",
    "        self.use_preactivation = use_preactivation  # apply preactivation in all except first dblock\n",
    "\n",
    "        self.downsample = downsample    # downsample occurs in all except last dblock\n",
    "        if downsample:\n",
    "            self.downsample_fn = nn.AvgPool2d(2)\n",
    "        self.mixin = (in_channels != out_channels) or downsample\n",
    "        if self.mixin:\n",
    "            self.conv_mixin = nn.utils.spectral_norm(nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0))\n",
    "\n",
    "    def _residual(self, x):\n",
    "        if self.use_preactivation:\n",
    "            if self.mixin:\n",
    "                x = self.conv_mixin(x)\n",
    "            if self.downsample:\n",
    "                x = self.downsample_fn(x)\n",
    "        else:\n",
    "            if self.downsample:\n",
    "                x = self.downsample_fn(x)\n",
    "            if self.mixin:\n",
    "                x = self.conv_mixin(x)\n",
    "        return x\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Apply preactivation if applicable\n",
    "        if self.use_preactivation:\n",
    "            h = F.relu(x)\n",
    "        else:\n",
    "            h = x\n",
    "\n",
    "        h = self.conv1(h)\n",
    "        h = self.activation(h)\n",
    "        if self.downsample:\n",
    "            h = self.downsample_fn(h)\n",
    "\n",
    "        return h + self._residual(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F4Re2GakokIU"
   },
   "source": [
    "Now implement the BigGAN discriminator in full!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "siICOXTaon2p"
   },
   "outputs": [],
   "source": [
    "class Discriminator(nn.Module):\n",
    "    '''\n",
    "    Discriminator Class\n",
    "    Values:\n",
    "    base_channels: the number of base channels, a scalar\n",
    "    n_classes: the number of image classes, a scalar\n",
    "    '''\n",
    "\n",
    "    def __init__(self, base_channels=96, n_classes=1000):\n",
    "        super().__init__()\n",
    "\n",
    "        # For adding class-conditional evidence\n",
    "        self.shared_emb = nn.utils.spectral_norm(nn.Embedding(n_classes, 16 * base_channels))\n",
    "\n",
    "        self.d_blocks = nn.Sequential(\n",
    "            DResidualBlock(3, base_channels, downsample=True, use_preactivation=False),\n",
    "            AttentionBlock(base_channels),\n",
    "\n",
    "            DResidualBlock(base_channels, 2 * base_channels, downsample=True, use_preactivation=True),\n",
    "            AttentionBlock(2 * base_channels),\n",
    "\n",
    "            DResidualBlock(2 * base_channels, 4 * base_channels, downsample=True, use_preactivation=True),\n",
    "            AttentionBlock(4 * base_channels),\n",
    "\n",
    "            DResidualBlock(4 * base_channels, 8 * base_channels, downsample=True, use_preactivation=True),\n",
    "            AttentionBlock(8 * base_channels),\n",
    "\n",
    "            DResidualBlock(8 * base_channels, 16 * base_channels, downsample=True, use_preactivation=True),\n",
    "            AttentionBlock(16 * base_channels),\n",
    "\n",
    "            DResidualBlock(16 * base_channels, 16 * base_channels, downsample=False, use_preactivation=True),\n",
    "            AttentionBlock(16 * base_channels),\n",
    "\n",
    "            nn.ReLU(inplace=True),\n",
    "        )\n",
    "        self.proj_o = nn.utils.spectral_norm(nn.Linear(16 * base_channels, 1))\n",
    "\n",
    "    def forward(self, x, y=None):\n",
    "        h = self.d_blocks(x)\n",
    "        h = torch.sum(h, dim=[2, 3])\n",
    "\n",
    "        # Class-unconditional output\n",
    "        uncond_out = self.proj_o(h)\n",
    "        if y is None:\n",
    "            return uncond_out\n",
    "\n",
    "        # Class-conditional output\n",
    "        cond_out = torch.sum(self.shared_emb(y) * h, dim=1, keepdim=True)\n",
    "        return uncond_out + cond_out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sj8aXxNLtGU_"
   },
   "source": [
    "## Setting Up BigGAN Training\n",
    "\n",
    "Now you're are ready to set up BigGAN for training! Unfortunately, this notebook will not provide actual training code due to the size of BigGAN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vdkfLu4wuHbh"
   },
   "outputs": [],
   "source": [
    "device = 'cpu'\n",
    "\n",
    "# Initialize models\n",
    "base_channels = 96\n",
    "z_dim = 120\n",
    "n_classes = 5   # 5 classes is used instead of the original 1000, for efficiency\n",
    "shared_dim = 128\n",
    "generator = Generator(base_channels=base_channels, bottom_width=4, z_dim=z_dim, shared_dim=shared_dim, n_classes=n_classes).to(device)\n",
    "discriminator = Discriminator(base_channels=base_channels, n_classes=n_classes).to(device)\n",
    "\n",
    "# Initialize weights orthogonally\n",
    "for module in generator.modules():\n",
    "    if (isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear) or isinstance(module, nn.Embedding)):\n",
    "        nn.init.orthogonal_(module.weight)\n",
    "for module in discriminator.modules():\n",
    "    if (isinstance(module, nn.Conv2d) or isinstance(module, nn.Linear) or isinstance(module, nn.Embedding)):\n",
    "        nn.init.orthogonal_(module.weight)\n",
    "\n",
    "# Initialize optimizers\n",
    "g_optimizer = torch.optim.Adam(generator.parameters(), lr=1e-4, betas=(0.0, 0.999), eps=1e-6)\n",
    "d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=4e-4, betas=(0.0, 0.999), eps=1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RHyU49y57PTb"
   },
   "source": [
    "Here is a sample forward pass:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E7YtKzgMxukz"
   },
   "outputs": [],
   "source": [
    "batch_size = n_classes\n",
    "\n",
    "z = torch.randn(batch_size, z_dim, device=device)                 # Generate random noise (z)\n",
    "y = torch.arange(start=0, end=n_classes, device=device).long()    # Generate a batch of labels (y), one for each class\n",
    "y_emb = generator.shared_emb(y)                                   # Retrieve class embeddings (y_emb) from generator\n",
    "\n",
    "x_gen = generator(z, y_emb)                                       # Generate fake images from z and y_emb\n",
    "score = discriminator(x_gen, y)                                   # Generate classification for fake images"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8-H2WoQDh-Fs"
   },
   "source": [
    "## BigGAN-deep\n",
    "\n",
    "Initially, the authors of the BigGAN paper didn't find much help in increasing the depth of the network. But they experimented further (research is always improving!) and added a few notes about an additional architecture, called BigGAN-deep. This modification of BigGAN is 4x deeper, sports a modified residual block architecture, and concatenates the entire $z$ vector to $c$ (as opposed to separate chunks at different resolutions).\n",
    "\n",
    "Typically on a difficult and complex task that you're unlikely to overfit, you expect better performance when a model has more parameters, because it has more room to learn. Surprisingly, BigGAN-deep has fewer parameters than its BigGAN counterpart. Architectural optimizations such as using depthwise separable convolutions and truncating/concatenating channels in skip connections (as opposed to using pointwise convolutions) decrease parameters without trading expressivity.\n",
    "\n",
    "For more details on the BigGAN-deep architecture, see Appendix B of the paper.\n",
    "\n",
    "And as for the implementation of the BigGAN-deep variant, well, that's left as an exercise for the reader. You're a smart cookie, you'll figure it out! Just keep in mind that with great power comes great responsibility ;)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "C2W3: Components of BigGAN (Optional).ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}