{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "view-in-github" }, "source": [ "\"Open" ] }, { "cell_type": "markdown", "metadata": { "id": "dEcSNKhrotPo" }, "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 various autoencoder models and run them on two datasets (dataset 1 and dataset 2). The expected outputs for dataset 1 are already provided to help as a sanity check.\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 image 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. It may take longer to run since we are using larger datasets." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "ER_o_stMdlIy", "outputId": "dc12a173-5428-4d55-e0ce-ac3411b21cee" }, "outputs": [], "source": [ "!if [ -d deepul ]; then rm -Rf deepul; fi\n", "!git clone https://github.com/rll/deepul.git\n", "# This link may break due to google changing how files are linked. You can manually download from https://drive.google.com/drive/folders/1BOLGo3CFDa_qJMLRgtTZ6zrAWezk9oJu?usp=sharing \n", "!curl \"https://drive.usercontent.google.com/download?id=1lWjGICwgzgcBDejo9S5g69hLAf0O3lGF&confirm=xxx\" -o hw2_data.zip\n", "!unzip -qq hw2_data.zip -d deepul/homeworks/hw2/data\n", "!pip install ./deepul\n", "from deepul.hw2_helper import *" ] }, { "cell_type": "markdown", "metadata": { "id": "IMx4IKP7dtNT" }, "source": [ "# Question 1: VAEs on 2D Data [20pts]\n", "In this question, you will train a simple VAE on 2D data, and look at situations in which latents are being used or not being used (i.e. when posterior collapse occurs)" ] }, { "cell_type": "markdown", "metadata": { "id": "2RTrhLn0Ofyd" }, "source": [ "## Part (a) Data from a Full Covariance Gaussian [10 pts]\n", "In this part, we train a VAE on data generated from a Gaussian with a full covariance matrix. Execute the cell below to visualize the datasets." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 922 }, "id": "IKdPunNzdvam", "outputId": "28777a58-2d2b-41ed-9b95-96ce4afda8ab" }, "outputs": [], "source": [ "visualize_q1_data('a', 1)\n", "visualize_q1_data('a', 2)" ] }, { "cell_type": "markdown", "metadata": { "id": "O5X3VHQ2rvCn" }, "source": [ "Consruct and train a VAE with the following characteristics\n", "* 2D latent variables $z$ with a standard normal prior, $p(z) = N(0, I)$\n", "* An approximate posterior $q_\\theta(z|x) = N(z; \\mu_\\theta(x), \\Sigma_\\theta(x))$, where $\\mu_\\theta(x)$ is the mean vector, and $\\Sigma_\\theta(x)$ is a diagonal covariance matrix\n", "* A decoder $p(x|z) = N(x; \\mu_\\phi(z), \\Sigma_\\phi(z))$, where $\\mu_\\phi(z)$ is the mean vector, and $\\Sigma_\\phi(z)$ is a diagonal covariance matrix\n", "\n", "**You will provide the following deliverables**\n", "\n", "\n", "1. Over the course of training, record the average full negative ELBO, reconstruction loss $E_xE_{z\\sim q(z|x)}[-p(x|z)]$, and KL term $E_x[D_{KL}(q(z|x)||p(z))]$ of the training data (per minibatch) and test data (for your entire test set). Code is provided that automatically plots the training curves.\n", "2. Report the final test set performance of your final model\n", "3. Samples of your trained VAE with ($z\\sim p(z), x\\sim N(x;\\mu_\\phi(z),\\Sigma_\\phi(z))$) and without ($z\\sim p(z), x = \\mu_\\phi(z)$) decoder noise\n", "\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "D3qice-vN65o" }, "source": [ "### Solution\n", "Fill out the function below and return the neccessary arguments. Feel free to create more cells if need be" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "MOYOUBMRrwTz" }, "outputs": [], "source": [ "def q1(train_data, test_data, part, dset_id):\n", " \"\"\"\n", " train_data: An (n_train, 2) numpy array of floats\n", " test_data: An (n_test, 2) numpy array of floats\n", "\n", " (You probably won't need to use the two inputs below, but they are there\n", " if you want to use them)\n", " part: An identifying string ('a' or 'b') of which part is being run. Most likely\n", " used to set different hyperparameters for different datasets\n", " dset_id: An identifying number of which dataset is given (1 or 2). Most likely\n", " used to set different hyperparameters for different datasets\n", "\n", " Returns\n", " - a (# of training iterations, 3) numpy array of full negative ELBO, reconstruction loss E[-p(x|z)],\n", " and KL term E[KL(q(z|x) | p(z))] evaluated every minibatch\n", " - a (# of epochs + 1, 3) numpy array of full negative ELBO, reconstruciton loss E[-p(x|z)],\n", " and KL term E[KL(q(z|x) | p(z))] evaluated once at initialization and after each epoch\n", " - a numpy array of size (1000, 2) of 1000 samples WITH decoder noise, i.e. sample z ~ p(z), x ~ p(x|z)\n", " - a numpy array of size (1000, 2) of 1000 samples WITHOUT decoder noise, i.e. sample z ~ p(z), x = mu(z)\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"" ] }, { "cell_type": "markdown", "metadata": { "id": "x_vb7Fn8ODiW" }, "source": [ "### Results\n", "Once you've finished `q1_a`, execute the cells below to visualize and save your results." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "O5oK-YZFOQEq", "outputId": "61999a87-5c2b-42d2-f84c-e0636421612a" }, "outputs": [], "source": [ "q1_save_results('a', 1, q1)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "fCdsKGqVi79i", "outputId": "7de3427e-f7a8-495f-c65e-b192924bf721" }, "outputs": [], "source": [ "q1_save_results('a', 2, q1)" ] }, { "cell_type": "markdown", "metadata": { "id": "5sEbxPOzQuWB" }, "source": [ "## Part (b) Data from a Diagonal Gaussian [10pts]\n", "In this part, we use your code from the previous part to train a VAE on data generated from a diagonal gaussian. Execute the cell below to visualize the datasets (note that they may look the same, but notice the axes)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 922 }, "id": "IQzLsgCYRDXz", "outputId": "63667912-9b8b-49eb-eee5-00fa766ac0f5" }, "outputs": [], "source": [ "visualize_q1_data('b', 1)\n", "visualize_q1_data('b', 2)" ] }, { "cell_type": "markdown", "metadata": { "id": "2t3dPbPsQ-ZI" }, "source": [ "### Results\n", "Execute the cells below to visualize and save your results." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "6aF_dZ-kRL8J", "outputId": "dfdd8b26-cb53-485a-e0cd-c6ad5956bc11" }, "outputs": [], "source": [ "q1_save_results('b', 1, q1)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "u1DMln9M1mPW", "outputId": "412aa03d-2c61-45ac-fc00-a7254a745da6" }, "outputs": [], "source": [ "q1_save_results('b', 2, q1)" ] }, { "cell_type": "markdown", "metadata": { "id": "mg22J90qROiO" }, "source": [ "### Reflection\n", "Compare the sampled xs with and without latents in parts (a) and (b). For which datasets are the latents being used? Why is this happening (i.e. why are the latents being ignored in some cases)? **Write your answer (1-2 sentences) in the given latex template**" ] }, { "cell_type": "markdown", "metadata": { "id": "xAAixSJ1dv7u" }, "source": [ "# Question 2: VAEs on Images [40pts]\n", "In this question, you will train different VAE models on image datasets. Execute the cell below to visualize the two datasets (colored shapes, and [SVHN](http://ufldl.stanford.edu/housenumbers/))." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 839 }, "id": "gj2CDM5bXBTG", "outputId": "a75eeb17-b26a-434d-be97-fa4aeac3a38b" }, "outputs": [], "source": [ "visualize_svhn()\n", "visualize_cifar10()" ] }, { "cell_type": "markdown", "metadata": { "id": "8aX115gIuMwB" }, "source": [ "## Part (a) VAE [20pts]\n", "In this part, implement a standard VAE with the following characteristics:\n", "\n", "* 16-dim latent variables $z$ with standard normal prior $p(z) = N(0,I)$\n", "* An approximate posterior $q_\\theta(z|x) = N(z; \\mu_\\theta(x), \\Sigma_\\theta(x))$, where $\\mu_\\theta(x)$ is the mean vector, and $\\Sigma_\\theta(x)$ is a diagonal covariance matrix\n", "* A decoder $p(x|z) = N(x; \\mu_\\phi(z), I)$, where $\\mu_\\phi(z)$ is the mean vector. (We are not learning the covariance of the decoder)\n", "\n", "You can play around with different architectures and try for better results, but the following encoder / decoder architecture below suffices (Note that image input is always $32\\times 32$.\n", "```\n", "conv2d(in_channels, out_channels, kernel_size, stride, padding)\n", "transpose_conv2d(in_channels, out_channels, kernel_size, stride, padding)\n", "linear(in_dim, out_dim)\n", "\n", "Encoder\n", " conv2d(3, 32, 3, 1, 1)\n", " relu()\n", " conv2d(32, 64, 3, 2, 1) # 16 x 16\n", " relu()\n", " conv2d(64, 128, 3, 2, 1) # 8 x 8\n", " relu()\n", " conv2d(128, 256, 3, 2, 1) # 4 x 4\n", " relu()\n", " flatten() # 16\n", " linear(4 * 4 * 256, 2 * latent_dim)\n", "\n", "Decoder\n", " linear(latent_dim, 4 * 4 * 128)\n", " relu()\n", " reshape(4, 4, 128)\n", " transpose_conv2d(128, 128, 4, 2, 1) # 8 x 8\n", " relu()\n", " transpose_conv2d(128, 64, 4, 2, 1) # 16 x 16\n", " relu()\n", " transpose_conv2d(64, 32, 4, 2, 1) # 32 x 32\n", " relu()\n", " conv2d(32, 3, 3, 1, 1)\n", "```\n", "\n", "You may find the following training tips helpful\n", "* When computing reconstruction loss and KL loss, average over the batch dimension and **sum** over the feature dimension\n", "* When computing reconstruction loss, it suffices to just compute MSE between the reconstructed $x$ and true $x$\n", "* Use batch size 128, learning rate $10^{-3}$, and an Adam optimizer\n", "\n", "\n", "\n", "\n", "**You will provide the following deliverables**\n", "\n", "\n", "1. Over the course of training, record the average full negative ELBO, reconstruction loss, and KL term of the training data (per minibatch) and test data (for your entire test set). Code is provided that automatically plots the training curves.\n", "2. Report the final test set performance of your final model\n", "3. 100 samples from your trained VAE\n", "4. 50 real-image / reconstruction pairs (for some $x$, encode and then decode)\n", "5. 10 interpolations of 10 images from your trained VAE (100 images total)\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "7_pYz61AfW4U" }, "source": [ "### Solution\n", "Fill out the function below and return the neccessary arguments. Feel free to create more cells if need be" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "LIW7tqSMdwg1" }, "outputs": [], "source": [ "def q2_a(train_data, test_data, dset_id):\n", " \"\"\"\n", " train_data: An (n_train, 32, 32, 3) uint8 numpy array of color images with values in {0, ..., 255}\n", " test_data: An (n_test, 32, 32, 3) uint8 numpy array of color images with values in {0, ..., 255}\n", " dset_id: An identifying number of which dataset is given (1 or 2). Most likely\n", " used to set different hyperparameters for different datasets\n", "\n", " Returns\n", " - a (# of training iterations, 3) numpy array of full negative ELBO, reconstruction loss E[-p(x|z)],\n", " and KL term E[KL(q(z|x) | p(z))] evaluated every minibatch\n", " - a (# of epochs + 1, 3) numpy array of full negative ELBO, reconstruciton loss E[-p(x|z)],\n", " and KL term E[KL(q(z|x) | p(z))] evaluated once at initialization and after each epoch\n", " - a (100, 32, 32, 3) numpy array of 100 samples from your VAE with values in {0, ..., 255}\n", " - a (100, 32, 32, 3) numpy array of 50 real image / reconstruction pairs\n", " FROM THE TEST SET with values in {0, ..., 255}\n", " - a (100, 32, 32, 3) numpy array of 10 interpolations of length 10 between\n", " pairs of test images. The output should be those 100 images flattened into\n", " the specified shape with values in {0, ..., 255}\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"" ] }, { "cell_type": "markdown", "metadata": { "id": "7QJvrGXyfnj0" }, "source": [ "### Results\n", "Once you've finished `q2_a`, execute the cells below to visualize and save your results." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "wL7s92ynfqil", "outputId": "6d28b6b9-7e52-464b-a780-280eb59a9f59" }, "outputs": [], "source": [ "q2_save_results('a', 1, q2_a)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 321 }, "id": "DCre0Qso1jcV", "outputId": "4d13f1e1-d466-46c4-f095-98f91751cfb9" }, "outputs": [], "source": [ "q2_save_results('a', 2, q2_a)" ] }, { "cell_type": "markdown", "metadata": { "id": "822YbCb2fsz8" }, "source": [ "## Part (b) Hierarchical VAE [20pts]\n", "\n", "In this part, we will explore a simplified version of the hierarchical VAE described in [NVAE](https://arxiv.org/pdf/2007.03898.pdf). We will not implement the full NVAE, but rather use some ideas from the paper to explore how to learn a prior distribution p(z).\n", "\n", "Implement a hierarchical VAE that follows the following structure.\n", "* $z1$ is a 2x2x12 latent vector where p(z1) is the unit Gaussian.\n", " * Learn the approximate posterior $q_\\theta(z|x) = N(z; \\mu_\\theta(x), \\Sigma_\\theta(x))$, where $\\mu_\\theta(x)$ is the mean vector, and $\\Sigma_\\theta(x)$ is a diagonal covariance matrix. I.e., same as a normal VAE, but use a matrix latent rather than a vector. Each dimension is independent.\n", "* $z2$ is a 2x2x12 latent vector.\n", " * $p_\\theta(z2|z1)$ is learned, and implemented as a neural network that parameterizes mean (and log std, optionally).\n", " * $q_\\theta(z2|z1,x)$ is also learned. Implement this as a Residual Normal [see NVAE] over the prior $p_\\theta(z2|z1)$.\n", "* The decoder should be a function of $z2$ only.\n", "\n", "Some helpful hints:\n", "* Two KL losses should be calculated. The first should match $q_\\theta(z|x)$ to the unit Gaussian. The second should match $q_\\theta(z2|z1,x)$ and $p_\\theta(z2|z1)$, and be taken with respect to $q$.\n", "* When calculating the second KL term, utilize the analytic form for the residual normal. When $q_\\theta(z2|z1,x) = N(z2; \\mu_\\theta(z1) + \\Delta \\mu_\\theta(z1,x), \\Sigma_\\theta(z1)) * \\Delta \\Sigma_\\theta(z1,x))$, use the following form: `kl_z2 = -z2_residual_logstd - 0.5 + (torch.exp(2 * z2_residual_logstd) + z2_residual_mu ** 2) * 0.5`\n", "* When calculating KL, remember to sum over the dimensions of the latent variable before taking the mean over batch.\n", "* For the prior $p_\\theta(z2|z1)$, fix standard deviation to be 1. Learn only the mean. This will help with stability in training.\n", "\n", "The following network structures may be useful:\n", "```\n", "conv2d(in_channels, out_channels, kernel_size, stride, padding)\n", "transpose_conv2d(in_channels, out_channels, kernel_size, stride, padding)\n", "linear(in_dim, out_dim)\n", "\n", "Encoder\n", " nn.Conv2d(3 + 12, 32, 3, padding=1), # [32, 32, 32]\n", " LayerNorm(32),\n", " nn.ReLU(),\n", " nn.Conv2d(32, 64, 3, stride=2, padding=1), # [64, 16, 16]\n", " LayerNorm(64),\n", " nn.ReLU(),\n", " nn.Conv2d(64, 64, 3, stride=2, padding=1), # [64, 8, 8]\n", " LayerNorm(64),\n", " nn.ReLU(),\n", " nn.Conv2d(64, 64, 3, stride=2, padding=1), # [64, 4, 4]\n", " LayerNorm(64),\n", " nn.ReLU(),\n", " nn.Conv2d(64, 64, 3, stride=2, padding=1), # [64, 2, 2]\n", " LayerNorm(64),\n", " nn.ReLU(),\n", " nn.Conv2d(64, 12*2, 3, padding=1), # [12*2, 2, 2]\n", "We assume encoder networks are of the form p(z'|z,x).\n", "When learning q(z1), an x of all zeros can be used as input.\n", "Upscale z with nearest-neighbor projection before concatenating with x.\n", "\n", "\n", "Decoder\n", " nn.ConvTranspose2d(12, 64, 3, padding=1), # [64, 2, 2]\n", " nn.ReLU(),\n", " nn.ConvTranspose2d(64, 64, 4, stride=2, padding=1), # [64, 4, 4]\n", " nn.ReLU(),\n", " nn.ConvTranspose2d(64, 64, 4, stride=2, padding=1), # [64, 8, 8]\n", " nn.ReLU(),\n", " nn.ConvTranspose2d(64, 64, 4, stride=2, padding=1), # [64, 16, 16]\n", " nn.ReLU(),\n", " nn.ConvTranspose2d(64, 32, 4, stride=2, padding=1), # [32, 32, 32]\n", " nn.ReLU(),\n", " nn.Conv2d(32, 3, 3, padding=1), # [3, 32, 32]\n", "```\n", "\n", "\n", "**You will provide the following deliverables**\n", "\n", "\n", "1. Over the course of training, record the average full negative ELBO, reconstruction loss, and KL term of the training data (per minibatch) and test data (for your entire test set). Code is provided that automatically plots the training curves.\n", "2. Report the final test set performance of your final model\n", "3. 100 samples from your trained VAE\n", "4. 50 real-image / reconstruction pairs (for some $x$, encode and then decode)\n", "5. 10 interpolations of 10 images from your trained VAE (100 images total)\n", "\n", "\n" ] }, { "cell_type": "markdown", "metadata": { "id": "xwT1tOdm0e84" }, "source": [ "### Solution\n", "Fill out the function below and return the neccessary arguments. Feel free to create more cells if need be" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "CpqhOqW1UDby" }, "outputs": [], "source": [ "def q2_b(train_data, test_data, dset_id):\n", " \"\"\"\n", " train_data: An (n_train, 32, 32, 3) uint8 numpy array of color images with values in {0, ..., 255}\n", " test_data: An (n_test, 32, 32, 3) uint8 numpy array of color images with values in {0, ..., 255}\n", " dset_id: An identifying number of which dataset is given (1 or 2). Most likely\n", " used to set different hyperparameters for different datasets\n", "\n", " Returns\n", " - a (# of training iterations, 3) numpy array of full negative ELBO, reconstruction loss E[-p(x|z)],\n", " and KL term E[KL(q(z|x) | p(z))] evaluated every minibatch\n", " - a (# of epochs + 1, 3) numpy array of full negative ELBO, reconstruciton loss E[-p(x|z)],\n", " and KL term E[KL(q(z|x) | p(z))] evaluated once at initialization and after each epoch\n", " - a (100, 32, 32, 3) numpy array of 100 samples from your VAE with values in {0, ..., 255}\n", " - a (100, 32, 32, 3) numpy array of 50 real image / reconstruction pairs\n", " FROM THE TEST SET with values in {0, ..., 255}\n", " - a (100, 32, 32, 3) numpy array of 10 interpolations of length 10 between\n", " pairs of test images. The output should be those 100 images flattened into\n", " the specified shape with values in {0, ..., 255}\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"" ] }, { "cell_type": "markdown", "metadata": { "id": "_qvP94pm0hYb" }, "source": [ "### Results\n", "Once you've finished `q2_b`, execute the cells below to visualize and save your results." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "SPtDMRpf0iAG", "outputId": "3a31f0be-432d-41d2-d568-a45c4cb49ed6" }, "outputs": [], "source": [ "q2_save_results('b', 1, q2_b)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "QOdX77hP1HMv", "outputId": "4e3d95f3-3497-47d1-da53-077da3ec2bed" }, "outputs": [], "source": [ "q2_save_results('b', 2, q2_b)" ] }, { "cell_type": "markdown", "metadata": { "id": "P-EZD8GCdx0B" }, "source": [ "# Question 3: VQ-VAE [40pts]\n", "In this question, you with train a [VQ-VAE](https://arxiv.org/abs/1711.00937) on the colored shapes dataset and SVHN. If you are confused on how the VQ-VAE works, you may find [Lilian Weng's blogpost](https://lilianweng.github.io/lil-log/2018/08/12/from-autoencoder-to-beta-vae.html#vq-vae-and-vq-vae-2) to be useful.\n", "\n", "You may experiment with different hyperparameters and architecture designs, but the following designs for the VQ-VAE architecture may be useful.\n", "\n", "```\n", "conv2d(in_channels, out_channels, kernel_size, stride, padding)\n", "transpose_conv2d(in_channels, out_channels, kernel_size, stride, padding)\n", "linear(in_dim, out_dim)\n", "batch_norm2d(dim)\n", "\n", "residual_block(dim)\n", " batch_norm2d(dim)\n", " relu()\n", " conv2d(dim, dim, 3, 1, 1)\n", " batch_norm2d(dim)\n", " relu()\n", " conv2d(dim, dim, 1, 1, 0)\n", "\n", "Encoder\n", " conv2d(3, 256, 4, 2, 1) 16 x 16\n", " batch_norm2d(256)\n", " relu()\n", " conv2d(256, 256, 4, 2, 1) 8 x 8\n", " residual_block(256)\n", " residual_block(256)\n", "\n", "Decoder\n", " residual_block(256)\n", " residual_block(256)\n", " batch_norm2d(256)\n", " relu()\n", " transpose_conv2d(256, 256, 4, 2, 1) 16 x 16\n", " batch_norm2d(256)\n", " relu()\n", " transpose_conv2d(256, 3, 4, 2, 1) 32 x 32\n", "```\n", "\n", "A few other tips:\n", "* Use a codebook with $K = 128$ latents each with a $D = 256$ dimensional embedding vector\n", "* You should initialize each element in your $K\\times D$ codebook to be uniformly random in $[-1/K, 1/K]$\n", "* Use batch size 128 with a learning rate of $10^{-3}$ and an Adam optimizer\n", "* Center and scale your images to $[-1, 1]$\n", "* Supposing that $z_e(x)$ is the encoder output, and $z_q(x)$ is the quantized output using the codebook, you can implement the straight-through estimator as follows (where below is fed into the decoder):\n", " * `(z_q(x) - z_e(x)).detach() + z_e(x)` in Pytorch\n", " * `tf.stop_gradient(z_q(x) - z_e(x)) + z_e(x)` in Tensorflow.\n", "\n", "In addition to training the VQ-VAE, you will also need to train a Transformer prior on the categorical latents in order to sample. Feel free to use your implementation for HW1! You should flatten the VQ-VAE tokens into a [H x W] sequence, and use a start token.\n", "\n", "**You will provide the following deliverables**\n", "\n", "\n", "1. Over the course of training, record the average loss of the training data (per minibatch) and test data (for your entire test set) **for both your VQ-VAE and Transformer prior**. Code is provided that automatically plots the training curves.\n", "2. Report the final test set performances of your final models\n", "3. 100 samples from your trained VQ-VAE and Transformer prior\n", "4. 50 real-image / reconstruction pairs (for some $x$, encode and then decode)" ] }, { "cell_type": "markdown", "metadata": { "id": "vHZsMrEw5wLN" }, "source": [ "## Solution\n", "Fill out the function below and return the neccessary arguments. Feel free to create more cells if need be" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "cUQ2V2hLdyUF" }, "outputs": [], "source": [ "def q3(train_data, test_data, dset_id):\n", " \"\"\"\n", " train_data: An (n_train, 32, 32, 3) uint8 numpy array of color images with values in [0, 255]\n", " test_data: An (n_test, 32, 32, 3) uint8 numpy array of color images with values in [0, 255]\n", " dset_id: An identifying number of which dataset is given (1 or 2). Most likely\n", " used to set different hyperparameters for different datasets\n", "\n", " Returns\n", " - a (# of training iterations,) numpy array of VQ-VAE train losess evaluated every minibatch\n", " - a (# of epochs + 1,) numpy array of VQ-VAE test losses evaluated once at initialization and after each epoch\n", " - a (# of training iterations,) numpy array of Transformer prior train losess evaluated every minibatch\n", " - a (# of epochs + 1,) numpy array of Transformer prior test losses evaluated once at initialization and after each epoch\n", " - a (100, 32, 32, 3) numpy array of 100 samples with values in {0, ... 255}\n", " - a (100, 32, 32, 3) numpy array of 50 real image / reconstruction pairs\n", " FROM THE TEST SET with values in [0, 255]\n", " \"\"\"\n", "\n", " \"\"\" YOUR CODE HERE \"\"\"" ] }, { "cell_type": "markdown", "metadata": { "id": "Nbn-r53G51X_" }, "source": [ "## Results\n", "Once you've finished `q3`, execute the cells below to visualize and save your results." ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 1000 }, "id": "ClKjwiAd535z", "outputId": "dc4baa34-e868-4e32-f35d-937f3aab2b9e" }, "outputs": [], "source": [ "q3_save_results(1, q3)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "5vDEvml-59zA" }, "outputs": [], "source": [ "q3_save_results(2, q3)" ] }, { "cell_type": "markdown", "metadata": { "id": "pOS0PKRKdtLS" }, "source": [ "# File Saving Utility\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "7lMqTnFfdur6" }, "outputs": [], "source": [ "# one way to save files\n", "from google.colab import files\n", "import os\n", "import os.path as osp\n", "for fname in os.listdir('results'):\n", " files.download('results/' + fname)" ] } ], "metadata": { "accelerator": "GPU", "colab": { "collapsed_sections": [ "16l6wsRW4k8d", "pOS0PKRKdtLS" ], "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" } }, "nbformat": 4, "nbformat_minor": 0 }