VitVqGanVae

A flexible and extensible framework for benchmarking data reconstruction using various Vector Quantization techniques with VAEs and GANs.

📖 Overview

This project provides a comprehensive benchmark for evaluating the performance of different Vector Quantization (VQ) methods within a VAE/GAN framework. It is designed to be modular and easily extensible, allowing researchers and developers to experiment with novel VQ techniques, models, and datasets.

🚀 Setup

🐍 Python

python3 -m venv .env
source .env/bin/activate
pip install -U pip wheel
pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cu129
pip install pytest # for development only
pip install .

🐾 Conda

For easier reproducibility, using conda is a recommended way to setup the environment.

bash setup.sh vitvqganvae
conda activate vitvqganvae

🐳 Docker

We provide docker image file in folder docker and services in docker-compose.yaml. All images are based on kohido/base_dl_cuda129:v0.0.7. In this image all required packages are built.

☸️ Kubernetes

We provide Kubernetes .yaml files in folder deploy. The image is also built upon the image kohido/base_dl_cuda129:v0.0.7. For multi-gpu/multi-node training with Kubernetes, we provide corresponding .yaml in folder deploy/mg.

📊 Logging with Wandb

This project uses wandb for logging metrics. You will need to log into wandb using their token. Perform wandb login then provide the token to login wandb.

For Docker user, you will need to create a file .env, so that the docker-compose will read the wandb token from that file to login wandb inside the container. A template for .env is as follows:

WANDB_PROJECT_NAME=<your_project_name>
WANDB_ENTITY=<your_project_entity>
WANDB_API_KEY=<your_wandb_api_key>

For Kubernetes user, you will need to create a file .env, which has the same template as above and add that file to a secret and then modify the secret name in .yaml file. An example is shown as follows:

env:
- name: GITHUB_TOKEN
    valueFrom:
    secretKeyRef:
        name: <your-secret-name>
        key: GITHUB_TOKEN
- name: WANDB_API_KEY
    valueFrom:
    secretKeyRef:
        name: <your-secret-name>
        key: WANDB_API_KEY

🎬 Results & Demo

We provide checkpoints and demo for using checkpoint via GoogleColab notebook. More checkpoints with different datasets can be found here.

VQVAE Result on CelebA

📚 Dataset	📏 Size	📉 Loss	🧠 Method	🏆 Checkpoint	🚀 Demo
CelebA	64 x 64	$$1.4\times 10^{-4}$$	🟩 VQ	🤗 HF
CelebA	64 x 64	$$2.4\times 10^{-4}$$	🟩 RVQ	🤗 HF	📓 Colab
CelebA	64 x 64	$$9\times 10^{-5}$$	🟩 GRVQ	🤗 HF	📓 Colab
CelebA	64 x 64	$$7\times 10^{-5}$$	🟩 MHVQ	🤗 HF
CelebA	128 x 128	$$6\times 10^{-4}$$	🟩 RVQ	🤗 HF
CelebA	128 x 128	$$4\times 10^{-4}$$	🟩 GRVQ	🤗 HF	📓 Colab
CelebA	128 x 128	$$4\times 10^{-4}$$	🟩 MHVQ	🤗 HF	📓 Colab
FFHQ	256 x 256	$$4.1\times 10^{-3}$$	🟩 MHVQ	🤗 HF	📓 Colab

🏋️‍♂️ Train

For data preparation, please refer to DATA.md

🐍 Python Environment & Conda

This project is built upon Accelerate to easily perform single, multi-gpus, and multi-nodes training. To train a Residual Vector Quantization (RVQ) VAE on the CelebA dataset, you can use the following command

accelerate launch \
    --mixed_precision=no \
    --num_processes=1 \
    --num_machines=1 \
    --dynamo_backend=no \
    main.py \
    --config config/celeba_vqvae_rvq.yaml \
    --train \
    trainer_kwargs.use_wandb_tracking=True

🐳 Docker

To train using Docker the compose fragments are split under docker/compose/ and a small helper script is provided to assemble them.

Recommended (helper script):

# list available fragments / service names
./compose-up.sh list

# bring up a single service (example)
./compose-up.sh up celeba128_vqvae_rvq --build --force-recreate

# bring up all fragments without a specific service
./compose-up.sh up

If you prefer to run docker compose directly, pass the top-level compose file and the fragments you want:

docker compose -f docker-compose.yaml -f docker/compose/img/celeba/celeba128_vqvae_rvq.yaml up --build --force-recreate

Note: ./compose-up.sh searches docker/compose recursively and will include all *.yaml/*.yml fragments it finds, so you don't need to list them manually.

☸️ Kubernetes

To train using Kubernetes, after setting up the environment, use the following command to run a pod.

cd deploy
kubectl apply -f <your_yaml_file>

To train a Residual Vector Quantization (RVQ) VAE on the CelebA dataset, you can use the following command

cd deploy
kubectl apply -f celeba_vqvae_rvq.yaml

🧪 Testing

To perform testing all core functions of vitvqganvae, please refer to folder test, which contains *.py test files. To easily test them all, perform the following command:

bash test.sh

🙏 Acknowledgement

• vector-quantize-pytorch
• x-transformers