🎉 Introduction to Variational Autoencoders (VAEs) with PyTorch

pip install pytorch-pretrained-biggan

pip install Pillow libsixel-python

import torch
import random

import numpy as np
import matplotlib.pylab as plt

import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

import torchvision
from torchvision import datasets, transforms

pip install matplotlib tqdm

pip install scipy

from pytorch_pretrained_biggan import one_hot_from_names
from tqdm.notebook import tqdm, trange

# @title Figure settings
import logging
logging.getLogger('matplotlib.font_manager').disabled = True

import ipywidgets as widgets
from ipywidgets import FloatSlider, IntSlider, HBox, Layout, VBox
from ipywidgets import interactive_output, Dropdown

%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/content-creation/main/nma.mplstyle")

!pip3 install vibecheck datatops

from vibecheck import DatatopsContentReviewContainer
def content_review(notebook_section: str):
    return DatatopsContentReviewContainer(
        "",  # No text prompt
        notebook_section,
        {
            "url": "https://pmyvdlilci.execute-api.us-east-1.amazonaws.com/klab",
            "name": "neuromatch_dl",
            "user_key": "f379rz8y",
        },
    ).render()


feedback_prefix = "W2D4_T1"

# Helper functions

DEVICE = torch.device("cuda:0")

def image_moments(image_batches, n_batches=None):
    """
    Compute mean and covariance of all pixels from batches of images

    Args:
        Image_batches: tuple
            Image batches
        n_batches: int
            Number of Batch size

    Returns:
        m1: float
            Mean of all pixels
        cov: float
            Covariance of all pixels
    """
    m1, m2 = torch.zeros((), device=DEVICE), torch.zeros((), device=DEVICE)
    n = 0
    for im in tqdm(image_batches, total=n_batches, leave=False, desc="Computing pixel mean and covariance"):
        im = im.to(DEVICE)
        b = im.size()[0]
        im = im.view(b, -1)
        m1 = m1 + im.sum(dim=0)
        m2 = m2 + (im.view(b, -1, 1) * im.view(b, 1, -1)).sum(dim=0)
        n += b
    m1, m2 = m1/n, m2/n
    cov = m2 - m1.view(-1, 1) * m1.view(1, -1)
    return m1.cpu(), cov.cpu()

def interpolate(A, B, num_interps):
    """
    Function to interpolate between images.
    It does this by linearly interpolating between the
    probability of each category you select and linearly
    interpolating between the latent vector values.

    Args:
        A: list
            List of categories
        B: list
            List of categories
        num_interps: int
            Quantity of pixel grids

    Returns:
        Interpolated np.ndarray
    """
    if A.shape != B.shape:
        raise ValueError("A and B must have the same shape")
    alphas = np.linspace(0, 1, num_interps)
    return np.array([(1-a) * A + a * B for a in alphas])

def kl_q_p(zs, phi):
    """
    Compute the KL divergence KL(q||p).

    Args:
        zs: torch.tensor
            Samples z drawn from q, shape [b, n, k].
        phi: torch.tensor
            Parameters of q, shape [b, k+1].

    Returns:
        torch.tensor
            KL divergence estimate.
    """
    b, n, k = zs.size()
    mu_q, log_sig_q = phi[:, :-1], phi[:, -1]
    
    # Compute log probabilities for p(z) ~ N(0, 1)
    log_p = -0.5 * (zs**2)

    # Compute log probabilities for q(z|phi)
    log_q = -0.5 * ((zs - mu_q.view(b, 1, k))**2 / log_sig_q.exp().view(b, 1, 1)**2) - log_sig_q.view(b, 1, -1)

    # Compute KL divergence
    return (log_q - log_p).sum(dim=-1).mean()

def log_p_x(x, mu_xs, sig_x):
    """
    Given [batch, ...] input x and [batch, n, ...] reconstructions, compute
    pixel-wise log Gaussian probability Sum over pixel dimensions, but mean over batch and samples.

    Args:
        x: np.ndarray
            Input Data
        mu_xs: np.ndarray
            Log of mean of samples
        sig_x: np.ndarray
            Log of standard deviation

    Returns:
        Mean over batch and samples.
    """
    b, n = mu_xs.size()[:2]
    # Flatten out pixels and add a singleton dimension [1] so that x will be
    # implicitly expanded when combined with mu_xs
    x = x.reshape(b, 1, -1)
    _, _, p = x.size()
    squared_error = (x - mu_xs.view(b, n, -1))**2 / (2 * sig_x**2)

    # Size of squared_error is [b, n, p]. log prob is by definition sum over [p].
    # Expected value requires mean over [n].
    # Handling different size batches requires mean over [b].
    return -(squared_error + torch.log(sig_x)).sum(dim=2).mean(dim=(0, 1))

def pca_encoder_decoder(mu, cov, k):
    mu = mu.view(1, -1)
    u, s, v = torch.svd_lowrank(cov, q=k)
    W_encode = v / torch.sqrt(s)
    W_decode = u * torch.sqrt(s)

    def pca_encode(x):
        return (x.view(-1, mu.numel()) - mu) @ W_encode

    def pca_decode(h):
        return (h @ W_decode.T) + mu

    return pca_encode, pca_decode

def cout(x, layer):
    """
    Unnecessarily complicated but complete way to
    calculate the output depth, height and width size for a Conv2D layer

    Args:
        x: tuple
            Input size (depth, height, width)
        layer: nn.Conv2d
            The Conv2D layer

    Returns:
        Tuple of out-depth/out-height and out-width 
        Output shape as given in [Ref]
        Ref: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
    """
    assert isinstance(layer, nn.Conv2d)
    p = layer.padding if isinstance(layer.padding, tuple) else (layer.padding, )
    k = layer.kernel_size if isinstance(layer.kernel_size, tuple) else (layer.kernel_size, )
    d = layer.dilation if isinstance(layer.dilation, tuple) else (layer.dilation, )
    s = layer.stride if isinstance(layer.stride, tuple) else (layer.stride, )
    in_depth, in_height, in_width = x
    out_depth = layer.out_channels
    out_height = 1 + (in_height + 2 * p[0] - (k[0] - 1) * d[0] - 1) // s[0]
    out_width = 1 + (in_width + 2 * p[-1] - (k[-1] - 1) * d[-1] - 1) // s[-1]
    return (out_depth, out_height, out_width)
    

def cout(x, layer):
    """
    Unnecessarily complicated but complete way to
    calculate the output depth, height and width size for a Conv2D layer

    Args:
        x: tuple
            Input size (depth, height, width)
        layer: nn.Conv2d
            The Conv2D layer

    Returns:
        Tuple of out-depth/out-height and out-width 
        Output shape as given in [Ref]
        Ref: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
    """
    assert isinstance(layer, nn.Conv2d)
    p = layer.padding if isinstance(layer.padding, tuple) else (layer.padding, )
    k = layer.kernel_size if isinstance(layer.kernel_size, tuple) else (layer.kernel_size, )
    d = layer.dilation if isinstance(layer.dilation, tuple) else (layer.dilation, )
    s = layer.stride if isinstance(layer.stride, tuple) else (layer.stride, )
    in_depth, in_height, in_width = x
    out_depth = layer.out_channels
    out_height = 1 + (in_height + 2 * p[0] - (k[0] - 1) * d[0] - 1) // s[0]
    out_width = 1 + (in_width + 2 * p[-1] - (k[-1] - 1) * d[-1] - 1) // s[-1]
    return (out_depth, out_height, out_width)
    

def cout(x, layer):
    """
    Unnecessarily complicated but complete way to
    calculate the output depth, height and width size for a Conv2D layer

    Args:
        x: tuple
            Input size (depth, height, width)
        layer: nn.Conv2d
            The Conv2D layer

    Returns:
        Tuple of out-depth/out-height and out-width 
        Output shape as given in [Ref]
        Ref: https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html
    """
    assert isinstance(layer, nn.Conv2d)
    p = layer.padding if isinstance(layer.padding, tuple) else (layer.padding, )
    k = layer.kernel_size if isinstance(layer.kernel_size, tuple) else (layer.kernel_size, )
    d = layer.dilation if isinstance(layer.dilation, tuple) else (layer.dilation, )
    s = layer.stride if isinstance(layer.stride, tuple) else (layer.stride, )
    in_depth, in_height, in_width = x
    out_depth = layer.out_channels
    out_height = 1 + (in_height + 2 * p[0] - (k[0] - 1) * d[0] - 1) // s[0]
    out_width = 1 + (in_width + 2 * p[-1] - (k[-1] - 1) * d[-1] - 1) // s[-1]
    return (out_depth, out_height, out_width)
    

# @title Plotting functions

def plot_gen_samples_ppca(therm1, therm2, therm_data_sim):
    """
    Plotting generated samples

    Args:
        therm1: list
            Thermometer 1
        therm2: list
            Thermometer 2
        therm_data_sim: list
            Generated (simulate, draw) `n_samples` from pPCA model

    Returns:
        Nothing
    """
    plt.plot(therm1, therm2, '.', c='c', label='training data')
    plt.plot(therm_data_sim[0], therm_data_sim[1], '.', c='m', label='"generated" data')
    plt.axis('equal')
    plt.xlabel(r'Thermometer 1 ($^\circ$C)')
    plt.ylabel(r'Thermometer 2 ($^\circ$C)')
    plt.legend()
    plt.show()


def plot_linear_ae(lin_losses):
    """
    Plotting linear autoencoder

    Args:
        lin_losses: list
            Log of linear autoencoder MSE losses

    Returns:
        Nothing
    """
    plt.figure()
    plt.plot(lin_losses)
    plt.ylim([0, 2 * torch.as_tensor(lin_losses).median()])
    plt.xlabel('Training batch')
    plt.ylabel('MSE Loss')
    plt.show()


def plot_conv_ae(lin_losses, conv_losses):
    """
    Plotting convolutional autoencoder

    Args:
        lin_losses: list
            Log of linear autoencoder MSE losses
        conv_losses: list
            Log of convolutional model MSE losses

    Returns:
        Nothing
    """
    plt.figure()
    plt.plot(lin_losses)
    plt.plot(conv_losses)
    plt.legend(['Lin AE', 'Conv AE'])
    plt.xlabel('Training batch')
    plt.ylabel('MSE Loss')
    plt.ylim([0,
              2 * max(torch.as_tensor(conv_losses).median(),
                      torch.as_tensor(lin_losses).median())])
    plt.show()


def plot_images(images, h=3, w=3, plt_title=''):
    """
    Helper function to plot images

    Args:
        images: torch.tensor
            Images
        h: int
            Image height
        w: int
            Image width
        plt_title: string
            Plot title

    Returns:
        Nothing
    """
    plt.figure(figsize=(h * 2, w * 2))
    plt.suptitle(plt_title, y=1.03)
    for i in range(h * w):
        plt.subplot(h, w, i + 1)
        plot_torch_image(images[i])
    plt.axis('off')
    plt.show()


def plot_phi(phi, num=4):
    """
    Contour plot of relative entropy across samples

    Args:
        phi: list
            Log of relative entropy changes
        num: int
            Number of iterations
    """
    plt.figure(figsize=(12, 3))
    for i in range(num):
        plt.subplot(1, num, i + 1)
        plt.scatter(zs[i, :, 0], zs[i, :, 1], marker='.')
        th = torch.linspace(0, 6.28318, 100)
        x, y = torch.cos(th), torch.sin(th)
        # Draw 2-sigma contours
        plt.plot(
            2 * x * phi[i, 2].exp().item() + phi[i, 0].item(),
            2 * y * phi[i, 2].exp().item() + phi[i, 1].item()
        )
        plt.xlim(-5, 5)
        plt.ylim(-5, 5)
        plt.grid()
        plt.axis('equal')
    plt.suptitle('If rsample() is correct, then most but not all points should lie in the circles')
    plt.show()


def plot_torch_image(image, ax=None):
    """
    Helper function to plot torch image

    Args:
        image: torch.tensor
            Image
        ax: plt object
            If None, plt.gca()

    Returns:
        Nothing
    """
    ax = ax if ax is not None else plt.gca()
    c, h, w = image.size()
    if c == 1:
        cm = 'gray'
    else:
        cm = None

    # Torch images have shape (channels, height, width)
    # but matplotlib expects
    # (height, width, channels) or just
    # (height, width) when grayscale
    im_plt = torch.clip(image.detach().cpu().permute(1, 2, 0).squeeze(), 0.0, 1.0)
    ax.imshow(im_plt, cmap=cm)
    ax.set_xticks([])
    ax.set_yticks([])
    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)

# @title Set random seed

# @markdown Executing `set_seed(seed=seed)` sets the random seed.

# For deep learning, it's critical to set the random seed so that students can 
# have a baseline to compare their results to expected results.
# Read more here: https://pytorch.org/docs/stable/notes/randomness.html

# Call `set_seed` function in exercises to ensure reproducibility.
import random
import torch
import numpy as np

def set_seed(seed=None, seed_torch=True):
    """
    Function that controls randomness. NumPy and random modules must be imported.

    Args:
        seed : int, optional
            A non-negative integer that defines the random state. Default is `None`.
        seed_torch : bool, optional
            If `True`, sets the random seed for PyTorch tensors. Default is `True`.

    Returns:
        None
    """
    if seed is None:
        seed = np.random.choice(2 ** 32)
    random.seed(seed)
    np.random.seed(seed)
    if seed_torch:
        torch.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.cuda.manual_seed(seed)
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True

    print(f'Random seed {seed} has been set.')

# In case that `DataLoader` is used
def seed_worker(worker_id):
    """
    DataLoader will reseed workers following randomness in
    multi-process data loading algorithm.

    Args:
        worker_id : int
            ID of subprocess to seed. `0` means that data will be loaded in 
            the main process.
            Refer to: https://pytorch.org/docs/stable/data.html#data-loading-randomness

    Returns:
        None
    """
    worker_seed = torch.initial_seed() % 2**32
    np.random.seed(worker_seed)
    random.seed(worker_seed)

def set_device():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    if device != "cuda":
        print("WARNING: Not running with GPU")
    else:
        print("GPU is enabled")
    return device

SEED = 2021
set_seed(seed=SEED)
DEVICE = set_device()

# @title Download `wordnet` dataset

"""
NLTK Download Instructions:
To download `wordnet` using NLTK, run the following:

import nltk
nltk.download('wordnet')
"""

import os
import requests
import zipfile

# Set the environment variable for NLTK data directory
os.environ['NLTK_DATA'] = 'nltk_data/'

# Filenames and corresponding download URLs
fnames = ['wordnet.zip', 'omw-1.4.zip']
urls = ['https://osf.io/ekjxy/download', 'https://osf.io/kuwep/download']

# Download and extract each file
for fname, url in zip(fnames, urls):
    # Download the file
    response = requests.get(url, allow_redirects=True)
    with open(fname, 'wb') as file:
        file.write(response.content)
    
    # Extract the contents of the zip file
    with zipfile.ZipFile(fname, 'r') as zip_ref:
        zip_ref.extractall('nltk_data/corpora')

print("WordNet dataset downloaded and extracted successfully.")

# @title Submit your feedback
content_review(f"{feedback_prefix}_Generative_Modeling_Video")

# @title Download BigGAN (a generative model) and a few standard image datasets

# Initially, the model could be downloaded directly via:
# biggan_model = BigGAN.from_pretrained('biggan-deep-256')

import requests
import torch

# URL for the BigGAN model
url = "https://osf.io/3yvhw/download"
# Filename for the downloaded model
fname = "biggan_deep_256"

# # Download the BigGAN model
# response = requests.get(url, allow_redirects=True)
# with open(fname, 'wb') as file:
#     file.write(response.content)

# Load the downloaded BigGAN model
biggan_model = torch.load(fname)

print("BigGAN model downloaded and loaded successfully.")

!pip install nltk --quiet

# @title BigGAN Image Generator (updates may take a few seconds, please be patient)

from scipy.stats import truncnorm
import torch
import numpy as np
import matplotlib.pyplot as plt
from ipywidgets import FloatSlider, Dropdown, VBox, interactive_output, Layout

def truncated_noise_sample(batch_size=1, dim_z=128, truncation=1.0, seed=None):
    """
    Create a truncated noise vector.

    Args:
        batch_size (int): Number of samples to generate.
        dim_z (int): Dimension of the latent vector z.
        truncation (float): Truncation value to control the range of z.
        seed (int, optional): Seed for the random generator.

    Returns:
        np.ndarray: Array of shape (batch_size, dim_z) with truncated noise values.
    """
    state = None if seed is None else np.random.RandomState(seed)
    values = truncnorm.rvs(-2, 2, size=(batch_size, dim_z), random_state=state).astype(np.float32)
    return truncation * values


def sample_from_biggan(category, z_magnitude):
    """
    Sample images from BigGAN Image Generator.

    Args:
        category (str): Selected category label.
        z_magnitude (float): Magnitude of variation vector (truncation value).

    Returns:
        None
    """
    truncation = z_magnitude
    z = truncated_noise_sample(truncation=truncation, batch_size=4)
    y = one_hot_from_names(category, batch_size=4)

    # Convert to tensors
    z = torch.from_numpy(z).float()
    y = torch.from_numpy(y)

    # Move tensors and model to GPU
    device = set_device()
    z = z.to(device)
    y = y.to(device)
    biggan_model.to(device)

    # Generate images
    with torch.no_grad():
        output = biggan_model(z, y, truncation)

    # Process the output to [0, 1] range
    output = output.to('cpu')
    output = torch.clip(((output.detach().clone() + 1) / 2.0), 0, 1)

    # Plot the generated images
    fig, axes = plt.subplots(2, 2, figsize=(8, 8))
    axes = axes.flatten()
    for i in range(4):
        axes[i].imshow(output[i].squeeze().permute(1, 2, 0))
        axes[i].axis('off')
    plt.show()


# UI Widgets for interactive control
z_slider = FloatSlider(
    min=0.1, max=1.0, step=0.1, value=0.1,
    continuous_update=False,
    description='Truncation Value',
    style={'description_width': '100px'},
    layout=Layout(width='440px')
)

category_dropdown = Dropdown(
    options=['tench', 'magpie', 'jellyfish', 'German shepherd', 'bee',
             'acoustic guitar', 'coffee mug', 'minibus', 'monitor'],
    value="German shepherd",
    description="Category:"
)

widgets_ui = VBox([category_dropdown, z_slider])

# Bind the UI widgets to the function
widgets_out = interactive_output(
    sample_from_biggan,
    {
        'z_magnitude': z_slider,
        'category': category_dropdown
    }
)

# Display the widgets and output
display(widgets_ui, widgets_out)

# @markdown BigGAN Interpolation Widget (the updates may take a few seconds)

def interpolate_biggan(category_A,
                       category_B):
    """
    Interpolation function with BigGan

    Args:
        category_A: string
            Category specification
        category_B: string
            Category specification

    Returns:
        Nothing
    """
    num_interps = 16

    # category_A = 'jellyfish' #@param ['tench', 'magpie', 'jellyfish', 'German shepherd', 'bee', 'acoustic guitar', 'coffee mug', 'minibus', 'monitor']
    # z_magnitude_A = 0 #@param {type:"slider", min:-10, max:10, step:1}

    # category_B = 'German shepherd' #@param ['tench', 'magpie', 'jellyfish', 'German shepherd', 'bee', 'acoustic guitar', 'coffee mug', 'minibus', 'monitor']
    # z_magnitude_B = 0 #@param {type:"slider", min:-10, max:10, step:1}


    def interpolate_and_shape(A, B, num_interps):
        """
        Function to interpolate and shape images.
        It does this by linearly interpolating between the
        probability of each category you select and linearly
        interpolating between the latent vector values.

        Args:
            A: list
                List of categories
            B: list
                List of categories
            num_interps: int
                Quantity of pixel grids

        Returns:
            Interpolated np.ndarray
        """
        interps = interpolate(A, B, num_interps)
        return (interps.transpose(1, 0, *range(2, len(interps.shape))).reshape(num_interps, *interps.shape[2:]))

    # unit_vector = np.ones((1, 128))/np.sqrt(128)
    # z_A = z_magnitude_A * unit_vector
    # z_B = z_magnitude_B * unit_vector
    truncation = .4
    z_A = truncated_noise_sample(truncation=truncation, batch_size=1)
    z_B = truncated_noise_sample(truncation=truncation, batch_size=1)
    y_A = one_hot_from_names(category_A, batch_size=1)
    y_B = one_hot_from_names(category_B, batch_size=1)

    z_interp = interpolate_and_shape(z_A, z_B, num_interps)
    y_interp = interpolate_and_shape(y_A, y_B, num_interps)

    # Convert to tensor
    z_interp = torch.from_numpy(z_interp).float()
    y_interp = torch.from_numpy(y_interp).float()

    # Move to GPU
    z_interp = z_interp.to(DEVICE)
    y_interp = y_interp.to(DEVICE)
    biggan_model.to(DEVICE)

    with torch.no_grad():
        output = biggan_model(z_interp, y_interp, 1)

    # Back to CPU
    output = output.to('cpu')

    # The output layer of BigGAN has a tanh layer,
    # resulting the range of [-1, 1] for the output image
    # Therefore, we normalize the images properly to
    # [0, 1] range.
    # Clipping is only in case of numerical instability
    # problems

    output = torch.clip(((output.detach().clone() + 1) / 2.0), 0, 1)
    output = output

    # Make grid and show generated samples
    output_grid = torchvision.utils.make_grid(output,
                                              nrow=min(4, output.shape[0]),
                                              padding=5)
    plt.axis('off');
    plt.imshow(output_grid.permute(1, 2, 0))
    plt.show()


# z_A_slider = IntSlider(min=-10, max=10, step=1, value=0,
#                        continuous_update=False, description='Z Magnitude A',
#                        layout=Layout(width='440px'),
#                        style={'description_width': 'initial'})

# z_B_slider = IntSlider(min=-10, max=10, step=1, value=0,
#                        continuous_update=False, description='Z Magntude B',
#                        layout=Layout(width='440px'),
#                        style={'description_width': 'initial'})

category_A_dropdown = Dropdown(
    options=['tench', 'magpie', 'jellyfish', 'German shepherd', 'bee',
             'acoustic guitar', 'coffee mug', 'minibus', 'monitor'],
             value="German shepherd",
             description="Category A: ")

category_B_dropdown = Dropdown(
    options=['tench', 'magpie', 'jellyfish', 'German shepherd', 'bee',
             'acoustic guitar', 'coffee mug', 'minibus', 'monitor'],
             value="jellyfish",
             description="Category B: ")



widgets_ui = VBox([HBox([category_A_dropdown]),
                   HBox([category_B_dropdown])])

widgets_out = interactive_output(interpolate_biggan,
                                 {'category_A': category_A_dropdown,
                                  # 'z_magnitude_A': z_A_slider,
                                  'category_B': category_B_dropdown})
                                  # 'z_magnitude_B': z_B_slider})

display(widgets_ui, widgets_out)

def generate_data(n_samples, mean_of_temps, cov_of_temps, seed):
    """
    Generate random data, normally distributed

    Args:
        n_samples : int
            The number of samples to be generated
        mean_of_temps : numpy.ndarray
            1D array with the mean of temparatures, Kx1
        cov_of_temps : numpy.ndarray
            2D array with the covariance, KxK
        seed : int
            Set random seed for the pseudo random generator

    Returns:
        therm1 : numpy.ndarray
            Thermometer 1
        therm2 : numpy.ndarray
            Thermometer 2
    """

    np.random.seed(seed)
    therm1, therm2 = np.random.multivariate_normal(mean_of_temps,
                                                   cov_of_temps,
                                                   n_samples).T
    return therm1, therm2


n_samples = 2000
mean_of_temps = np.array([25, 25])
cov_of_temps = np.array([[10, 5], [5, 10]])
therm1, therm2 = generate_data(n_samples, mean_of_temps, cov_of_temps, seed=SEED)

plt.plot(therm1, therm2, '.')
plt.axis('equal')
plt.xlabel('Thermometer 1 ($^\\circ$C)')
plt.ylabel('Thermometer 2 ($^\\circ$C)')
plt.show()

# @markdown Add first PC axes to the plot

plt.plot(therm1, therm2, '.')
plt.axis('equal')
plt.xlabel('Thermometer 1 ($^\\circ$C)')
plt.ylabel('Thermometer 2 ($^\\circ$C)')
plt.plot([plt.axis()[0], plt.axis()[1]],
         [plt.axis()[0], plt.axis()[1]])
plt.show()

# Project Data onto the principal component axes.
# We could have "learned" this from the data by applying PCA,
# but we "know" the value from the problem definition.
pc_axes = np.array([1.0, 1.0]) / np.sqrt(2.0)

# Thermometers data
therm_data = np.array([therm1, therm2])

# Zero center the data
therm_data_mean = np.mean(therm_data, 1)
therm_data_center = np.outer(therm_data_mean, np.ones(therm_data.shape[1]))
therm_data_zero_centered = therm_data - therm_data_center

# Calculate the variance of the projection on the PC axes
pc_projection = np.matmul(pc_axes, therm_data_zero_centered)
pc_axes_variance = np.var(pc_projection)

# Calculate the residual variance (variance not accounted for by projection on the PC axes)
sensor_noise_std = np.mean(np.linalg.norm(therm_data_zero_centered - np.outer(pc_axes, pc_projection), axis=0, ord=2))
sensor_noise_var = sensor_noise_std ** 2

def gen_from_pPCA(noise_var, data_mean, pc_axes, pc_variance):
    """
    Generate samples from pPCA

    Args:
        noise_var: np.ndarray
            Sensor noise variance
        data_mean: np.ndarray
            Thermometer data mean
        pc_axes: np.ndarray
            Principal component axes
        pc_variance: np.ndarray
            The variance of the projection on the PC axes

    Returns:
        therm_data_sim: np.ndarray
            Generated (simulate, draw) `n_samples` from pPCA model
    """
    # We are matching this value to the thermometer data so the visualizations look similar
    n_samples = 1000

    # Randomly sample from z (latent space value)
    z = np.random.normal(0.0, np.sqrt(pc_variance), n_samples)

    # Sensor noise covariance matrix (∑)
    epsilon_cov = [[noise_var, 0.0], [0.0, noise_var]]

    # Data mean reshaped for the generation
    sim_mean = np.outer(data_mean, np.ones(n_samples))

    # Draw `n_samples` from `np.random.multivariate_normal`
    rand_eps = np.random.multivariate_normal([0.0, 0.0], epsilon_cov, n_samples)
    rand_eps = rand_eps.T

    # Generate (simulate, draw) `n_samples` from pPCA model
    therm_data_sim = sim_mean + np.outer(pc_axes, z) + rand_eps

    return therm_data_sim


## Uncomment to test your code
therm_data_sim = gen_from_pPCA(sensor_noise_var, therm_data_mean, pc_axes, pc_axes_variance)
plot_gen_samples_ppca(therm1, therm2, therm_data_sim)

# @markdown Download MNIST and CIFAR10 datasets
import tarfile, requests, os

fname = 'MNIST.tar.gz'
name = 'mnist'
url = 'https://osf.io/y2fj6/download'

if not os.path.exists(name):
    print('\nDownloading MNIST dataset...')
    r = requests.get(url, allow_redirects=True)
    with open(fname, 'wb') as fh:
        fh.write(r.content)
    print('\nDownloading MNIST completed!\n')

if not os.path.exists(name):
    with tarfile.open(fname) as tar:
        tar.extractall(name, filter=None)
        os.remove(fname)
else:
    print('MNIST dataset has been downloaded.\n')

fname = 'cifar-10-python.tar.gz'
name = 'cifar10'
url = 'https://osf.io/jbpme/download'

if not os.path.exists(name):
    print('\nDownloading CIFAR10 dataset...')
    r = requests.get(url, allow_redirects=True)
    with open(fname, 'wb') as fh:
        fh.write(r.content)
    print('\nDownloading CIFAR10 completed!')

if not os.path.exists(name):
    with tarfile.open(fname) as tar:
        tar.extractall(name, filter=None)
        os.remove(fname)
else:
    print('CIFAR10 dataset has been downloaded.')

# @markdown Load MNIST and CIFAR10 image datasets
# See https://pytorch.org/docs/stable/torchvision/datasets.html

# MNIST
mnist = datasets.MNIST('./mnist/',
                       train=True,
                       transform=transforms.ToTensor(),
                       download=False)
mnist_val = datasets.MNIST('./mnist/',
                           train=False,
                           transform=transforms.ToTensor(),
                           download=False)

# CIFAR 10
cifar10 = datasets.CIFAR10('./cifar10/',
                           train=True,
                           transform=transforms.ToTensor(),
                           download=False)
cifar10_val = datasets.CIFAR10('./cifar10/',
                               train=False,
                               transform=transforms.ToTensor(),
                               download=False)

# @markdown Execute this cell to enable helper function `get_data`

def get_data(name='mnist'):
    """
    Get data

    Args:
        name: string
            Name of the dataset

    Returns:
        my_dataset: dataset instance
            Instance of dataset
        my_dataset_name: string
            Name of the dataset
        my_dataset_shape: tuple
            Shape of dataset
        my_dataset_size: int
            Size of dataset
        my_valset: torch.loader
            Validation loader
    """
    if name == 'mnist':
        my_dataset_name = "MNIST"
        my_dataset = mnist
        my_valset = mnist_val
        my_dataset_shape = (1, 28, 28)
        my_dataset_size = 28 * 28
    elif name == 'cifar10':
        my_dataset_name = "CIFAR10"
        my_dataset = cifar10
        my_valset = cifar10_val
        my_dataset_shape = (3, 32, 32)
        my_dataset_size = 3 * 32 * 32

    return my_dataset, my_dataset_name, my_dataset_shape, my_dataset_size, my_valset

# @markdown #### Run to define the `train_autoencoder` function.
# @markdown Feel free to inspect the training function if the time allows.

# @markdown `train_autoencoder(autoencoder, dataset, device, epochs=20, batch_size=250, seed=0)`

def train_autoencoder(autoencoder, dataset, device, epochs=20, batch_size=250, seed=0):
    """
    Function to train autoencoder

    Args:
        autoencoder: nn.module
            Autoencoder instance
        dataset: function
            Dataset
        device: string
            GPU if available. CPU otherwise
        epochs: int
            Number of epochs [default: 20]
        batch_size: int
            Batch size
        seed: int
            Set seed for reproducibility; [default: 0]

    Returns:
        mse_loss: float
            MSE Loss
    """
    autoencoder.to(device)
    optim = torch.optim.Adam(autoencoder.parameters(),
                             lr=1e-3,
                             weight_decay=1e-5)
    loss_fn = nn.MSELoss()
    g_seed = torch.Generator()
    g_seed.manual_seed(seed)
    loader = DataLoader(dataset,
                        batch_size=batch_size,
                        shuffle=True,
                        pin_memory=True,
                        num_workers=2,
                        worker_init_fn=seed_worker,
                        generator=g_seed)

    mse_loss = torch.zeros(epochs * len(dataset) // batch_size, device=device)
    i = 0
    for epoch in trange(epochs, desc='Epoch'):
        for im_batch, _ in loader:
            im_batch = im_batch.to(device)
            optim.zero_grad()
            reconstruction = autoencoder(im_batch)
            # Loss calculation
            loss = loss_fn(reconstruction.view(batch_size, -1),
                           target=im_batch.view(batch_size, -1))
            loss.backward()
            optim.step()

            mse_loss[i] = loss.detach()
            i += 1
    # After training completes,
    # make sure the model is on CPU so we can easily
    # do more visualizations and demos.
    autoencoder.to('cpu')
    return mse_loss.cpu()

class LinearAutoEncoder(nn.Module):
    def __init__(self, x_dim, h_dim):
        super().__init__()
        self.enc_lin = nn.Linear(x_dim, h_dim)
        self.dec_lin = nn.Linear(h_dim, x_dim)

    def encode(self, x):
        return self.enc_lin(x)

    def decode(self, h):
        return self.dec_lin(h)

    def forward(self, x):
        flat_x = x.view(x.size()[0], -1)
        h = self.encode(flat_x)
        return self.decode(h).view(x.size())

dataset_name = 'mnist'  # This can be mnist or cifar10
train_set, dataset_name, data_shape, data_size, valid_set = get_data(name=dataset_name)

# Pick your own K
K = 20
set_seed(seed=SEED)
## Uncomment to test your code
lin_ae = LinearAutoEncoder(data_size, K)
lin_losses = train_autoencoder(lin_ae, train_set, device=DEVICE, seed=SEED)
plot_linear_ae(lin_losses)

# PCA requires finding the top K eigenvectors of the data covariance. Start by
# finding the mean and covariance of the pixels in our dataset
g_seed = torch.Generator()
g_seed.manual_seed(SEED)

loader = DataLoader(train_set,
                    batch_size=32,
                    pin_memory=True,
                    num_workers=2,
                    worker_init_fn=seed_worker,
                    generator=g_seed)
mu, cov = image_moments((im for im, _ in loader), n_batches=len(train_set) // 32)
pca_encode, pca_decode = pca_encoder_decoder(mu, cov, K)

# @markdown Visualize the reconstructions $\mathbf{x}'$, run this code a few times to see different examples.

n_plot = 7
plt.figure(figsize=(10, 4.5))
for i in range(n_plot):
    idx = torch.randint(len(train_set), size=())
    image, _ = train_set[idx]
    # Get reconstructed image from autoencoder
    with torch.no_grad():
        reconstruction = lin_ae(image.unsqueeze(0)).reshape(image.size())

    # Get reconstruction from PCA dimensionality reduction
    h_pca = pca_encode(image)
    recon_pca = pca_decode(h_pca).reshape(image.size())

    plt.subplot(3, n_plot, i + 1)
    plot_torch_image(image)
    if i == 0:
        plt.ylabel('Original\nImage')

    plt.subplot(3, n_plot, i + 1 + n_plot)
    plot_torch_image(reconstruction)
    if i == 0:
        plt.ylabel(f'Lin AE\n(K={K})')

    plt.subplot(3, n_plot, i + 1 + 2 * n_plot)
    plot_torch_image(recon_pca)
    if i == 0:
        plt.ylabel(f'PCA\n(K={K})')
plt.show()

class BiasLayer(nn.Module):
    def __init__(self, shape):
        super(BiasLayer, self).__init__()
        init_bias = torch.zeros(shape)
        self.bias = nn.Parameter(init_bias, requires_grad=True)

    def forward(self, x):
        return x + self.bias

dummy_image = torch.rand(data_shape).unsqueeze(0)
in_channels = data_shape[0]
out_channels = 7

dummy_conv = nn.Conv2d(in_channels=in_channels,
                       out_channels=out_channels,
                       kernel_size=5)

dummy_deconv = nn.ConvTranspose2d(in_channels=out_channels,
                                  out_channels=in_channels,
                                  kernel_size=5)

print(f'Size of image is {dummy_image.shape}')
print(f'Size of Conv2D(image) {dummy_conv(dummy_image).shape}')
print(f'Size of ConvTranspose2D(Conv2D(image)) {dummy_deconv(dummy_conv(dummy_image)).shape}')

class ConvAutoEncoder(nn.Module):
    def __init__(self, x_dim, h_dim, n_filters=32, filter_size=5):
        super().__init__()
        channels, height, widths = x_dim

        # encoder
        self.enc_bias = BiasLayer(x_dim)

        self.enc_conv_1 = nn.Conv2d(channels, n_filters, filter_size)
        conv1_shape = cout(x_dim, self.enc_conv_1)

        self.enc_conv_2 = nn.Conv2d(n_filters, n_filters, filter_size)
        conv2_shape = cout(conv1_shape, self.enc_conv_2)

        self.enc_flatten = nn.Flatten()
        flat_after_conv = conv2_shape[0] * conv2_shape[1] * conv2_shape[2]

        self.enc_lin = nn.Linear(flat_after_conv, h_dim)


        # decoder
        self.dec_lin = nn.Linear(h_dim, flat_after_conv)

        self.dec_unflatten = nn.Unflatten(dim=-1, unflattened_size=conv2_shape)

        self.dec_deconv_1 = nn.ConvTranspose2d(n_filters, n_filters, filter_size)
        self.dec_deconv_2 = nn.ConvTranspose2d(n_filters, channels, filter_size)

        self.dec_bias = BiasLayer(x_dim)

    def encode(self, x):
        s = self.enc_bias(x)
        s = F.relu(self.enc_conv_1(s))
        s = F.relu(self.enc_conv_2(s))
        s = self.enc_flatten(s)
        h = self.enc_lin(s)
        return h

    def decode(self, h):
        s = F.relu(self.dec_lin(h))
        s = self.dec_unflatten(s)
        s = F.relu(self.dec_deconv_1(s))
        s = self.dec_deconv_2(s)
        x_prime = self.dec_bias(s)
        return x_prime

    def forward(self, x):
        h = self.encode(x)
        x_prime = self.decode(h)
        return x_prime

set_seed(seed=SEED)
## Uncomment to test your solution
trained_conv_AE = ConvAutoEncoder(data_shape, K)
assert trained_conv_AE.encode(train_set[0][0].unsqueeze(0)).numel() == K, "Encoder output size should be K!"
conv_losses = train_autoencoder(trained_conv_AE, train_set, device=DEVICE, seed=SEED)
with plt.xkcd():
  plot_conv_ae(lin_losses, conv_losses)

# @markdown Visualize the linear and nonlinear AE outputs
if lin_ae.enc_lin.out_features != trained_conv_AE.enc_lin.out_features:
    raise ValueError('ERROR: your linear and convolutional autoencoders have different values of K')

n_plot = 7
plt.figure(figsize=(10, 4.5))
for i in range(n_plot):
    idx = torch.randint(len(train_set), size=())
    image, _ = train_set[idx]
    with torch.no_grad():
        # Get reconstructed image from linear autoencoder
        lin_recon = lin_ae(image.unsqueeze(0))[0]

        # Get reconstruction from deep (nonlinear) autoencoder
        nonlin_recon = trained_conv_AE(image.unsqueeze(0))[0]

    plt.subplot(3, n_plot, i + 1)
    plot_torch_image(image)
    if i == 0:
        plt.ylabel('Original\nImage')

    plt.subplot(3, n_plot, i + 1 + n_plot)
    plot_torch_image(lin_recon)
    if i == 0:
        plt.ylabel(f'Lin AE\n(K={K})')

    plt.subplot(3, n_plot, i + 1 + 2 * n_plot)
    plot_torch_image(nonlin_recon)
    if i == 0:
        plt.ylabel(f'NonLin AE\n(K={K})')
plt.show()