This repository contains a trained Variational Autoencoder (VAE) model on the CelebA dataset. The model is designed to encode and decode facial images, enabling tasks such as image reconstruction, latent space interpolation, and attribute manipulation.
The training process was tracked using
Weights and Biases. You can view the full training logs and metrics
here.
1import torch
2from vae_model import VAE # Ensure the VAE class is defined in vae_model.py
3
4# Define the latent dimension
5latent_dim = 200
6
7# Initialize the model
8model = VAE(latent_dim=latent_dim)
9
10# Load the trained weights
11model_path = "./vae_celeba_latent_200_epochs_10_batch_64_subset_80000.pth"
12model.load_state_dict(torch.load(model_path))
13model.eval()
Below is a reconstruction example where the first row represents the original images, and the second row represents the reconstructed images:
This project is licensed under the MIT License. See the LICENSE file for details.