{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github", "colab_type": "text" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "id": "87ce1c2d", "metadata": { "id": "87ce1c2d" }, "source": [ "# Getting Started\n", "\n", "## Overview\n", "This semester, all homeworks will be conducted through Google Colab notebooks. All code for the homework assignment will be written and run in this notebook. Running in Colab will automatically provide a GPU, but you may also run this notebook locally by following [these instructions](https://research.google.com/colaboratory/local-runtimes.html) if you wish to use your own GPU.\n", "\n", "You will save images in the notebooks to use and fill out a given LaTeX template which will be submitted to Gradescope, along with your notebook code.\n", "\n", "## Using Colab\n", "On the left-hand side, you can click the different icons to see a Table of Contents of the assignment, as well as local files accessible through the notebook.\n", "\n", "Make sure to go to **Runtime -> Change runtime type** and select **GPU** as the hardware accelerator. This allows you to use a GPU. Run the cells below to get started on the assignment. Note that a session is open for a maximum of 12 hours, and using too much GPU compute may result in restricted access for a short period of time. Please start the homework early so you have ample time to work.\n", "\n", "**If you loaded this notebook from clicking \"Open in Colab\" from github, you will need to save it to your own Google Drive to keep your work.**\n", "\n", "## General Tips\n", "In each homework problem, you will implement and train various diffusion models.\n", "\n", "Feel free to print whatever output (e.g. debugging code, training code, etc) you want, as the graded submission will be the submitted pdf with images.\n", "\n", "After you complete the assignment, download all of the images outputted in the results/ folder and upload them to the figure folder in the given latex template.\n", "\n", "Run the cells below to download and load up the starter code." ] }, { "cell_type": "code", "execution_count": null, "id": "d7d20590", "metadata": { "id": "d7d20590" }, "outputs": [], "source": [ "!if [ -d deepul ]; then rm -Rf deepul; fi\n", "!git clone https://github.com/rll/deepul.git\n", "!pip install ./deepul\n", "!pip install scikit-learn" ] }, { "cell_type": "code", "execution_count": null, "id": "3f55a3dd", "metadata": { "id": "3f55a3dd" }, "outputs": [], "source": [ "from deepul.hw4_helper import *\n", "import warnings\n", "warnings.filterwarnings('ignore')" ] }, { "cell_type": "markdown", "id": "3c2427a6", "metadata": { "id": "3c2427a6" }, "source": [ "# Question 1: Toy Dataset [30pt]\n", "\n", "In this question, we will train a simple diffusion models a toy 2D dataset.\n", "\n", "Execute the cell below to visualize our datasets." ] }, { "cell_type": "code", "execution_count": null, "id": "4838953c", "metadata": { "id": "4838953c" }, "outputs": [], "source": [ "visualize_q1_dataset()" ] }, { "cell_type": "markdown", "id": "e4766a8a", "metadata": { "id": "e4766a8a" }, "source": [ "For code simplicity, we will train a continuous-time variant of the diffusion prompt. In practice training objectives and code between discrete-time and continuous-time diffusion models are similar.\n", "\n", "Given a data element $x$ and neural net $f_\\theta(x, t)$, implement the following diffusion training steps:\n", "1. Sample the diffusion timestep: $t \\sim \\text{Uniform}(0, 1)$\n", "2. Compute the noise-strength following a cosine schedule: $\\alpha_t = \\cos\\left(\\frac{\\pi}{2}t\\right), \\sigma_t = \\sin\\left(\\frac{\\pi}{2}t\\right)$\n", "3. Apply the forward process - Sample noise $\\epsilon \\sim N(0,I)$ (same shape as $x$) and compute noised $x_t = \\alpha_t x + \\sigma_t \\epsilon$\n", "4. Estimate $\\hat{\\epsilon} = f_\\theta(x_t, t)$\n", "5. Optimize the loss $L = \\lVert \\epsilon - \\hat{\\epsilon} \\rVert_2^2$. Here, it suffices to just take the mean over all dimensions.\n", "\n", "Note that for the case of continuous-time diffusion, the forward process is $x_{0\\to1}$ and reverse process is $x_{1\\to0}$\n", "\n", "Use an MLP for $f_\\theta$ to optimize the loss. You may find the following details helpful.\n", "* Normalize the data using mean and std computed from the train dataset\n", "* Train 100 epochs, batch size 1024, Adam with LR 1e-3 (100 warmup steps, cosine decay to 0)\n", "* MLP with 4 hidden layers and hidden size 64\n", "* Condition on t by concatenating it with input x (i.e. 2D x + 1D t = 3D cat(x, t))\n", "\n", "To sample, implement the standard DDPM sampler. You may find the equation from the [DDIM paper](https://arxiv.org/pdf/2010.02502.pdf) helpful, rewritten and re-formatted here for convenience.\n", "$$x_{t-1} = \\alpha_{t-1}\\left(\\frac{x_t - \\sigma_t\\hat{\\epsilon}}{\\alpha_t}\\right) + \\sqrt{\\sigma_{t-1}^2 - \\eta_t^2}\\hat{\\epsilon} + \\eta_t\\epsilon_t$$\n", "where $\\epsilon_t \\sim N(0, I)$ is random Gaussian noise. For DDPM, let\n", "$$\\eta_t = \\sigma_{t-1}/\\sigma_t\\sqrt{1 - \\alpha_t^2/\\alpha_{t-1}^2}$$\n", "To run the reverse process, start from $x_1 \\sim N(0, I)$ and perform `num_steps` DDPM updates (a hyperparameter), pseudocode below.\n", "```\n", "ts = linspace(1 - 1e-4, 1e-4, num_steps + 1)\n", "x = sample_normal\n", "for i in range(num_steps):\n", " t = ts[i]\n", " tm1 = ts[i + 1]\n", " eps_hat = model(x, t)\n", " x = DDPM_UPDATE(x, eps_hat, t, tm1)\n", "return x\n", "```\n", "Note: If you encounter NaNs, you may need to clip $\\sigma_{t-1}^2 - \\eta_t^2$ to 0 if it goes negative, as machine precision issues can make it a very small negative number (e.g. -1e-12) if its too close to 0" ] }, { "cell_type": "code", "execution_count": null, "id": "b4e19d1c", "metadata": { "id": "b4e19d1c" }, "outputs": [], "source": [ "def q1(train_data, test_data):\n", " \"\"\"\n", " train_data: A (100000, 2) numpy array of 2D points\n", " test_data: A (10000, 2) numpy array of 2D points\n", "\n", " Returns\n", " - a (# of training iterations,) numpy array of train losses evaluated every minibatch\n", " - a (# of num_epochs + 1,) numpy array of test losses evaluated at the start of training and the end of every epoch\n", " - a numpy array of size (9, 2000, 2) of samples drawn from your model.\n", " Draw 2000 samples for each of 9 different number of diffusion sampling steps\n", " of evenly logarithmically spaced integers 1 to 512\n", " hint: np.power(2, np.linspace(0, 9, 9)).astype(int)\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"\n", "\n", " return train_losses, test_losses, all_samples" ] }, { "cell_type": "code", "execution_count": null, "id": "ed50dbad", "metadata": { "scrolled": true, "id": "ed50dbad" }, "outputs": [], "source": [ "q1_save_results(q1)" ] }, { "cell_type": "markdown", "id": "fc8fcc41", "metadata": { "id": "fc8fcc41" }, "source": [ "# Question 2: Pixel-Space Diffusion on CIFAR-10 [30pt]\n", "\n", "In this question, we will train pixel-space UNet diffusion model on CIFAR-10\n", "\n", "Execute the cell below to visualize our datasets." ] }, { "cell_type": "code", "execution_count": null, "id": "44f9d642", "metadata": { "id": "44f9d642" }, "outputs": [], "source": [ "visualize_q2_data()" ] }, { "cell_type": "markdown", "id": "370c91ef", "metadata": { "id": "370c91ef" }, "source": [ "We'll use a UNet architecture similar to the original [DDPM](https://arxiv.org/abs/2006.11239) paper. We provide the following pseudocode for each part of the model:\n", "```\n", "def timestep_embedding(timesteps, dim, max_period=10000):\n", " half = dim // 2\n", " freqs = np.exp(-np.log(max_period) * np.arange(0, half, dtype=float32) / half)\n", " args = timesteps[:, None].astype(float32) * freqs[None]\n", " embedding = cat([np.cos(args), np.sin(args)], axis=-1)\n", " if dim % 2:\n", " embedding = cat([embedding, np.zeros_like(embedding[:, :1])], axis=-1)\n", " return embedding\n", "\n", "ResidualBlock(in_channels, out_channels, temb_channels)\n", " Given x, temb\n", " h = Conv2d(in_channels, out_channels, 3, padding=1)(x)\n", " h = GroupNorm(num_groups=8, num_channels=out_channels)(h)\n", " h = SiLU()(h)\n", " \n", " temb = Linear(temb_channels, out_channels)(temb)\n", " h += temb[:, :, None, None] # h is BxDxHxW, temb is BxDx1x1\n", " \n", " h = Conv2d(out_channels, out_channels, 3, padding=1)(h)\n", " h = GroupNorm(num_groups=8, num_channels=out_channels)(h)\n", " h = SiLU()(h)\n", " \n", " if in_channels != out_channels:\n", " x = Conv2d(in_channels, out_channels, 1)(x)\n", " return x + h\n", " \n", "Downsample(in_channels)\n", " Given x\n", " return Conv2d(in_channels, in_channels, 3, stride=2, padding=1)(x)\n", "\n", "Upsample(in_channels)\n", " Given x\n", " x = interpolate(x, scale_factor=2)\n", " x = Conv2d(in_channels, in_channels, 3, padding=1)(x)\n", " return x\n", " \n", "UNet(in_channels, hidden_dims, blocks_per_dim)\n", " Given x, t\n", " temb_channels = hidden_dims[0] * 4\n", " emb = timestep_embedding(t, hidden_dims[0])\n", " emb = Sequential(Linear(hidden_dims[0], temb_channels), SiLU(), Linear(temb_channels, temb_channels))(emb)\n", " \n", " h = Conv2d(in_channels, hidden_dims[0], 3, padding=1)(x)\n", " hs = [h]\n", " prev_ch = hidden_dims[0]\n", " down_block_chans = [prev_ch]\n", " for i, hidden_dim in enumerate(hidden_dims):\n", " for _ in range(blocks_per_dim):\n", " h = ResidualBlock(prev_ch, hidden_dim, temb_channels)(h, emb)\n", " hs.append(h)\n", " prev_ch = hidden_dim\n", " down_block_chans.append(prev_ch)\n", " if i != len(hidden_dims) - 1:\n", " h = Downsample(prev_ch)(h)\n", " hs.append(h)\n", " down_block_chans.append(prev_ch)\n", " \n", " h = ResidualBlock(prev_ch, prev_ch, temb_channels)(h, emb)\n", " h = ResidualBlock(prev_ch, prev_ch, temb_channels)(h, emb)\n", " \n", " for i, hidden_dim in list(enumerate(hidden_dims))[::-1]:\n", " for j in range(blocks_per_dim + 1):\n", " dch = down_block_chans.pop()\n", " h = ResidualBlock(prev_ch + dch, hidden_dim, temb_channels)(cat(h, hs.pop()), emb)\n", " prev_ch = hidden_dim\n", " if i and j == blocks_per_dim:\n", " h = Upsample(prev_ch)(h)\n", " \n", " h = GroupNorm(num_groups=8, num_channels=prev_ch)(h)\n", " h = SiLU()(h)\n", " out = Conv2d(prev_ch, in_channels, 3, padding=1)(h)\n", " return out\n", "```" ] }, { "cell_type": "markdown", "id": "7ff725c4", "metadata": { "id": "7ff725c4" }, "source": [ "**Hyperparameter details**\n", "* Normalize data to [-1, 1]\n", "* UNET with hidden_dims as [64, 128, 256, 512] and 2 blocks_per_dim\n", "* Train 60 epochs, batch size 256, Adam with LR 1e-3 (100 warmup steps, cosine decay to 0)\n", "* For diffusion schedule, sampling and loss, use the same setup as Q1\n", "\n", "You may also find it helpful to clip $\\hat{x} = \\frac{x_t - \\sigma_t \\hat{\\epsilon}}{\\alpha_t}$ to [-1, 1] during each sampling step." ] }, { "cell_type": "code", "execution_count": null, "id": "20bf77f2", "metadata": { "id": "20bf77f2" }, "outputs": [], "source": [ "def q2(train_data, test_data):\n", " \"\"\"\n", " train_data: A (50000, 32, 32, 3) numpy array of images in [0, 1]\n", " test_data: A (10000, 32, 32, 3) numpy array of images in [0, 1]\n", "\n", " Returns\n", " - a (# of training iterations,) numpy array of train losses evaluated every minibatch\n", " - a (# of num_epochs + 1,) numpy array of test losses evaluated at the start of training and the end of every epoch\n", " - a numpy array of size (10, 10, 32, 32, 3) of samples in [0, 1] drawn from your model.\n", " The array represents a 10 x 10 grid of generated samples. Each row represents 10 samples generated\n", " for a specific number of diffusion timesteps. Do this for 10 evenly logarithmically spaced integers\n", " 1 to 512, i.e. np.power(2, np.linspace(0, 9, 10)).astype(int)\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"\n", "\n", " return train_losses, test_losses, samples" ] }, { "cell_type": "code", "execution_count": null, "id": "76ea5a82", "metadata": { "id": "76ea5a82" }, "outputs": [], "source": [ "q2_save_results(q2)" ] }, { "cell_type": "markdown", "id": "d445a644", "metadata": { "id": "d445a644" }, "source": [ "# Question 3: Class-Conditional Latent-Space Diffusion on CIFAR-10 with DiT [60pt]\n", "\n", "In this question, we will train latent-space [Diffusion Transformer (DiT)](https://arxiv.org/abs/2212.09748) model on CIFAR-10 **with class conditioning.**\n", "\n", "Execute the cell below to visualize our datasets." ] }, { "cell_type": "code", "execution_count": null, "id": "7bb87c66", "metadata": { "id": "7bb87c66" }, "outputs": [], "source": [ "visualize_q3_data()" ] }, { "cell_type": "markdown", "id": "5478a849", "metadata": { "id": "5478a849" }, "source": [ "## Part 3(a) VAE reconstructions and Scale Factor [10pt]\n", "\n", "Similar to how we learned a AR model in VQGAN latent space for homework 1, in this question, you will train a diffusion model in the latent space of a VAE. Note that since diffusion models can model continuous distributions, we do not need a discretization bottleneck in the VAE, and the latent space itself is continuous.\n", "\n", "Below, we specify each of the relevant properties or functions that you may need." ] }, { "cell_type": "code", "execution_count": null, "id": "a5fc0bfd", "metadata": { "id": "a5fc0bfd" }, "outputs": [], "source": [ "# @property\n", "# def latent_shape(self) -> Tuple[int, int, int]:\n", "# \"\"\"Size of the encoded representation\"\"\"\n", "#\n", "# def encode(self, x: np.ndarray) -> np.ndarray:\n", "# \"\"\"Encode an image x. Note: Channel dim is in dim 1\n", "#\n", "# Args:\n", "# x (np.ndarray, dtype=float32): Image to encode. shape=(batch_size, 3, 32, 32). Values in [-1, 1]\n", "#\n", "# Returns:\n", "# np.ndarray: Encoded image. shape=(batch_size, 4, 8, 8). Unbounded values\n", "# \"\"\"\n", "#\n", "# def decode(self, z: np.ndarray) -> np.ndarray:\n", "# \"\"\"Decode an encoded image.\n", "#\n", "# Args:\n", "# z (np.ndarray, dtype=float32): Encoded image. shape=(batch_size, 4, 8, 8). Unbounded values.\n", "#\n", "# Returns:\n", "# np.ndarray: Decoded image. shape=(batch_size, 3, 32, 32). Values in [-1, 1]\n", "# \"\"\"\n", "#" ] }, { "cell_type": "markdown", "id": "923f5fe7", "metadata": { "id": "923f5fe7" }, "source": [ "In this part, feed the given images through the VAE to compute and visualize reconstructions. In addition, you will compute a scale factor that will be needed during diffusion training to help normalize the data.\n", "\n", "To estimate the scale factor, encode 1000 images into the VAE latent space, flatten the entire tensor along all dimensions, and compute the standard deviation." ] }, { "cell_type": "code", "execution_count": null, "id": "ef106e4e", "metadata": { "id": "ef106e4e" }, "outputs": [], "source": [ "def q3_a(images, vae):\n", " \"\"\"\n", " images: (1000, 32, 32, 3) numpy array in [0, 1], the images to pass through the encoder and decoder of the vae\n", " vae: a vae model, trained on the relevant dataset\n", "\n", " Returns\n", " - a numpy array of size (50, 2, 32, 32, 3) of the decoded image in [0, 1] consisting of pairs\n", " of real and reconstructed images\n", " - a float that is the scale factor\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"\n", "\n", " return autoencoded_images, scale_factor" ] }, { "cell_type": "code", "execution_count": null, "id": "1ae6798e", "metadata": { "id": "1ae6798e" }, "outputs": [], "source": [ "q3a_save_results(q3_a)" ] }, { "cell_type": "markdown", "id": "c572cf95", "metadata": { "id": "c572cf95" }, "source": [ "## Part 3(b) Diffusion Transformer [30pt]\n", "In this part, you will train a Diffusion Transformer (Dit) on the latent space of the above pretrained VAE. You can use your Transformer implementation from HW1 as the core part of the DiT implementation.\n", "\n", "Below, we outline the key modifications needed on top of the standard Transformer for DiT.\n", "```\n", "def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):\n", " assert embed_dim % 2 == 0\n", "\n", " # use half of dimensions to encode grid_h\n", " emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)\n", " emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)\n", "\n", " emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)\n", " return emb\n", "\n", "\n", "def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):\n", " assert embed_dim % 2 == 0\n", " omega = np.arange(embed_dim // 2, dtype=np.float64)\n", " omega /= embed_dim / 2.\n", " omega = 1. / 10000**omega # (D/2,)\n", "\n", " pos = pos.reshape(-1) # (M,)\n", " out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product\n", "\n", " emb_sin = np.sin(out) # (M, D/2)\n", " emb_cos = np.cos(out) # (M, D/2)\n", "\n", " emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)\n", " return emb\n", "\n", "def get_2d_sincos_pos_embed(embed_dim, grid_size):\n", " grid_h = np.arange(grid_size, dtype=np.float32)\n", " grid_w = np.arange(grid_size, dtype=np.float32)\n", " grid = np.meshgrid(grid_w, grid_h) # here w goes first\n", " grid = np.stack(grid, axis=0)\n", "\n", " grid = grid.reshape([2, 1, grid_size, grid_size])\n", " pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)\n", " return pos_embed\n", "\n", "def modulate(x, shift, scale):\n", " return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)\n", "\n", "DiTBlock(hidden_size, num_heads)\n", " Given x (B x L x D), c (B x D)\n", " c = SiLU()(c)\n", " c = Linear(hidden_size, 6 * hidden_size)(c)\n", " shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = c.chunk(6, dim=1)\n", " \n", " h = LayerNorm(hidden_size, elementwise_affine=False)(x)\n", " h = modulate(h, shift_msa, scale_msa)\n", " x = x + gate_msa.unsqueeze(1) * Attention(hidden_size, num_heads)(h)\n", " \n", " h = LayerNorm(hidden_size, elementwise_affine=False)(x)\n", " h = modulate(h, shift_mlp, scale_mlp)\n", " x = x + gate_mlp.unsqueeze(1) * MLP(hidden_size)(h)\n", " \n", " return x\n", " \n", "FinalLayer(hidden_size, patch_size, out_channels)\n", " Given x (B x L x D), c (B x D)\n", " c = SiLU()(c)\n", " c = Linear(hidden_size, 2 * hidden_size)(c)\n", " shift, scale = c.chunk(2, dim=1)\n", " x = LayerNorm(hidden_size, elementwise_affine=False)(x)\n", " x = modulate(x, shift, scale)\n", " x = Linear(hidden_size, patch_size * patch_size * out_channels)(x)\n", " return x\n", " \n", "DiT(input_shape, patch_size, hidden_size, num_heads, num_layers, num_classes, cfg_dropout_prob)\n", " Given x (B x C x H x W) - image, y (B) - class label, t (B) - diffusion timestep\n", " x = patchify_flatten(x) # B x C x H x W -> B x (H // P * W // P) x D, P is patch_size\n", " x += pos_embed # see get_2d_sincos_pos_embed\n", " \n", " t = compute_timestep_embedding(t) # Same as in UNet\n", " if training:\n", " y = dropout_classes(y, cfg_dropout_prob) # Randomly dropout to train unconditional image generation\n", " y = Embedding(num_classes + 1, hidden_size)(y)\n", " c = t + y\n", " \n", " for _ in range(num_layers):\n", " x = DiTBlock(hidden_size, num_heads)(x, c)\n", " \n", " x = FinalLayer(hidden_size, patch_size, out_channels)(x)\n", " x = unpatchify(x) # B x (H // P * W // P) x (P * P * C) -> B x C x H x W\n", " return x\n", "```" ] }, { "cell_type": "markdown", "id": "b3b8c7ff", "metadata": { "id": "b3b8c7ff" }, "source": [ "**Hyperparameter details**\n", "* Normalize image to [-1, 1], (2) Encode using the VAE, (3) divide latents by the scale_factor compute in part (a)\n", "* Transformer with patch_size 2, hidden_size 512, num_heads 8, num_layers 12\n", "* Train 60 epochs, batch size 256, Adam with LR 1e-3 (100 warmup steps, cosine decay to 0)\n", "* When sampling, remember to multiple the final generated latents by the scale_factor before feeding it through the decoder\n", "* For diffusion schedule, sampling and loss, use the same setup as Q1\n", "\n", "For class conditioning, learn an embedding for each class, and an extra embedding to represent the null class. To condition, add the class embedding to the timestep embedding before feeding it into the transformer blocks (see pseudocode). **Train your class conditional diffusion models while dropping out the class (replace with null class) 10% of the time. This will be necessary for part (c).**\n", "\n", "**Remember to save your model parameters after training, as you will need them for part (c)**" ] }, { "cell_type": "code", "execution_count": null, "id": "891156e0", "metadata": { "id": "891156e0" }, "outputs": [], "source": [ "def q3_b(train_data, train_labels, test_data, test_labels, vae):\n", " \"\"\"\n", " train_data: A (50000, 32, 32, 3) numpy array of images in [0, 1]\n", " train_labels: A (50000,) numpy array of class labels\n", " test_data: A (10000, 32, 32, 3) numpy array of images in [0, 1]\n", " test_labels: A (10000,) numpy array of class labels\n", " vae: a pretrained VAE\n", "\n", " Returns\n", " - a (# of training iterations,) numpy array of train losses evaluated every minibatch\n", " - a (# of num_epochs + 1,) numpy array of test losses evaluated at the start of training and the end of every epoch\n", " - a numpy array of size (10, 10, 32, 32, 3) of samples in [0, 1] drawn from your model.\n", " The array represents a 10 x 10 grid of generated samples. Each row represents 10 samples generated\n", " for a specific class (i.e. row 0 is class 0, row 1 class 1, ...). Use 512 diffusion timesteps\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"\n", "\n", " return train_losses, test_losses, samples" ] }, { "cell_type": "code", "execution_count": null, "id": "b556f1a4", "metadata": { "id": "b556f1a4" }, "outputs": [], "source": [ "q3b_save_results(q3_b)" ] }, { "cell_type": "markdown", "id": "3e228cf4", "metadata": { "id": "3e228cf4" }, "source": [ "## Part 3(c) Classifier-Free Guidance [20pt]\n", "In this part, you will implement [Classifier-Free Guidance](https://arxiv.org/abs/2207.12598) (CFG). CFG is a widely used method during diffusion model sampling to push samples towards more accurately aligning with the conditioning information (e.g. class, text caption).\n", "\n", "Implement CFG requires a small modification to the diffusion sampling code. Given a CIFAR-10 class label, instead of using $\\hat{\\epsilon} = f_\\theta(x_t, t, y)$ to sample, use:\n", "$$\\hat{\\epsilon} = f_\\theta(x_t, t, \\varnothing) + w(f_\\theta(x_t, t, y) - f_\\theta(x_t, t, \\varnothing))$$\n", "where $w$ is a sampling hyperparameter that controls the strength of CFG. $\\varnothing$ indicates the unconditional model with the class label dropped out, which your pre-trained UNet from 3(b) should support. Note that $w = 1$ recovers standard sampling.\n", "\n", "Note: It may be expected to see worse samples (e.g. sautrated images) when CFG value is too high. Generation quality is closer to a U-shape when increasing CFG values (gets better, then worse)" ] }, { "cell_type": "code", "execution_count": null, "id": "eba62b5d", "metadata": { "id": "eba62b5d" }, "outputs": [], "source": [ "def q3_c(vae):\n", " \"\"\"\n", " vae: a pretrained vae\n", "\n", " Returns\n", " - a numpy array of size (4, 10, 10, 32, 32, 3) of samples in [0, 1] drawn from your model.\n", " The array represents a 4 x 10 x 10 grid of generated samples - 4 10 x 10 grid of samples\n", " with 4 different CFG values of w = {1.0, 3.0, 5.0, 7.5}. Each row of the 10 x 10 grid\n", " should contain samples of a different class. Use 512 diffusion sampling timesteps.\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"\n", "\n", " return samples" ] }, { "cell_type": "code", "execution_count": null, "id": "683d63c7", "metadata": { "id": "683d63c7" }, "outputs": [], "source": [ "q3c_save_results(q3_c)" ] }, { "cell_type": "code", "execution_count": null, "id": "6998ef15-6493-4aeb-a9b3-cd70f9c2b3fa", "metadata": { "id": "6998ef15-6493-4aeb-a9b3-cd70f9c2b3fa" }, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "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.9.18" }, "colab": { "provenance": [], "include_colab_link": true } }, "nbformat": 4, "nbformat_minor": 5 }