Simple implementations of different types of Autoencoders in PyTorch

Overview

This repository contains PyTorch implementations of different type of autoencoders.

Try the interactive demo: https://julienexr-autoencooder-mnist-app-sehuso.streamlit.app/

Table of contents

• Quick project layout
• Modules
• Usage instructions
• Small tips

Quick project layout

• modules/autoencoder/ae.py -- AE model (encoder + decoder).
• modules/autoencoder/training.py -- AE training loop.
• modules/vae/vae.py -- VAE model (probabilistic encoder producing μ and logvar + decoder). Two encoder/decoder variants exist (default and pp).
• modules/vae/training.py -- VAE training loop.
• modules/cvae/cvae.py -- CVAE model (VAE conditioned on class labels).
• modules/cvae/training.py -- CVAE training loop.
• modules/vq_vae/vq_vae.py -- VQ-VAE model (encoder + decoder + vector quantizer).
• modules/vq_vae/tranformer_prior.py -- Transformer prior over VQ codebook indices.
• modules/vq_vae/training.py -- VQ-VAE + Transformer prior training loops.
• main.py -- example entry points. By default it runs AE training then VAE training (see note below).
• src/visualization.py -- Visualizer helper used by training.py to save reconstructions, PCA plots, interpolations and noise samples into visu/.
• src/data.py -- MNIST dataloader helpers.
• modules/autoencoder/ -- AE module README and usage notes.
• modules/vae/ -- VAE module README and usage notes.
• modules/cvae/ -- CVAE module README and usage notes.
• modules/vq_vae/ -- VQ-VAE + Transformer prior README and usage notes.
• models/AE/, models/VAE/, models/CVAE/, models/VQ-VAE/ -- expected checkpoints are saved here (encoder/decoder/codebook state dicts).

Modules

Each model has its own README with full explanations and figures :

• Autoencoder (AE): modules/autoencoder/README.md
• Variational Autoencoder (VAE): modules/vae/README.md
• Conditional VAE (CVAE): modules/cvae/README.md
• VQ-VAE + Transformer Prior: modules/vq_vae/README.md

Usage instructions

1. Clone the repository

git clone git@github.com:JulienExr/Autoencoder-MNIST.git
(HTTPS : git clone https://github.com/JulienExr/Autoencoder-MNIST.git)
cd Autoencoder-MNIST

2. Create and activate a virtual environment :

python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

3. Prepare data

MNIST will be downloaded automatically by torchvision into ./data when you run training.

4. Run training (CLI)

main.py you can choose the model, dataset, and latent dimension:

You can visualize training outputs by adding the --visualize flag.

python main.py --model AE --dataset mnist --latent_dim 256
python main.py --model VAE --dataset mnist --latent_dim 32

Dataset options:

• mnist (default)
• fashion_mnist (more challenging, grayscale clothing items)

Example with Fashion-MNIST:

python main.py --model VAE --dataset fashion_mnist --latent_dim 128

4. Outputs

• Model checkpoints are saved under models/AE/, models/VAE/, models/CVAE/, and models/VQ-VAE/.
• Visual outputs are saved under visu/<dataset>_<model>/ with subfolders recon, pca, umap, interp, and noise.
• If you want to try with your own model saved on models/* use : streamlit run app.py.