An implementation of the Soft-Intro VAE with the beta-TC-VAE disentanglement term.
| File | Description |
|---|---|
| main.py | Entry point for training script, parses command line arguments |
| train.py | Train loop with various parameters to change behaviour |
| models.py | Implements autoencoder with encoder and decoder modules |
| dataset.py | Implements the data loader with preprocessing and downsampling |
| utils.py | Loss functions and miscellaneous helper functions |
| ModelVis.ipynb | Visualizes the output of a specified model |
| tests | Test cases, including a test run of the training loop |
| README.md | This file |
python main.py --dataset ukiyo_e64 --device 0 --lr 2e-4 --num_epochs 250 --beta_kl 0.5 --beta_rec 0.75 --beta_neg 512 --z_dim 128 --batch_size 64 --amp --arch conv
Note that this requires a GPU to work. To train on the CPU use --device -1.
To run the tests use pytest:
pytest ./tests