Awesome
A Continuous Time Framework for Discrete Denoising Models
Install
conda env create --file tldr_env.yml
Sampling
To reproduce the FID score reported in the paper, first download the pre-trained model from https://www.dropbox.com/scl/fo/zmwsav82kgqtc0tzgpj3l/h?dl=0&rlkey=k6d2bp73k4ifavcg9ldjhgu0s
Then go to config/eval/cifar10.py and change the model_location field to point to the cifar10 pytorch checkpoint.
Also change model_config_location to the config included with the pytorch checkpoint.
We provide a sampling script to sample 50000 images from the model.
Go to scripts/sample.py, and change save_samples_path to a location where samples can be saved.
Then run
python scripts/sample.py
Then to compute the FID score, you will also need to download the CIFAR10 dataset. Once you have the CIFAR10 pngs in another folder you can compute the FID score using
python -m pytorch_fid --device cuda:0 path/to/tauLDR_samples path/to/cifar10pngs
ELBO evaluation
To reproduce the ELBO value reported in the paper for the CIFAR10 model, obtain the checkpoint as above.
Then go to config/eval/cifar10_elbo.py and change config.experiment_dir to point to a directory that has the following structure
experiment_dir
- checkpoints
- ckpt_0001999999.pt
- config
- config_001.yaml
using the files downloaded from the dropbox. Also change config.checkpoint_path to point to /path/to/experiment_dir/checkpoints/ckpt_0001999999.pt.
Change config.cifar10_path to point to somewhere where the CIFAR10 dataset can be downloaded (it will be automatically downloaded to that location).
Finally, run the following command
python elbo_evaluation.py
This will save the ELBO value in a eval folder within the experiment_dir. The ELBO value is written in the file neg_elbo, the first 0 can be ignored and the second number is the negative ELBO averaged over the pixels. It should be around 3.59.
Notebooks
To generate CIFAR10 samples, open the notebooks/image.ipynb notebook.
Change the paths at the top of the config/eval/cifar10.py config file to point to a folder where CIFAR10 can be downloaded and the paths to the model and config downloaded from the dropbox link.
To generate piano samples, open the notebooks/piano.ipynb notebook.
Change the paths at the top of the config/eval/piano.py config file to point to the dataset downloaded from the dropbox link as well as the model weights and config file.
The sampling settings can be set in the config files, switching between standard tau-leaping and with predictor-corrector steps.
Training
CIFAR10
The CIFAR10 model can be trained using
python train.py cifar10
Paths to store the output and to download the CIFAR10 dataset should be set in the training config, config/train/cifar10.py.
To train the model over multiple GPUs, use
python dist_train.py cifar10
with settings found in the config/train/cifar10_distributed.py config file.
Piano
The piano model can be trained using
python train.py piano
Paths to store the output and to the dataset downloaded from the dropbox link should be set in config/train/piano.py.
Audio Samples
These are 4 pairs of audio samples. The first is a music sequence generated by the model conditioned on the first 2 bars (~2.5 secs) of the piece. The second is the ground truth song from the test dataset.