Rotation-invariant convolutional VAE

This repository trains a 1-channel convolutional VAE with windowed self-attention in the encoder and decoder (vae.py). Training scripts support PyTorch DDP for multi-GPU runs (including Slurm-style launches with srun).

Features

• ConvVAE with local 2D window attention and reparameterized Gaussian latent space.
• Rotation-aware objective : paired forward passes on original and augmented images plus latent alignment.
• YAML-driven config via YParams (default_500.yaml is an example).
• Slurm batch script example: VAE_gpu_500_regular_512.sh (site-specific paths; adjust before use).

Requirements

• Python 3.10+ (3.10 is used on the reference Perlmutter workflow).
• PyTorch with CUDA if you train on GPU (install from pytorch.org if the default wheel does not match your CUDA driver).

Install

From the repository root:

python -m venv .venv
source .venv/bin/activate   # Windows: .venv\Scripts\activate
pip install -U pip
pip install -e .

For a CUDA-enabled PyTorch build, install torch/torchvision first per the official matrix, then:

pip install -e ".[dev]"   # optional: ruff

Configuration

1. YAML — Copy default_500.yaml and set at least:
   • image_dir: recursive folder of grayscale-friendly images (.png, .jpg, …).
   • base_dir, logs_dir, checkpoint_dir, etc. Paths may include {latent_dim}; the training scripts substitute that from --latent_dim.
2. .env — Optional convenience variables (MASTER_ADDR, MASTER_PORT, etc.) for distributed jobs. Export them in your shell or Slurm prolog if you do not load dotenv in Python (the training entrypoints do not call python-dotenv by default).

Run training

Single GPU (no Slurm), from repo root:

python training_2.py \
  --yaml_config ./default_500.yaml \
  --config default \
  --latent_dim 512

Adjust image_dir and output paths in the YAML first; defaults in the shipped file point at HPC scratch paths.

Slurm / multi-node: See VAE_gpu_500_regular_512.sh for an srun example. Set MASTER_ADDR, MASTER_PORT, and GPU bindings the way your cluster expects.

Scripts

File	Role
training_2.py	Main DDP trainer with Trainer.train() loop (mean MSE/KL variant).
VAE_gpu_500_regular.py	Alternate Slurm-oriented setup (Perlmutter-style env parsing); ensure a train loop exists on Trainer if you use this entrypoint as-is.
data_loader.py	ImageFolder-style recursive loading, 80/20 split, DistributedSampler.
vae.py	Model and simple_autoencoder(params) factory.
YParams.py	Load a named section from YAML into params.

Project layout

├── default_500.yaml          # Example hyperparameters and paths
├── data_loader.py
├── vae.py
├── YParams.py
├── training_2.py
├── training_2-Copy1.py
├── VAE_gpu_500_regular.py
├── VAE_gpu_500_regular_512.sh
├── pyproject.toml
├── .env
└── README.md