Vision Transformer Autoencoder (ViT AE)

This repository contains a distributed training setup for a grayscale image autoencoder based on a Vision Transformer (ViT) backbone. It supports SLURM-based multi-GPU training using PyTorch DDP and includes modular components for model, data loading, configuration, and training management.

Directory Structure

├── train.py                     # Main training entrypoint
├── trainer.py                   # Trainer class for training, validation, checkpointing
├── models/
│   └── vt.py                    # ViT-based autoencoder and attention modules
├── data/
│   └── data_loader.py           # Masked dataset and DataLoader setup
├── configs/
│   ├── vit_config.yaml          # YAML config with latent_dim placeholder
│   └── YParams.py               # Config loader class
├── utils/
│   ├── logger.py                # Logger setup by rank
│   └── ddp_setup.py             # DDP init from SLURM env
├── scripts/
│   └── VT_64.sh 		 # SLURM job submission script

Usage

srun --ntasks=1 --cpus-per-task=32 --gpus-per-task=1 \
python train.py \
--yaml_config=/path/to/ViT_project/configs/vit_config.yaml \
--config=default \
--latent_dim=64 \
--root_dir=/path/to/ViT_project/ \
--run_num=ViT_latent_64