Variational Autoencoder (VAE)
A Variational Autoencoder (VAE) is a type of artificial neural network that belongs to the family of generative models. Unlike traditional autoencoders that learn a deterministic mapping from the input to a lower-dimensional latent space, VAEs learn a probabilistic distribution over this latent space. This key difference allows VAEs not only to compress and reconstruct data but also to generate new data samples that resemble the training data.
Here's a breakdown of the key components and concepts behind VAEs:
1. Architecture:
Like standard autoencoders, a VAE consists of two main parts:
Encoder: The encoder takes the input data and maps it into the parameters of a probability distribution in the latent space. Typically, this distribution is assumed to be a multivariate Gaussian, and the encoder outputs the mean (μ) and the standard deviation (σ) (or variance σ2) of this distribution for each latent dimension.
Decoder: The decoder takes a sample from the latent distribution (obtained using the parameters from the encoder) and maps it back to the original data space, attempting to reconstruct the input.
2. The Probabilistic Latent Space:
The crucial difference from standard autoencoders is that the encoder in a VAE doesn't output a single fixed vector representing the compressed input. Instead, it outputs the parameters of a probability distribution. This has several important implications:
Learning a Distribution: The model learns the underlying probability distribution of the data, capturing the inherent variability and structure.
Smooth Latent Space: By learning distributions, VAEs encourage a continuous and smooth latent space. This means that points close to each other in the latent space should correspond to similar data points when decoded, allowing for meaningful interpolation between data samples.
Generative Capability: Because the latent space is a probability distribution, you can generate new data by randomly sampling points from this space and feeding them into the decoder.
3. The Reparameterization Trick:
A key challenge in training VAEs is that sampling from a probability distribution is a non-deterministic operation, which makes it difficult to backpropagate gradients through the sampling process. To address this, VAEs use a technique called the reparameterization trick.
For a Gaussian distribution with mean μ and standard deviation σ, the reparameterization trick involves expressing a sample z as:
z=μ+σ⋅ϵ
where ϵ is a random variable sampled from a standard normal distribution (N(0,1)). This allows the gradients to flow through the mean and standard deviation parameters during backpropagation, as the randomness is now isolated in ϵ, which is treated as a constant during the backward pass.
4. Loss Function:
The loss function for a VAE has two main components:
Reconstruction Loss: This term measures how well the decoder can reconstruct the original input from the sampled latent vector z. Common reconstruction loss functions include Mean Squared Error (MSE) for continuous data or binary cross-entropy for binary data. It encourages the model to learn a latent representation that retains the important information from the input.
KL Divergence (Regularization Term): The Kullback-Leibler (KL) divergence measures the difference between the learned latent distribution q(z∣x) (the distribution parameterized by the encoder's output for a given input x) and a prior distribution p(z) over the latent space. Typically, the prior p(z) is chosen to be a standard normal distribution (N(0,1)). This term acts as a regularizer, encouraging the learned latent distributions to be close to the prior. This helps ensure that the latent space is well-structured and continuous, which is crucial for the generative capabilities of the VAE.
The total loss function is a weighted sum of these two terms:
L(x,z)=Ez∼q(z∣x)[logp(x∣z)]−DKL(q(z∣x)∣∣p(z))
where:
Ez∼q(z∣x)[logp(x∣z)] is the expected log-likelihood of the reconstruction (related to the reconstruction loss).
DKL(q(z∣x)∣∣p(z)) is the KL divergence between the learned approximate posterior q(z∣x) and the prior p(z).
The goal of training is to minimize this loss function, which simultaneously encourages accurate reconstruction and a well-behaved latent space.
In summary, a Variational Autoencoder is a generative model that learns a probabilistic latent representation of the input data. By encoding data into the parameters of a probability distribution and using a carefully designed loss function with a regularization term (KL divergence), VAEs can learn a smooth and continuous latent space that allows for both effective data reconstruction and the generation of new, realistic data samples.
Implementation
Step 1:-
importing the labarys and Setting up the environment
Code :-
Implementation
Step 1:-
importing the labarys and Setting up the environment
Code :-
import torch # tensor calculations and GPU support
import numpy as nn # numerical operations like arrays
import torch.nn as nn # Building neural networks (layers, models)
from torch.optim import Adam # model training
import matplotlib.pyplot as plt # creating plots and graphs
from torchvision.datasets import MNIST # accessing the MNIST dataset
from torch.utils.data import DataLoader # loading data in batches
import torchvision.transforms as transforms # image transformations (like converting to tensor)
from mpl_toolkits.axes_grid1 import ImageGrid # creating image grids for visualization
from torchvision.utils import save_image, make_grid # saving and displaying images in a grid
Step 2 :-
#create a transofrm to apply to each datapoint
transform = transforms.Compose([transforms.ToTensor()])
# download the MNIST datasets
path = '~/datasets'
# Load training and test datasets with transform applied
train_dataset = MNIST(path, transform=transform, download=True)
test_dataset = MNIST(path, transform=transform, download=True)
batch_size = 100
# Create data loaders for training and testing
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
1. transform = transforms.Compose([transforms.ToTensor()]): This line defines a sequence of transformations to be applied to each image in the MNIST dataset.
transforms.Compose([...]): This is a PyTorch function that allows you to chain multiple image transformations together. The transformations are applied in the order they appear in the list.
transforms.ToTensor(): This is a crucial transformation. It converts the PIL Image or NumPy array representing an image into a PyTorch tensor. This is necessary because PyTorch works with tensors. It also automatically scales the pixel values from the range [0, 255] to the range [0.0, 1.0], which is often beneficial for training neural networks.
2. path = '~/datasets': This line sets the directory where the MNIST dataset will be downloaded and stored. The ~ typically refers to your user's home directory.
train_dataset = MNIST(path, transform=transform, download=True): This line loads the training portion of the MNIST dataset.
MNIST(path, ...): This is a class from torchvision.datasets that provides convenient access to the MNIST dataset.
path=path: Specifies the root directory where the dataset will be stored. If the dataset is not found at this path, it will be downloaded.
transform=transform: Applies the transform you defined earlier to each image in the training dataset. This means each image will be converted to a PyTorch tensor.
download=True: If the MNIST dataset is not already present in the specified path, this argument tells PyTorch to download it from the internet.
test_dataset = MNIST(path, transform=transform, download=True): This line does the same as above, but for the testing portion of the MNIST dataset. It ensures that the test images are also loaded as PyTorch tensors with pixel values scaled between 0 and 1.
3.batch_size = 100: This line defines the number of samples (images) that will be processed together in one batch during training and testing. A batch size of 100 means that the model will see 100 images at a time before updating its weights.
train_loader = DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True): This line creates a data loader for the training dataset.
DataLoader(dataset, ...): This class from torch.utils.data provides an iterable over the dataset, making it easy to access batches of data.
dataset=train_dataset: Specifies the dataset to load data from (in this case, the training dataset).
batch_size=batch_size: Specifies the number of samples to include in each batch.
shuffle=True: This is important for training. Setting shuffle=True shuffles the order of the data at the beginning of each epoch (a full pass through the training data). This helps to prevent the model from learning the order of the data and improves generalization.
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False): This line creates a data loader for the testing dataset.
It uses the test_dataset.
It uses the same batch_size.
shuffle=False: For the test set, we usually don't need to shuffle the data because we are just evaluating the model's performance. The order of the test data doesn't affect the overall evaluation metrics.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu"): This line sets up the device that will be used for training and running your neural network.
torch.device("cuda" if torch.cuda.is_available() else "cpu"): This checks if a CUDA-enabled GPU is available. If it is, the device will be set to "cuda", which means computations will be performed on the GPU for faster processing. If a GPU is not available, the device will be set to "cpu", and computations will be done on the CPU.
100%|██████████| 26.4M/26.4M [00:01<00:00, 13.3MB/s]
100%|██████████| 29.5k/29.5k [00:00<00:00, 212kB/s]
100%|██████████| 4.42M/4.42M [00:01<00:00, 3.88MB/s]
100%|██████████| 5.15k/5.15k [00:00<00:00, 12.2MB/s]
Step 3
Code
dataiter = iter(train_loader) # Get a batch of training data from the DataLoader
image = next(dataiter)
num_samples = 25
sample_images = [image[0] [1,0] for i in range(num_samples)] # Extract grayscale images
fig = plt.figure(figsize=(4, 4))
grid = ImageGrid(fig, 111, nrows_ncols=(5, 5), axes_pad=0.1)
#Plot each image in the grid
for ax, im in zip(grid, sample_images):
ax.imshow(im, cmap='gray') # Display image in grayscale
ax.axis('off')
plt.show()
dataiter = iter(train_loader): This line creates an iterator object from your train_loader. An iterator allows you to access the elements of the data loader one at a time.
image = next(dataiter): This line retrieves the next element from the dataiter. Since train_loader shuffles the data and provides it in batches, image will be a list (or a tuple) containing one batch of data. Typically, for the MNIST dataset loaded with DataLoader, this image will contain two elements:
The first element is a tensor of the batch of input images. Its shape will be (batch_size, 1, 28, 28) in your case (100 samples, 1 channel, 28x28 pixels).
The second element is a tensor of the corresponding labels for those images. Its shape will be (batch_size).
num_samples = 25: This line sets the number of sample images you want to visualize to 25.
sample_images = [image[0][i, 0] for i in range(num_samples)]: This is a list comprehension that extracts the first num_samples grayscale images from the batch. Let's break it down further:
image[0]: This accesses the first element of the image tuple, which is the tensor containing the batch of input images (shape (100, 1, 28, 28)).
[i, 0]: For each i from 0 to 24 (inclusive), this indexing selects the i-th image from the batch. Since the images are grayscale and have a channel dimension of 1, the 0 index after i selects the single channel. So, image[0][i, 0] gives you a 2D tensor of shape (28, 28) representing the i-th grayscale image.
fig = plt.figure(figsize=(4, 4)): This creates a new Matplotlib figure with a specified size of 4x4 inches. This will be the canvas for your grid of images.
grid = ImageGrid(fig, 111, nrows_ncols=(5, 5), axes_pad=0.1): This uses the ImageGrid class from mpl_toolkits.axes_grid1 to create a grid of subplots for displaying the images.
fig: The Matplotlib figure to which the grid will be added.
111: This is a subplot specification (equivalent to (1, 1, 1)), indicating a single grid of subplots that will take up the entire figure.
nrows_ncols=(5, 5): This specifies that the grid should have 5 rows and 5 columns, resulting in a total of 25 subplots, which matches the num_samples you defined.
axes_pad=0.1: This sets the padding between the individual subplots (axes) in the grid to 0.1 inches.
for ax, im in zip(grid, sample_images):: This loop iterates through the grid of subplots (grid) and the list of sample images (sample_images) simultaneously using the zip function. In each iteration, ax will be an individual subplot axis object, and im will be one of the 28x28 grayscale image tensors.
ax.imshow(im, cmap='gray'): This displays the current image (im) on the current subplot (ax).
imshow() is a Matplotlib function for displaying image data.
cmap='gray' specifies that the image should be displayed in grayscale.
ax.axis('off'): This turns off the axis ticks and labels for the current subplot, providing a cleaner visualization of the images.
plt.show(): This function displays the Matplotlib figure containing the grid of sample images.
Step 4:-
code
class VAE(nn.Module):
def __init__(self, input_dim=784, hidden_dim=400, latent_dim=200, device=device):
super(VAE, self).__init__()
#encoder
self.encoder = nn.Sequential(
nn.Linear(input_dim, hidden_dim), #Fully connected layer
nn.LeakyReLU(0.2), #Activation function
nn.Linear (hidden_dim, latent_dim),
nn.LeakyReLU(0.2)
)
#Layers to output mean and log variance for the latent space
self.mean_layer = nn.Linear(latent_dim, 2)
self.logvar_layer = nn.Linear(latent_dim, 2)
#Decoder: Reconstructs the input from latent space
self.decoder = nn.Sequential(
nn.Linear(2, latent_dim), #Fixed typo: nn. Linear to nn.Linear
nn.LeakyReLU(0.2),
nn.Linear (latent_dim, hidden_dim),
nn.LeakyReLU(0.2),
nn.Linear (hidden_dim, input_dim),
nn.Sigmoid()#Output in range [0,1] to match original input
)
def encode(self, x): # Added self
x=self.encoder(x)
mean, logvar=self.mean_layer(x), self.logvar_layer(x) #Predict mean and log variance of latent distribution
return mean, logvar
def reparameterization(self, mean, var): # Added self
epsilon = torch.randn_like(var).to(device) # Sample epsilon from standard normal
z = mean + epsilon * torch.exp(0.5 * var) # Calculate z using reparameterization trick
return z # Return the calculated z
def decode(self, x): # Added self
return self.decoder(x) #Decode latent vector back to original input space
def forward(self, x): # Added self to make it a method of the VAE class
mean, logvar = self.encode(x) #Encode input into latent space
z = self.reparameterization(mean, logvar) #Fixed typo: Logvar to logvar
x_hat = self.decode(z)
return x_hat, mean, logvar #Return reconstructed input, mean, and logvar
1. Class Definition:
class VAE(nn.Module): Defines a class named VAE that inherits from nn.Module, which is the base class for all neural network modules in PyTorch.
2. Initialization (__init__):
def __init__(self, input_dim=784, hidden_dim=400, latent_dim=200, device=device):
Constructor of the VAE class.
input_dim: The dimension of the input data (784 for flattened 28x28 MNIST images).
hidden_dim: The number of neurons in the hidden layers of the encoder and decoder.
latent_dim: The dimension of the latent space (the compressed representation of the input). The code was corrected to use this, but the original code defined the mean and logvar layers to have an output dimension of 2. I've left that as 2, but it's suspicious.
device: The device to which the tensors and model will be moved (CPU or GPU).
super(VAE, self).__init__(): Calls the constructor of the parent class (nn.Module) to properly initialize the VAE as a PyTorch module.
Encoder:
self.encoder: A nn.Sequential module that defines the encoder network.
It consists of two fully connected (nn.Linear) layers with nn.LeakyReLU activation functions.
The encoder maps the input data (input_dim) to a lower-dimensional representation (latent_dim).
Mean and Log Variance Layers:
self.mean_layer: A linear layer that maps the output of the encoder to the mean of the latent distribution. The output dimension is 2.
self.logvar_layer: A linear layer that maps the output of the encoder to the log variance of the latent distribution. The output dimension is 2.
Decoder:
self.decoder: A nn.Sequential module that defines the decoder network.
It mirrors the encoder architecture but maps from the latent space back to the original input space.
The final layer uses a nn.Sigmoid() activation function to ensure the output values are in the range [0, 1], which is suitable for representing pixel intensities.
3. Encoder (encode):
def encode(self, x):
Takes an input x and passes it through the encoder network.
Returns the mean and log variance of the latent distribution.
4. Reparameterization Trick (reparameterization):
def reparameterization(self, mean, var):
Implements the reparameterization trick, which is crucial for training VAEs.
It samples a value z from the latent distribution N(μ,σ2) using:
ϵ∼N(0,I) (sample from a standard normal distribution)
z=μ+ϵ⋅exp(0.5⋅log(σ2)) (Note: The code uses var, which is log variance. So, exp(0.5 * var) effectively calculates the standard deviation)
This allows backpropagation through the sampling process.
Returns the sampled latent vector z.
5. Decoder (decode):
def decode(self, x):
Takes a latent vector x and passes it through the decoder network.
Returns the reconstructed input.
6. Forward Pass (forward):
def forward(self, x):
Defines the forward pass of the VAE.
mean, logvar = self.encode(x): Encodes the input x to get the mean and log variance of the latent distribution.
z = self.reparameterization(mean, logvar): Samples a latent vector z using the reparameterization trick.
x_hat = self.decode(z): Decodes the latent vector z to get the reconstructed input x_hat.
Returns the reconstructed input x_hat, the mean, and the log variance.
Step 5
code
model= VAE().to(device)
optimizer=Adam(model.parameters(),lr=1e-3)
The first line creates your VAE model and places it on the appropriate hardware (GPU or CPU).
The second line sets up the Adam optimizer, telling it which parameters to optimize (the model's parameters) and how quickly to adjust them (the learning rate).
Step 6
code
def loss_function(x, x_hat, mean, log_var):
reconstruction_loss = nn.functional.binary_cross_entropy(x_hat, x, reduction='sum') # Reconstruction loss (how close the output is to the input)
KLD = - 0.5 * torch.sum(1+ log_var - mean.pow(2) - log_var.exp()) # KL Divergence: regularizes the latent space to follow a standard normal distribution
# Total loss = reconstruction loss + KL divergence
return reconstruction_loss + KLD
fig = plt.figure(figsize=(4, 4)): This creates a new Matplotlib figure with a specified size of 4x4 inches. This will be the canvas for your grid of images.
grid = ImageGrid(fig, 111, nrows_ncols=(5, 5), axes_pad=0.1): This uses the ImageGrid class from mpl_toolkits.axes_grid1 to create a grid of subplots for displaying the images.
fig: The Matplotlib figure to which the grid will be added.
111: This is a subplot specification (equivalent to (1, 1, 1)), indicating a single grid of subplots that will take up the entire figure.
nrows_ncols=(5, 5): This specifies that the grid should have 5 rows and 5 columns, resulting in a total of 25 subplots, which matches the num_samples you defined.
axes_pad=0.1: This sets the padding between the individual subplots (axes) in the grid to 0.1 inches.
Step 7:-
Code
def train(model, optimizer, epochs, device, x_dim=784):
model.train()
for epoch in range(epochs):
overall_loss = 0
for batch_idx, (x, _) in enumerate(train_loader):
x = x.view(batch_size, x_dim).to(device) # Flatten the image and move to device
optimizer.zero_grad()
#Forward pass through the VAE
x_hat, mean, log_var = model(x)
loss = loss_function(x, x_hat, mean, log_var) # Compute total loss (reconstruction + KL divergence)
overall_loss += loss.item()
loss.backward() # Backpropagation
optimizer.step() # Update model parameters
# Print average loss per epoch
print("\tEpoch", epoch + 1, "\tAverage Loss: ", overall_loss/(batch_idx*batch_size))
return overall_loss # Return total loss after training
def train(model, optimizer, epochs, device, x_dim=784):
Defines a function named train that takes the following arguments:
model: The VAE model to be trained.
optimizer: The optimizer (e.g., Adam) used for updating the model's parameters.
epochs: The number of times to iterate over the entire training dataset.
device: The device (CPU or GPU) to use for training.
x_dim: The dimension of the input data (784 for flattened 28x28 MNIST images). It has a default value of 784.
model.train():
Sets the VAE model to training mode. This is important because some layers, such as dropout and batch normalization, behave differently during training and evaluation.
for epoch in range(epochs)::
This is the outer loop that iterates over the training data epochs number of times. Each iteration is called an "epoch."
overall_loss = 0:
Initializes a variable overall_loss to 0 at the beginning of each epoch. This variable will accumulate the loss over all batches in the current epoch.
for batch_idx, (x, _) in enumerate(train_loader)::
This is the inner loop that iterates over the training data in batches using the train_loader.
enumerate(train_loader) returns pairs of (batch index, data).
(x, _) unpacks each batch:
x: The batch of input images. The shape of x is (batch_size, 1, 28, 28).
_: The labels for the images. In this code, the labels are ignored (using _) because VAEs are unsupervised models and don't use the labels for training.
x = x.view(batch_size, x_dim).to(device):
Reshapes the input images x to have the shape (batch_size, x_dim).
x.view(batch_size, x_dim) flattens each image in the batch from its original shape (1, 28, 28) to a 1D vector of length 784. This is necessary because the first layer of your encoder is a linear layer.
.to(device) moves the flattened input images to the specified device (CPU or GPU).
optimizer.zero_grad():
Resets the gradients of all the model's parameters to zero. This is a crucial step before calculating the gradients for the current batch. If you don't reset the gradients, PyTorch will accumulate them from previous iterations, which can lead to incorrect gradient updates.
x_hat, mean, log_var = model(x):
Performs a forward pass through the VAE model:
The flattened input images x are passed to the model (your VAE).
The model returns the reconstructed images x_hat, the mean of the latent distribution mean, and the log variance of the latent distribution log_var.
loss = loss_function(x, x_hat, mean, log_var):
Calculates the VAE loss using the loss_function you defined earlier. The loss consists of the reconstruction loss and the KL divergence.
overall_loss += loss.item():
Adds the current batch's loss to the overall_loss. .item() extracts the scalar value of the loss tensor.
loss.backward():
Performs backpropagation: calculates the gradients of the loss with respect to all the model's parameters. These gradients indicate how much each parameter needs to be adjusted to reduce the loss.
optimizer.step():
Updates the model's parameters using the calculated gradients and the optimization algorithm specified by the optimizer (e.g., Adam). This is where the model learns from the data.
print("\tEpoch", epoch + 1, "\tAverage Loss: ", overall_loss / (batch_idx * batch_size)):
Prints the average loss for the current epoch.
epoch + 1: Displays the epoch number (starting from 1).
overall_loss / (batch_idx * batch_size): Calculates the average loss by dividing the total loss for the epoch by the total number of samples processed so far in that epoch. batch_idx starts at 0, so we should use (batch_idx + 1) * batch_size to get the correct number of samples.
return overall_loss:
Returns the total loss after training for all epochs. While the average loss per epoch is printed during training, the final total loss across all epochs is returned. This is somewhat unusual; it might be more useful to return the average loss of the last epoch.
Step 8: -
train(model,optimizer,epochs=50,device=device)
we are passing parameter to train function
Epoch 1 Average Loss: 175.2999954840098
Epoch 2 Average Loss: 157.67686688099957
Epoch 3 Average Loss: 152.73615309369782
Epoch 4 Average Loss: 149.44358592445744
Epoch 5 Average Loss: 147.10354744887312
Epoch 6 Average Loss: 145.28939875378234
Epoch 7 Average Loss: 144.02065579677065
Epoch 8 Average Loss: 143.12622828411938
Epoch 9 Average Loss: 142.21737885081907
Epoch 10 Average Loss: 141.5994077694595
Epoch 11 Average Loss: 140.8984326090359
Epoch 12 Average Loss: 140.3726635049301
Epoch 13 Average Loss: 139.94625259221098
Epoch 14 Average Loss: 139.36170787575648
Epoch 15 Average Loss: 139.05357218084828
Epoch 16 Average Loss: 138.56550262807804
Epoch 17 Average Loss: 138.25561981557806
Epoch 18 Average Loss: 137.89444425605174
Epoch 19 Average Loss: 137.64193343071787
Epoch 20 Average Loss: 137.18865467510955
Epoch 21 Average Loss: 137.04402120395972
Epoch 22 Average Loss: 136.8527757036467
Epoch 23 Average Loss: 136.5826956711707
Epoch 24 Average Loss: 136.42437767372704
Epoch 25 Average Loss: 135.92366537001774
Epoch 26 Average Loss: 135.94024063543407
Epoch 27 Average Loss: 135.87051704011895
Epoch 28 Average Loss: 135.7120373310987
Epoch 29 Average Loss: 135.43031020124687
Epoch 30 Average Loss: 135.34221462854757
Epoch 31 Average Loss: 135.29933124217445
Epoch 32 Average Loss: 134.86808101392947
Epoch 33 Average Loss: 134.90404621308952
Epoch 34 Average Loss: 134.7131630190943
Epoch 35 Average Loss: 134.5444851771181
Epoch 36 Average Loss: 134.36532006469116
Epoch 37 Average Loss: 134.35329986827003
Epoch 38 Average Loss: 134.14271618061352
Epoch 39 Average Loss: 133.93251103727567
Epoch 40 Average Loss: 134.16042473132305
Epoch 41 Average Loss: 133.82699106257826
Epoch 42 Average Loss: 133.7781305920023
Epoch 43 Average Loss: 133.66783494952526
Epoch 44 Average Loss: 133.42391774376566
Epoch 45 Average Loss: 133.47797280950022
Epoch 46 Average Loss: 133.41502753612792
Epoch 47 Average Loss: 133.44754979001462
Epoch 48 Average Loss: 132.9319998011008
Epoch 49 Average Loss: 133.2059645144251
Epoch 50 Average Loss: 133.0060967171849
7967065.193359375
Step 7:-
Code:-
def generate_digit(mean, var):
# Create a latent vector z from the provided mean and variance
z_sample = torch.tensor([[mean, var]], dtype=torch.float).to(device)
x_decoded = model.decode(z_sample) # Decode the latent vector into an image
digit = x_decoded.detach().cpu().reshape(28, 28) # reshape vector to 2d array
# Plot the generated digit
plt.title(f'[{mean}, {var}]')
plt.imshow(digit, cmap='gray')
plt.axis('off')
plt.show()
def generate_digit(mean, logvar)::
Defines a function named generate_digit that now takes mean and logvar (logarithm of variance) as input. This is crucial because VAEs parameterize the latent space distribution using the log of the variance for numerical stability.
std = torch.exp(0.5 * logvar):
Calculates the standard deviation (std) from the log variance (logvar).
The VAE outputs the log variance, so we need to exponentiate half of it to get the standard deviation: σ=exp(0.5⋅log(σ2)).
epsilon = torch.randn_like(std):
Generates a tensor epsilon of the same shape and data type as std, filled with samples from a standard normal distribution (mean 0, variance 1).
This epsilon is used for the reparameterization trick.
z_sample = mean + epsilon * std:
Performs the reparameterization trick to sample a latent vector z_sample from the distribution N(mean,std2).
This is the core of how VAEs generate new samples. Instead of directly sampling from a distribution with a learned mean and variance, we sample from a standard normal distribution and then transform it. This allows us to backpropagate through the sampling process.
z_sample = z_sample.to(device):
Moves the sampled latent vector z_sample to the same device (CPU or GPU) as the VAE model. This ensures that the decoding operation can be performed on the correct device.
x_decoded = model.decode(z_sample):
Decodes the latent vector z_sample using the decoder part of the VAE model.
model.decode() takes the latent vector as input and outputs a reconstructed image (in the original data space).
digit = x_decoded.detach().cpu().view(-1, 28, 28):
Processes the decoded output to make it suitable for display with Matplotlib.
.detach(): Creates a new tensor from x_decoded that is detached from the computation graph. This prevents gradients from being calculated for this tensor, which is necessary because we're just displaying the image and not training the model.
.cpu(): Moves the detached tensor to the CPU. Matplotlib expects data to be in CPU memory.
.view(-1, 28, 28): Reshapes the decoded output into the correct image dimensions. The -1 is a placeholder that tells PyTorch to automatically infer the size of the first dimension (the batch dimension) based on the size of the original tensor and the other specified dimensions (28, 28). This ensures that the output is interpreted as one or more 28x28 images.
plt.title(f'Mean: {mean.mean().item():.2f}, Logvar: {logvar.mean().item():.2f}'):
Sets the title of the plot to display the mean and log variance values used to generate the digit.
mean.mean().item(): Calculates the mean of the mean tensor and extracts it as a Python number using .item().
logvar.mean().item(): Calculates the mean of the logvar tensor and extracts it as a Python number.
:.2f: Formats the numbers to two decimal places for cleaner display.
plt.imshow(digit[0], cmap='gray'):
Displays the generated digit image using Matplotlib.
digit[0]: Selects the first image from the batch. Even if you're generating only one image, the output of the decoder might have a batch dimension. We need to index it to get the single image.
cmap='gray': Specifies that the image should be displayed in grayscale.
plt.axis('off'):
Turns off the axis labels and ticks for the plot, providing a cleaner visualization.
plt.show():
Displays the plot containing the generated digit.
Step 7:-
Code
#img1:
generate_digit(0.0,1.0),generate_digit(1.0,0.0)
Step 8:-
Code
def plot_latent_space(model, scale=5.0, n=25, digit_size=28, figsize=15):
# display a n*n 2D manifold of digits
figure = np.zeros((digit_size * n, digit_size * n))
# construct a grid
grid_x = np.linspace(-scale, scale, n)
grid_y = np.linspace(-scale, scale, n) [ : :- 1]
for i, yi in enumerate(grid_y):
for j, xi in enumerate(grid_x):
# Create a latent vector z from grid coordinates
z_sample = torch.tensor([[xi, yi]], dtype=torch.float).to(device)
x_decoded = model.decode(z_sample) # Decode the latent vector into an image
digit = x_decoded[0].detach().cpu().reshape(digit_size, digit_size)
# Place the digit into the correct position in the grid figure
figure[i * digit_size : (i + 1) * digit_size, j * digit_size : (j + 1) * digit_size, ] = digit
plt.figure(figsize=(figsize, figsize))
plt.title('VAE Latent Space Visualization')
start_range = digit_size // 2
end_range = n * digit_size + start_range
pixel_range = np.arange(start_range, end_range, digit_size)
sample_range_x = np.round(grid_x, 1)
sample_range_y = np.round(grid_y, 1)
plt.xticks(pixel_range, sample_range_x)
plt.yticks(pixel_range, sample_range_y)
plt.xlabel("mean, z [0]")
plt.ylabel("var, z [1]")
plt.imshow(figure, cmap="Greys_r") # Show the image grid in grayscale
plt.show()
1. Function Definition and Initialization:
def plot_latent_space(model, scale=5.0, n=25, digit_size=28, figsize=15):
Defines the function plot_latent_space with the following parameters:
model: Your trained VAE model.
scale: The range of values in the latent space to sample (default: 5.0). This determines how far you go from the center of the latent space.
n: The number of points along each axis of the grid (default: 25). This will result in an n x n grid of generated digits.
digit_size: The size of the generated digit images (default: 28 pixels).
figsize: The size of the figure to be plotted (default: 15 inches).
figure = np.zeros((digit_size * n, digit_size * n))
Creates an empty NumPy array called figure to store the grid of generated digits.
The size of the array is (digit_size * n, digit_size * n). For example, if n is 25 and digit_size is 28, the array will be 700x700. This array will hold the individual digit images arranged in a grid.
2. Creating the Latent Space Grid:
grid_x = np.linspace(-scale, scale, n)
Creates a 1D NumPy array grid_x containing n evenly spaced values between -scale and scale. This represents the x-coordinates in the 2D latent space.
grid_y = np.linspace(-scale, scale, n)[::-1]
Creates a 1D NumPy array grid_y similar to grid_x, but the values are in reverse order ([::-1]). This represents the y-coordinates in the 2D latent space. Reversing grid_y is a common practice to make the origin of the grid (0,0) appear in the top-left corner of the plot, which aligns with how images are typically indexed.
3. Generating Digits and Filling the Grid:
The code then iterates through the grid coordinates using nested loops:
for i, yi in enumerate(grid_y): Iterates through the grid_y values. i is the row index.
for j, xi in enumerate(grid_x): Iterates through the grid_x values. j is the column index.
Inside the loops:
z_sample = torch.tensor([[xi, yi]], dtype=torch.float).to(device)
Creates a latent vector z_sample as a PyTorch tensor, using the current xi and yi values from the grid. Important Note: This code assumes a 2-dimensional latent space. If your VAE has a higher-dimensional latent space, this part needs to be modified.
The tensor is moved to the appropriate device (CPU or GPU).
x_decoded = model.decode(z_sample)
Decodes the latent vector z_sample using the VAE's decoder. This generates a reconstructed image.
digit = x_decoded[0].detach().cpu().reshape(digit_size, digit_size)
Processes the decoded output:
x_decoded[0]: Selects the first (and presumably only) image from the output.
.detach(): Detaches the tensor from the computation graph.
.cpu(): Moves the tensor to the CPU for use with Matplotlib.
.reshape(digit_size, digit_size): Reshapes the output into a square image.
figure[i * digit_size: (i + 1) * digit_size, j * digit_size: (j + 1) * digit_size, ] = digit
Places the generated digit image into the correct position in the figure array.
This uses array slicing to put the digit (which is a digit_size x digit_size image) into the i-th row and j-th column of the larger figure array.
4. Plotting the Grid:
plt.figure(figsize=(figsize, figsize))
Creates a Matplotlib figure with the specified figsize.
plt.title('VAE Latent Space Visualization')
Sets the title of the plot.
start_range = digit_size // 2
Calculates the starting pixel position for the x and y axis labels.
end_range = n * digit_size + start_range
Calculates the ending pixel position for the x and y axis labels.
pixel_range = np.arange(start_range, end_range, digit_size)
Creates an array of pixel positions for the x and y axis ticks.
sample_range_x = np.round(grid_x, 1)
Rounds the grid_x values to one decimal place for cleaner display as x-axis labels.
sample_range_y = np.round(grid_y, 1)
Rounds the grid_y values to one decimal place for display as y-axis labels.
plt.xticks(pixel_range, sample_range_x)
Sets the x-axis tick positions and labels.
plt.yticks(pixel_range, sample_range_y)
Sets the y-axis tick positions and labels.
plt.xlabel("mean, z [0]")
Sets the label for the x-axis.
plt.ylabel("var, z [1]")
Sets the label for the y-axis. Important Note: The y-axis label says "var", but the code uses grid_y, which represents a value in the latent space. It's more appropriate to label this as "mean, z [1]" if you are using the latent space coordinates directly. If you are plotting the log variance, then the label should be "logvar, z[1]".
plt.imshow(figure, cmap="Greys_r")
Displays the figure array as an image, using the "Greys_r" colormap (reversed grayscale).
plt.show()
Displays the plot.
Step 9:-
Code
!pip install numpy
import numpy as np # Import numpy and assign it to the alias 'np'
def plot_latent_space(model, scale=5.0, n=25, digit_size=28, figsize=15):
# display a n*n 2D manifold of digits
figure = np.zeros((digit_size * n, digit_size * n)) # Now np is defined and can be used
# construct a grid
grid_x = np.linspace(-scale, scale, n)
grid_y = np.linspace(-scale, scale, n) [ : :- 1]
for i, yi in enumerate(grid_y):
for j, xi in enumerate(grid_x):
# Create a latent vector z from grid coordinates
z_sample = torch.tensor([[xi, yi]], dtype=torch.float).to(device)
x_decoded = model.decode(z_sample) # Decode the latent vector into an image
digit = x_decoded[0].detach().cpu().reshape(digit_size, digit_size)
# Place the digit into the correct position in the grid figure
figure[i * digit_size : (i + 1) * digit_size, j * digit_size : (j + 1) * digit_size, ] = digit
plt.figure(figsize=(figsize, figsize))
plt.title('VAE Latent Space Visualization')
start_range = digit_size // 2
end_range = n * digit_size + start_range
pixel_range = np.arange(start_range, end_range, digit_size)
sample_range_x = np.round(grid_x, 1)
sample_range_y = np.round(grid_y, 1)
plt.xticks(pixel_range, sample_range_x)
plt.yticks(pixel_range, sample_range_y)
plt.xlabel("mean, z [0]")
plt.ylabel("var, z [1]")
plt.imshow(figure, cmap="Greys_r") # Show the image grid in grayscale
plt.show()
Function Definition:
Defines the plot_latent_space function with parameters for the VAE model, the scale of the latent space, the number of grid points, the digit size, and the figure size.
Create Empty Grid:
figure = np.zeros((digit_size * n, digit_size * n)): Creates a NumPy array filled with zeros to store the grid of generated digits. The dimensions of the grid are calculated based on the digit size and the number of grid points.
Generate Latent Space Coordinates:
grid_x = np.linspace(-scale, scale, n): Creates a 1D array of evenly spaced values for the x-coordinates in the latent space.
grid_y = np.linspace(-scale, scale, n)[::-1]: Creates a similar array for the y-coordinates, but in reverse order.
Iterate and Generate Digits:
The code iterates through the x and y coordinates of the grid. For each coordinate pair:
It creates a latent vector z_sample using torch.tensor.
It decodes this latent vector into an image using model.decode().
It reshapes the decoded output to the correct digit size.
It places the generated digit image into the corresponding position in the figure array.
Plot the Grid:
The code uses Matplotlib (plt) to create a figure and display the grid of digits:
It sets the figure size and title.
It calculates the tick positions and labels for the axes.
It displays the figure array as an image using plt.imshow().
It sets the x and y axis labels.
It shows the plot.
Step 9:-
Code
plot_latent_space(model, scale=1.0)
we are passing the parameter