{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "cYHI6R4Ax3Fw"
},
"source": [
"# Examen Partiel: Introduction to Normalizing Flow \n",
"\n",
"This is a partial exam for the course: Deep Learning II. \n",
"The exam is divided in two parts:\n",
"- Implementing a normalizing flow model\n",
"- Training the model to generate samples from a target distribution\n",
"\n",
"The exam consists of several tasks to be implemented sequentially, each task builds on the previous one.\n",
"The more tasks you complete, the more points you will get.\n",
"\n",
"\n",
"The notebook as an `HTML` file has to be sent at 11:40 at the following email address: `alexandre.verine@dauphine.psl.eu`. Make sure I received your email before leaving the room ! \n",
"\n",
"Before starting, please make sure you disable Copilot and any other AI assistance tool.\n",
"On Colab, you can do this by going to `Tools > AI Assistance` and check `Mask AI assistance`.\n",
"Every time, the use of AI assistance will be considered as a fraud and will be penalized by 5 points. \n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "upOFQgn9JwPW"
},
"source": [
"Required modules :"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "PIqLciMKxN0S"
},
"outputs": [],
"source": [
"import torch\n",
"import torchvision\n",
"from torch import nn\n",
"import torch.nn.functional as F\n",
"import torch.nn.init as init\n",
"import torch.optim as optim\n",
"from torch.autograd import Variable\n",
"\n",
"import math\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"\n",
"device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 1. Normalizing Flows\n",
"\n",
"A Normalizing Flow is a generative model that learns a mapping from a simple distribution to a more complex one. The model is trained to maximize the log-likelihood of the data. Given a dataset $X = \\{x_1, x_2, ..., x_N\\}$, a model $F_{\\theta}$, and a latent distribution $Q$ with density $q$, the log-likelihood of the data is given by:\n",
"$$\n",
"\\log p(X) = \\sum_{i=1}^{N} \\log p(x_i) = \\sum_{i=1}^{N} \\log q(F_{\\theta}(x_i)) + \\log |det(\\frac{\\partial F_{\\theta}(x_i)}{\\partial x_i})|\n",
"$$\n",
"\n",
"Once the model is trained, we can sample from the target distribution $P$ by sampling from the latent distribution and applying the inverse transformation of the model:\n",
"- Sample $z \\sim Q$\n",
"- Compute $x = F_{\\theta}^{-1}(z)$\n",
"\n",
"The distribution $Q$ is usually chosen to be a simple distribution such as a standard Gaussian distribution. The model $F_{\\theta}$ defines a distribution $\\widehat{P}$ and the goal is to learn a mapping that makes $\\widehat{P}$ as close as possible to the true distribution $P$.\n",
"\n",
"For that reason, Normalizing Flows are a mapping learned by a neural networ that satifies two properties:\n",
"- The mapping is invertible\n",
"- The determinant of the Jacobian of the mapping can be computed efficiently\n",
"To acheive this, the model is usually composed of a sequence of invertible transformations that satisfy these properties. We can write the transformation as a composition of functions:\n",
"$$\n",
"F_{\\theta}(x) = f_K \\circ f_{K-1} \\circ ... \\circ f_1(x)\n",
"$$\n",
"and thus, the log determinant of the jacobian of the transformation is given by:\n",
"$$\n",
"\\log |det(\\frac{\\partial F_{\\theta}(x)}{\\partial x})| = \\sum_{i=1}^{K} \\log |det(\\frac{\\partial f_i(x)}{\\partial x})|\n",
"$$\n",
"and the inverse transformation is given by:\n",
"$$\n",
"F_{\\theta}^{-1}(z) = f_1^{-1} \\circ f_2^{-1} \\circ ... \\circ f_K^{-1}(z).\n",
"$$\n",
"\n",
"We can define multiple type of atomic transformations that can be used to build a normalizing flow. The most common ones are:\n",
"- Affine Coupling Layer\n",
"- Planar Flow\n",
"- Radial Flow\n",
"- Residual Flow\n",
"\n",
"In the following, we will implement two different normalizing flows using Pytorch and train them to generate samples from a 2D Gaussian distribution using:\n",
"- Linear transformations\n",
"- Affine Coupling Layers"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DdVU3d62Ki2D"
},
"source": [
"## 2. Coding coupling based Normalizing Flow"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"In this section, we will implement different types of blocks that can be used to build a normalizing flow. The block will be implemented as a Pytorch module with this backbone:\n",
"```python\n",
"\n",
"class Block(nn.Module):\n",
" def __init__(self, dim):\n",
" super(Block, self).__init__()\n",
" self.dim = dim\n",
" \n",
" def forward(self, x, inverse=False):\n",
" raise NotImplementedError\n",
"```\n",
"\n",
"The ```__init__``` method will take as input the dimension of the input and output of the block and will serve to initialize the parameters of the block. The ```forward``` method will take as input the input tensor and a boolean that indicates if the forward or inverse transformation should be applied. It should return the transformed tensor and the log determinant of the jacobian of the transformation. "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "oDvdh3jXQFF0"
},
"source": [
"### 1.1 Coupling Functions\n",
"\n",
"The Coupling Layer is a type of block that is used to build a normalizing flow. The idea is to split the input tensor $\\boldsymbol{x}$ into two parts $\\boldsymbol{x}_1$ and $\\boldsymbol{x}_2$ and apply a transformation to $\\boldsymbol{x}_1$ that depends on $\\boldsymbol{x}_2$. The transformation is defined as:\n",
"$$\n",
"\\begin{align}\n",
"\\boldsymbol{y}_1 &= \\boldsymbol{x}_1 \\\\\n",
"\\boldsymbol{y}_2 &= \\boldsymbol{x}_2 \\odot \\exp(s(\\boldsymbol{x}_1)) + t(\\boldsymbol{x}_1)\n",
"\\end{align}\n",
"$$\n",
"where $s$ and $t$ are neural networks that take $\\boldsymbol{x}_1$ as input and output the scale and translation parameters of the transformation. $\\odot$ denotes the element-wise product.\n",
"The log determinant of the jacobian of the transformation is given by:\n",
"$$\n",
"\\log \\left| \\det \\left(\\frac{\\partial f(\\boldsymbol{x})}{\\partial \\boldsymbol{x}} \\right) \\right| = \\sum_{i=1}^{D} s_i(\\boldsymbol{x}_1)\n",
"$$\n",
"where $D$ is the dimension of $\\boldsymbol{x}_2$. Finally, the inverse transformation is given by:\n",
"$$\n",
"\\begin{align}\n",
"\\boldsymbol{x}_1 &= \\boldsymbol{y}_1 \\\\\n",
"\\boldsymbol{x}_2 &= (\\boldsymbol{y}_2 - t(\\boldsymbol{y}_1)) \\odot \\exp(-s(\\boldsymbol{y}_1))\n",
"\\end{align}\n",
"$$\n",
"\n",
"To build a Normalizing Flow, we can stack multiple coupling layers to build a deep neural network. The coupling layer is usually applied to half of the dimensions of the input tensor and the other half is left unchanged. But first, we need to implement a neural network that will be used to compute the scale and translation parameters of the transformation.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**TODO:**\n",
"- Complete the forward method of the NN class. It should take as input the input tensor, apply 4 linear layers with ReLU activations and return the output tensor and no output activation."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "OVrACMakSrpb"
},
"outputs": [],
"source": [
"class NN(nn.Module):\n",
" \"\"\"Small neural network used to compute scale or translate factors.\n",
" Args:\n",
" input_dimension (int): Number of dimensions in the input.\n",
" hidden_dimension (int): Number of dimensions in the hidden layers.\n",
" output_dimension (int): Number of dimensions in the output.\n",
" \"\"\"\n",
" def __init__(self, input_dimension, hidden_dimension, output_dimension):\n",
" super(NN, self).__init__()\n",
"\n",
" self.in_layer = nn.Linear(in_features=input_dimension, out_features=hidden_dimension)\n",
" self.mid_layer1 = nn.Linear(in_features=hidden_dimension, out_features=hidden_dimension)\n",
" self.mid_layer2 = nn.Linear(in_features=hidden_dimension, out_features=hidden_dimension)\n",
" self.out_layer = nn.Linear(in_features=hidden_dimension, out_features=output_dimension)\n",
" self.activation = nn.ReLU()\n",
"\n",
" def forward(self, x):\n",
" ### YOUR CODE HERE\n",
" return x\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KVMOMOFxOCG0"
},
"source": [
"**TODO:**\n",
"- Complete the AffineCouplingBlock class. Implement the forward and inverse transformations using the equations above. \n",
"- Run the next cell to check if your code works."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "YcVP5SsuSuPX"
},
"outputs": [],
"source": [
"class AffineCouplingBlock(nn.Module):\n",
" \"\"\"Affine Coupling Layer. \n",
" Split x into xA, xB and return (s(xB)*xA+t(xB), xB). \n",
" Args:\n",
" input_dimension (int): Number of dimensions in the input.\n",
" hidden_dimension (int): Number of dimensions in the hidden layers.\n",
" alternate (bool): reverse xA and xB.\n",
" activation (bool): use activation in t.\n",
" \"\"\"\n",
" def __init__(self, input_dimension, \n",
" hidden_dimension, alternate):\n",
" super(AffineCouplingBlock, self).__init__()\n",
" assert input_dimension%2 == 0\n",
" self.st = NN(input_dimension//2, hidden_dimension,\n",
" input_dimension)\n",
" self.scale = nn.Parameter(torch.ones(input_dimension//2))\n",
" self.alternate = alternate\n",
" \n",
" def forward(self, x, ldj, reverse=False):\n",
"\n",
" if self.alternate:\n",
" xB, xA = x.chunk(2, dim=1)\n",
" else:\n",
" xA, xB = x.chunk(2, dim=1)\n",
"\n",
" st = self.st(xB)\n",
" s, t = st[:, 0::2], st[:, 1::2]\n",
" s = self.scale * torch.tanh(s)\n",
"\n",
" if reverse:\n",
" ldj = ldj - torch.sum(s, dim=1).view(-1, 1)\n",
" s = torch.exp(-s)\n",
" yA = ### YOUR CODE HERE\n",
" else:\n",
" ldj = ldj + torch.sum(s, dim=1).view(-1, 1)\n",
" s = torch.exp(s)\n",
" yA = ### YOUR CODE HERE\n",
"\n",
" if self.alternate:\n",
" y = torch.cat((xB, yA), dim=1)\n",
" else:\n",
" y = torch.cat((yA, xB), dim=1)\n",
" \n",
" return y, ldj\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2ztFHJRwFfCu"
},
"outputs": [],
"source": [
"f = AffineCouplingBlock(2, 64, True) # Define the invertible block\n",
"if torch.cuda.is_available(): # Move the block to the GPU\n",
" f = f.cuda()\n",
"x = torch.randn((10,2)).to(device) # Random input\n",
"ldj = torch.zeros((x.size(0), 1)).to(device) # Log-determinant Jacobian\n",
"\n",
"y, ldjo = f(x, ldj.clone()) # Forward pass\n",
"xi, ldji = f(y, ldj.clone(), True) # Inverse pass\n",
"torch.testing.assert_close(xi, x, atol=0.01, rtol=0.01) # Check if the inverse is correct\n",
"torch.testing.assert_close(ldji+ ldjo, ldj, atol=0.01, rtol=0.01) # Check if the log-determinant Jacobian is correct\n",
"\n",
"del f, x, xi, y, ldjo, ldj, ldji"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"As previously explained, a Normalizing Flow is a composition of multiple transformations. To build a normalizing flow, we can stack multiple coupling layers to build a deep neural network. The coupling layer is usually applied to half of the dimensions of the input tensor and the other half is left unchanged, alternatively. We can build a normalizing flow by stacking multiple coupling layers:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "4eAu0WZ0xXYY"
},
"outputs": [],
"source": [
"class NormalizingFlow(nn.Module):\n",
" \"\"\"Normalizing Flow class.\n",
" Takes x as input and return F(x) and sldj. If F(y, reverse=True) is \n",
" called then it returns the inverse of F. \n",
" Args:\n",
" input_dimension (int): Number of dimensions in the input.\n",
" hidden_dimension (int): Number of dimensions in the hidden layers.\n",
" \"\"\"\n",
" def __init__(self, dimension=1, hidden_dimension=256, num_steps=10):\n",
" super(NormalizingFlow, self).__init__()\n",
" self.dimension = dimension\n",
" self.flows = [AffineCouplingBlock(input_dimension=dimension,\n",
" hidden_dimension=hidden_dimension,\n",
" alternate = (depth%2==0))\n",
" for depth in range(num_steps)]\n",
" self.flows = nn.ModuleList(self.flows)\n",
"\n",
" def forward(self, x, reverse=False):\n",
" sldj = torch.zeros((x.size(0), 1)).to(device)\n",
" if reverse:\n",
" flows = reversed(self.flows)\n",
" else:\n",
" flows = self.flows\n",
" for block in flows:\n",
" x, sldj = block(x, sldj, reverse)\n",
" return x, sldj "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "24-rOCxCl2FQ"
},
"source": [
"**TODO:**\n",
"- Run the next cell to check if your code works."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "e1WTaPrtZ2Rn"
},
"outputs": [],
"source": [
"f = NormalizingFlow(dimension=10, num_steps=10)\n",
"if torch.cuda.is_available():\n",
" f = f.cuda()\n",
"x = torch.randn((5,10)).to(device)\n",
"ldj = torch.zeros((x.size(0), 1)).to(device)\n",
"y, ldjo = f(x)\n",
"xi, ldji = f(y, True)\n",
"torch.testing.assert_close(xi, x, atol=0.01, rtol=0.01)\n",
"torch.testing.assert_close(ldji+ldjo, ldj, atol=0.1, rtol=0.1)\n",
"del f, x, xi, y, ldjo, ldj, ldji"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "eTR1deFgpgJ5"
},
"source": [
"## 3. Training Functions\n",
"\n",
"In this section, we will implement the training loop for the normalizing flow. The training loop will be implemented as a function that takes as input the model ```model```, the optimizer ```optimizer```, the number of epochs ```n_epochs```, the initial learning rate ```lr```, the list of epochs to decay the learning rate ```epochs_to_update_lr```, the frequence of epoch to print the loss ```freq_verbose```, the frequence to plot the samples ```freq_plot``` and the function used to plot the samples ```plot_function```.\n",
"\n",
"We will consider that the loaders for the training and validation datasets ```loader_train``` and ```loader_test``` are given. \n",
"First we will define the loss function. As previously explained, the log-likelihood of the data is given by:\n",
"$$\n",
"\\log p(X) = \\sum_{i=1}^{N} \\log q(F_{\\theta}(x_i)) + \\log |det(\\frac{\\partial F_{\\theta}(x_i)}{\\partial x_i})|\n",
"$$\n",
"where $q$ is the density of the latent distribution. The loss function is the negative log-likelihood of the data.\n",
"Therefore, we will define a function ```loss_function``` that takes as input a latent vector ```z``` and the log determinant of the jacobian of the transformation ```logdetjac``` and returns the average negative log-likelihood of the data:\n",
"$$\n",
"\\text{loss} = -\\frac{1}{N} \\sum_{i=1}^{N}\\left[ \\log q(z_i) + \\text{logdetjac}_i \\right].\n",
"$$\n",
"As the latent distribution is a standard normal distribution, the density $q$ is given by:\n",
"$$\n",
"q(z) = \\frac{1}{\\sqrt{2\\pi}^D } \\exp\\left(-\\frac{1}{2}\\Vert z \\Vert_2^2\\right),\n",
"$$\n",
"where $D$ is the dimension of the latent vector. Therefore, the log density is given by:\n",
"$$\n",
"\\log q(z) = -\\frac{1}{2} D \\log(2\\pi) - \\frac{1}{2} \\Vert z \\Vert_2^2.\n",
"$$\n",
"It can be split into two terms:\n",
"- The constant term $-\\frac{1}{2} D \\log(2\\pi)$\n",
"- The term $-\\frac{1}{2} \\Vert z \\Vert_2^2$\n",
"\n",
"**TODO:**\n",
"- Complete the loss_function function. It should return the average negative log-likelihood of the data.\n",
"- Run the next cell to check if your code works."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def loss_function(z, logdetjac):\n",
" \"\"\"\n",
" Computes the Negative Log-Likelihood (NLL) Loss for Normalizing Flows.\n",
"\n",
" Parameters:\n",
" - z: The transformed data in latent space, with shape (batch_size, dimensions).\n",
" This represents the result after applying the flow transformation to the input data.\n",
" - logdetjac: The log determinant of the Jacobian of the transformation, \n",
" with shape (batch_size, 1).\n",
"\n",
" Returns:\n",
" - The average negative log-likelihood loss for the batch.\n",
" \"\"\"\n",
"\n",
" # Step 1: Calculate the likelihood contribution from z\n",
" likelihood_term = ### YOUR CODE HERE\n",
"\n",
" # Step 2: Calculate the constant term\n",
" constant_term = ### YOUR CODE HERE\n",
"\n",
" # Step 3: Add the log determinant of the Jacobian (logdetjac)\n",
"\n",
" jacobian_term = logdetjac.view(-1) # Reshape to ensure it's a 1D vector\n",
"\n",
" # Step 4: Combine all terms to get the negative log-likelihood (NLL) for each sample\n",
" loss = ### YOUR CODE HERE\n",
"\n",
" # Step 5: Take the mean over all samples in the batch to get the average NLL\n",
" # Multiplying by -1 here because we minimize the loss during training\n",
" loss = ### YOUR CODE HERE\n",
"\n",
" return loss"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"x = torch.Tensor([[-0.5689, 1.2247, 1.4252, -2.1294, -1.3356, 0.6384, 0.8114, 0.9037,\n",
" -0.7982, -0.0422],\n",
" [ 1.0852, -0.7060, -0.7232, -0.4275, -0.2482, -0.1510, -1.5881, -1.1483,\n",
" -1.2538, 0.3152],\n",
" [ 1.5991, 0.8852, 0.9818, -0.2948, 0.0666, -0.8725, 1.3663, -0.9406,\n",
" -0.4303, 0.4287],\n",
" [-1.7749, 0.2848, 0.1286, -1.1444, 1.3227, -0.8872, 0.2311, -0.8701,\n",
" -1.0925, -1.2644],\n",
" [ 2.9917, -1.7241, 0.3513, -0.1070, -2.3848, 0.3795, 1.1048, 1.1429,\n",
" -0.7369, 0.5604]]).to(device)\n",
"logdetjac = torch.Tensor([[ 0.1145],\n",
" [ 1.0377],\n",
" [ 2.8035],\n",
" [-0.0696],\n",
" [ 0.0674]]).to(device)\n",
"\n",
"assert torch.isclose(loss_function(x, logdetjac), torch.tensor(14.4907).to(device), atol=1e-5)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We give the function to evaluate the model on the test set. The function ```evaluate``` takes as input the model ```model``` returns the average negative log-likelihood of the data."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def eval_model(model):\n",
" model.eval()\n",
" losses = []\n",
" with torch.no_grad():\n",
" for batch_idx, x in enumerate(loader_test):\n",
" x = x.to(device)\n",
" z, logdetjac = model(x)\n",
" loss = loss_function(z, logdetjac)\n",
" losses.append(loss.item())\n",
" model.train()\n",
" return np.mean(losses) "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**TODO:**\n",
"- Complent the function train to perform the backpropagation and optimization of the model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def train(model, optimizer, lr, n_epochs = 20, freq_verbose=2, freq_plot=2, plot_function=None):\n",
" train_loss = [] \n",
" test_loss = []\n",
" for epoch in range(0, n_epochs+1):\n",
" losses = []\n",
" for batch_idx, x in enumerate(loader_train):\n",
"\n",
" ### YOUR CODE HERE\n",
" loss = ### YOUR CODE HERE\n",
" ### YOUR CODE HERE\n",
"\n",
" losses.append(loss.item())\n",
" train_loss.append(np.mean(losses))\n",
" test_loss.append(eval_model(model))\n",
" \n",
" if epoch %freq_verbose==0:\n",
" print('[%d/%d]: \\tloss: %.3f \\tloss eval: %.5f' % ((epoch), n_epochs, train_loss[-1], test_loss[-1]))\n",
" if epoch %freq_plot==0 and epoch>0:\n",
" if plot_function is not None:\n",
" plot_function(model, epoch)\n",
" plt.plot(train_loss, label='train')\n",
" plt.plot(test_loss, label='test')\n",
" plt.legend()\n",
" plt.xlabel('Epoch')\n",
" plt.ylabel('NLL')\n",
" plt.show() "
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IR5r7U4Quu4l"
},
"source": [
"**TODO:**\n",
"- Run the next cell to check if your code works.\n",
"Run this cellule to check if your code works :"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9clAmKi3rjJR"
},
"outputs": [],
"source": [
"batch_size = 256\n",
"z = torch.randn((batch_size, 25)).to(device)\n",
"logdetjac = torch.rand((batch_size, 1)).to(device)\n",
"loss1 = loss_function(z, logdetjac)\n",
"loss2 = loss_function(z*2+1, 2*logdetjac)\n",
"assert loss1 < loss2, \"Loss1 should be lower than Loss2\"\n",
"del batch_size, z, logdetjac, loss1, loss2"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7bQ3zQVqxGc4"
},
"source": [
"## 4. Low Dimensional Dataset"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "K1D462nPxORx"
},
"source": [
"We train the Normalizing Flow to generate samples from a 2D distribution (A Gaussian Mixure Model). We will use the following dataset:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "sQy59SHzxNvE"
},
"outputs": [],
"source": [
"def Gaussians(batch_size):\n",
" scale = 4.\n",
" centers = [(-1, 0.5),(-0.35, -0.5), (0.35, -0.5), (1, 0.5)]\n",
" centers = [(scale * x, scale * y) for x, y in centers]\n",
"\n",
" dataset = []\n",
" for i in range(batch_size):\n",
" point = np.random.randn(2)*1.2\n",
" idx = np.random.randint(4)\n",
" center = centers[idx]\n",
" point[0] += center[0]\n",
" point[1] += center[1]\n",
" dataset.append(point)\n",
" dataset = np.array(dataset, dtype=\"float32\")\n",
" dataset /= 1.414\n",
" return torch.Tensor(dataset)\n",
"\n",
"plt.figure(figsize=(5,5))\n",
"x = Gaussians(5000).numpy()\n",
"plt.hist2d(x[:,0], x[:,1],\n",
" range=[[-5, 5], [-5, 5]], bins=100, cmap='inferno')\n",
"plt.xticks([])\n",
"plt.yticks([])\n",
"plt.xlim(-5, 5)\n",
"plt.ylim(-5, 5)\n",
"plt.title(\"Dataset $X_r\\sim P$\")\n",
"plt.show()\n",
"del x"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "gq_3LaVDx-AU"
},
"source": [
"We define a function to plot the dataset $X_r$ sampled from the target distribution, the samples $X_g=F_{\\theta}^{-1}(Z_g)\\sim \\widehat{P}$ generated by the model by applying the reverse transformation to samples $Z_g\\sim Q$. We also plot the samples $Z_r=F_{\\theta}(X_r)$ and $Z_g\\sim Q$."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cm1q7eiUpeTZ"
},
"outputs": [],
"source": [
"def plot_2D(model, epoch, sample=Gaussians, device=device):\n",
" plt.clf()\n",
" plt.figure(1, figsize=(7,10))\n",
" plt.suptitle(\"Epoch \"+str(epoch), y=0.95)\n",
" plt.subplot(2,2,2)\n",
" plt.title(\"$X_g = F^{-1}(Z_g)$\")\n",
" z = Variable(torch.randn((1000, 2)).to(device))\n",
" x_gen, ldj = model(z,True)\n",
" x_gen = x_gen.detach().cpu().numpy()\n",
" plt.hist2d(x_gen[:,0], x_gen[:,1],\n",
" range=[[-5, 5], [-5, 5]], bins=100, cmap='inferno')\n",
" plt.xticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.yticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.xlim(-5, 5)\n",
" plt.ylim(-5, 5)\n",
" plt.subplot(2,2,1)\n",
" plt.title(\"$X_r$\")\n",
" x = sample(1000)\n",
" plt.hist2d(x[:,0].numpy(), x[:,1].numpy(),\n",
" range=[[-5, 5], [-5, 5]], bins=100, cmap='inferno')\n",
" plt.xticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.yticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.xlim(-5, 5)\n",
" plt.ylim(-5, 5)\n",
" plt.subplot(2,2,3)\n",
" plt.title('$Z_r=F(X_r)$')\n",
" x = Variable(x.to(device))\n",
" z, ldj = model(x)\n",
" z = z.detach().cpu().numpy()\n",
" plt.hist2d(z[:,0],z[:,1] , \n",
" range=[[-5, 5], [-5, 5]], bins=100, cmap='inferno')\n",
" plt.xlim(-5, 5)\n",
" plt.ylim(-5, 5)\n",
" plt.xticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.yticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.subplot(2,2,4)\n",
" plt.title('$Z_g$')\n",
" z = torch.randn((1000, 2))\n",
" plt.hist2d(z[:,0].numpy(),z[:,1].numpy(), \n",
" range=[[-5, 5], [-5, 5]], bins=100, cmap='inferno')\n",
" plt.xlim(-5, 5)\n",
" plt.ylim(-5, 5)\n",
" plt.xticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.yticks(ticks=[-5, -2.5, 0, 2.5, 5], labels=[\"\",\"\",\"\",\"\",\"\"])\n",
" plt.tight_layout()\n",
" plt.show()\n",
" plt.close()\n",
"\n",
"model = NormalizingFlow(dimension=2, hidden_dimension=10, num_steps=20).to(device)\n",
"plot_2D(model, 0)\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"To interpret the results:\n",
"- The samples $X_r$ should be close to the samples $X_g$.\n",
"- The samples $Z_r$ should be close to the samples $Z_g$.\n",
"\n",
"\n",
"We will now define the dataset and the loaders for the training and validation datasets.\n",
"**TODO:**\n",
"- Define the loaders for the training and validation datasets with a batch size of 1024. Shuffle the training dataset and do not shuffle the validation dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dataset_train = Gaussians(102400)\n",
"dataset_test = Gaussians(10240)\n",
"loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=1024, shuffle=True)\n",
"loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1024, shuffle=False)\n",
"\n",
"print(f'Number of training samples: {len(dataset_train)}')\n",
"print(f'Number of test samples: {len(dataset_test)}')\n",
"print(f'Number of training batches: {len(loader_train)}')\n",
"print(f'Number of test batches: {len(loader_test)}')\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**TODO:**\n",
"- Train the model with the constraints given below.\n",
"\n",
"Train the model with the following parameters:\n",
"- The dimension of the latent space is 2\n",
"- The number of coupling layers is 2\n",
"- The hidden dimension of the neural networks is 10\n",
"\n",
"Using the following optimizer:\n",
"- Adam optimizer with a learning rate of between 0.01 and 0.0001.\n",
"- The number of epochs is 10\n",
"- Print the loss every epoch\n",
"- Plot the samples/losses at the end of the training\n",
"\n",
"Don't forget to copy the model on the desired device (CPU or GPU).\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "6UMI9Uw_yfX1"
},
"outputs": [],
"source": [
"model = ### YOUR CODE HERE\n",
"lr = ### YOUR CODE HERE\n",
"optimizer = ### YOUR CODE HERE\n",
"train(model, optimizer, lr, n_epochs=10, freq_verbose=1, freq_plot=10)\n",
"assert eval_model(model) < 3.9, \"The model should have a test loss lower than 3.9\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**TODO:**\n",
"- Train a deeper model with 10 coupling layers and hidden dimension of 10. Use the same optimizer and hyperparameters as before."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = ### YOUR CODE HERE\n",
"lr = ### YOUR CODE HERE\n",
"optimizer = ### YOUR CODE HERE\n",
"train(model, optimizer, lr, n_epochs=10, freq_verbose=1, freq_plot=10)\n",
"\n",
"assert eval_model(model) < 3.75, \"The model should have a test loss lower than 3.75\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**TODO:**\n",
"- Finally, train a model with 10 coupling layers and hidden dimension of 100. (You might need to reduce the learning rate)."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = ### YOUR CODE HERE\n",
"lr = ### YOUR CODE HERE\n",
"optimizer = ### YOUR CODE HERE\n",
"train(model, optimizer, lr, n_epochs=10, freq_verbose=1, freq_plot=10)\n",
"assert eval_model(model) < 3.7, \"The model should have a test loss lower than 3.7\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## 5. High Dimensional Dataset\n",
"\n",
"We will now train the Normalizing Flow to generate samples from a higher dimensional distribution. We will use the small version of the MNIST dataset. To do so, we used a rescaled version of the dataset with a size of $size \\times size$ pixels of 0 and 1."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torchvision.transforms as transforms\n",
"\n",
"size = 12\n",
"# Transform to resize, convert to tensor, and flatten to a vector of size (size)\n",
"transform = transforms.Compose([\n",
" transforms.Resize((size, size)), # Resize to sizexsize pixels\n",
" transforms.ToTensor(), # Convert to a tensor with shape (1, size, size)\n",
" transforms.Lambda(lambda x: x.view(size*size)), # Flatten to (size*size,) and normalize\n",
" transforms.Lambda(lambda x: (x+5e-2)/(1+1e-1)) # Normalize to [0.05, 0.95] to avoid overflow\n",
"])\n",
"\n",
"# Custom dataset to filter only labels 0 and 1\n",
"class FilteredMNIST(torchvision.datasets.MNIST):\n",
" def __init__(self, *args, **kwargs):\n",
" super().__init__(*args, **kwargs)\n",
" # Filter to keep only images with labels 0 and 1\n",
" indices = [i for i, label in enumerate(self.targets) if label in [0, 1]]\n",
" self.data = self.data[indices]\n",
" self.targets = self.targets[indices]\n",
"\n",
" def __getitem__(self, index):\n",
" # Return only the image, not the label\n",
" image, _ = super().__getitem__(index)\n",
" return image\n",
"\n",
"# Create datasets and loaders\n",
"dataset_train = FilteredMNIST(root='./data', train=True, download=True, transform=transform)\n",
"dataset_test = FilteredMNIST(root='./data', train=False, download=True, transform=transform)\n",
"\n",
"loader_train = torch.utils.data.DataLoader(dataset_train, batch_size=128, shuffle=True)\n",
"loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=128, shuffle=True)\n",
"\n",
"first_batch = next(iter(loader_test))\n",
"\n",
"plt.figure(1, figsize=(10,10))\n",
"plt.suptitle(\"Test Set\", y=0.6)\n",
"for i in range(10):\n",
" plt.subplot(1, 10, i + 1)\n",
" plt.imshow(first_batch[i].view(size, size).numpy(), cmap='gray')\n",
" plt.axis('off')\n",
"plt.show()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The image are slightly scaled to be between 0.05 and 0.95 approximately. Moreover, they are directly flattened to be of size $size^2$:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"print(f'Size of the images: {size}x{size}')\n",
"print(f'First batch shape: {first_batch.shape}')\n",
"print(f'Min/Max pixel value: {first_batch.min():.2f}, {first_batch.max():.2f}')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will use the function to plot the generated samples:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def plot_mnist(model, epoch):\n",
" plt.clf()\n",
" plt.figure(1, figsize=(10,10))\n",
" plt.suptitle(\"Epoch \"+str(epoch), y=0.6)\n",
" z = torch.randn((10, model.dimension)).to(device)\n",
" x_gen, ldj = model(z,True)\n",
" x_gen = x_gen.detach().cpu().numpy().reshape(-1, np.sqrt(model.dimension).astype(int), np.sqrt(model.dimension).astype(int))\n",
" for i in range(10):\n",
" plt.subplot(1,10,i+1)\n",
" plt.imshow(x_gen[i], cmap='gray')\n",
" plt.axis('off')\n",
" plt.show()\n",
"\n",
"model = NormalizingFlow(num_steps=10, dimension=size*size).to(device)\n",
"plot_mnist(model, 0)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The image generated should be similat to the images from the dataset. However, the image are now completly random since the model is not trained yet.\n",
"\n",
"**TODO:**\n",
"- Train a model with 10 coupling layers and hidden dimension of 256 using Adam optimizer with a learning rate of 0.003. The number of epochs is 20. Print the loss every epoch and plot the samples/losses at the end of the training."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = ### YOUR CODE HERE\n",
"lr = ### YOUR CODE HERE\n",
"optimizer = ### YOUR CODE HERE\n",
"train(model, optimizer, lr, n_epochs=20, freq_verbose=1, freq_plot=20, plot_function=plot_mnist)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can see that the model is generating samples in a larger domain than expected:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"z = torch.randn((10, size*size)).to(device)\n",
"x_gen, ldj = model(z,True)\n",
"print(f'Min/Max pixel value: {x_gen.min()}, {x_gen.max()}')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can therefore help the model by making sure that the last operation made by the reversed transformation is a sigmoid function that will make sure that the output is between 0 and 1. To do that, we need to implement a block that will apply the sigmoid function to the output of the model. However, the block must be implemented as a ReverseSigmoid to be applied as the first block of the model. \n",
"\n",
"The sigmoid function is defined as follows:\n",
"$$\n",
"\\text{sigmoid}(x) = \\frac{1}{1+\\exp(-x)}.\n",
"$$\n",
"The derivative of the sigmoid function is given by:\n",
"$$\n",
"\\partial_x \\text{sigmoid}(x) = \\text{sigmoid}(x)\\times(1-\\text{sigmoid}(x)).\n",
"$$\n",
"Since the function is element-wise, the log determinant of the jacobian of the transformation is given by:\n",
"$$\n",
"\\log \\left| \\det \\left(\\frac{\\partial (\\text{sigmoid}\\boldsymbol{x})}{\\partial \\boldsymbol{x}} \\right) \\right| = \\sum_{i=1}^{D} \\log \\left( \\partial_{x_i} \\text{sigmoid}(x_i) \\right).\n",
"$$\n",
"\n",
"The inverse transformation is given by:\n",
"$$\n",
"\\text{sigmoid}^{-1}(x) = \\log\\left(\\frac{x}{1-x}\\right).\n",
"$$\n",
"and its derivative is given by:\n",
"$$\n",
"\\partial_x \\text{sigmoid}^{-1}(x) = \\frac{1}{x(1-x)}.\n",
"$$\n",
"\n",
"**TODO:**\n",
"- Complete the InvertibleReverseSigmoid class. Implement the forward and inverse transformations using the equations above.\n",
"- Run the next cell to check if your code works.\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class InvertibleReverseSigmoid(nn.Module):\n",
" def __init__(self, dim):\n",
" super(InvertibleReverseSigmoid, self).__init__()\n",
" self.dim = dim\n",
"\n",
" def forward(self, x, ldj, inverse=False):\n",
" if inverse:\n",
" # Apply the sigmoid function\n",
" y = ### YOUR CODE HERE\n",
" # Compute the log-determinant of the Jacobian for sigmoid\n",
" jacobian_diag = ### YOUR CODE HERE\n",
" ldj += ### YOUR CODE HERE\n",
" return y, ldj\n",
" else:\n",
" # Apply the logit (inverse sigmoid) function\n",
" y = ### YOUR CODE HERE\n",
" # Compute the log-determinant of the Jacobian for the logit\n",
" jacobian_diag = ### YOUR CODE HERE\n",
" ldj += ### YOUR CODE HERE\n",
" return y, ldj"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"f = InvertibleReverseSigmoid(10)\n",
"x = torch.rand((10, 10))\n",
"ldj = torch.zeros((10, 1))\n",
"y, ldjo = f(x, ldj.clone())\n",
"xi, ldji = f(y, ldj.clone(), True)\n",
"torch.testing.assert_close(xi, x, atol=0.01, rtol=0.01)\n",
"torch.testing.assert_close(ldji+ldjo, ldj, atol=0.01, rtol=0.01)\n",
"del f, x, xi, y, ldjo, ldj, ldji"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**TODO:**\n",
"- Complete the code for the NormalizingFlowWithSigmoid Class by adding the ReverseSigmoid block at the beginning of the model.\n",
"- Run the next cell to check if your code works."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"class NormalizingFlowWithSigmoid(nn.Module):\n",
" def __init__(self, dimension=1, hidden_dimension=256, num_steps=10):\n",
" super(NormalizingFlowWithSigmoid, self).__init__()\n",
" self.dimension = dimension\n",
" self.activation = ### YOUR CODE HERE\n",
" self.flows_blocks = [AffineCouplingBlock(input_dimension=dimension,\n",
" hidden_dimension=hidden_dimension,\n",
" alternate=(depth % 2 == 0))\n",
" for depth in range(num_steps)]\n",
" self.flows = ### YOUR CODE HERE\n",
"\n",
" def forward(self, x, reverse=False):\n",
" sldj = torch.zeros((x.size(0), 1)).to(device)\n",
" if reverse:\n",
" flows = reversed(self.flows)\n",
" else:\n",
" flows = self.flows\n",
" for block in flows:\n",
" x, sldj = block(x, sldj, reverse)\n",
" return x, sldj"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"f = NormalizingFlowWithSigmoid(dimension=10, num_steps=10)\n",
"if torch.cuda.is_available():\n",
" f = f.cuda()\n",
"x = torch.rand((5, 10)).to(device)\n",
"ldj = torch.zeros((x.size(0), 1)).to(device)\n",
"y, ldjo = f(x)\n",
"xi, ldji = f(y, True)\n",
"torch.testing.assert_close(xi, x, atol=0.01, rtol=0.01)\n",
"torch.testing.assert_close(ldji+ldjo, ldj, atol=0.01, rtol=0.01)\n",
"torch.testing.assert_close(xi, xi.clamp(0, 1), atol=0.01, rtol=0.01)\n",
"del f, x, xi, y, ldjo, ldj, ldji"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**TODO:**\n",
"- Train the model with the same hyperparameters as before. (You might consider a smaller learning rate).\n",
"- Observe the results and compare them with the previous model."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"model = ### YOUR CODE HERE\n",
"lr = ### YOUR CODE HERE\n",
"optimizer = ### YOUR CODE HERE\n",
"train(model, optimizer, lr, n_epochs=20, freq_verbose=1, freq_plot=20, plot_function=plot_mnist)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"z = torch.randn((10, size*size)).to(device)\n",
"x_gen, ldj = model(z,True)\n",
"print(f'Min/Max pixel value: {x_gen.min()}, {x_gen.max()}')\n",
"plt.figure(1, figsize=(10,10))\n",
"plt.suptitle(\"Generated Images\", y=0.6)\n",
"for i in range(10):\n",
" plt.subplot(1, 10, i + 1)\n",
" plt.imshow(x_gen[i].view(size, size).detach().cpu().numpy(), cmap='gray')\n",
" plt.axis('off')\n",
"plt.show()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
" "
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"authorship_tag": "ABX9TyNxpDDiOn9pOoEPML16KIV1",
"collapsed_sections": [],
"name": "NormalizingFlow_AdvancedML_Verine2022_Correction.ipynb",
"private_outputs": true,
"provenance": [],
"toc_visible": true
},
"kernelspec": {
"display_name": "venv",
"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.11.2"
}
},
"nbformat": 4,
"nbformat_minor": 0
}