TRACT: Denoising Diffusion Models with Transitive Closure Time-Distillation
This software project accompanies the research paper TRACT: Denoising Diffusion Models with Transitive Closure Time-Distillation
Citation:
@article{berthelot2023tract,
title={TRACT: Denoising Diffusion Models with Transitive Closure Time-Distillation},
author={Berthelot, David and Autef, Arnaud and Lin, Jierui and Yap, Dian Ang and Zhai, Shuangfei and Hu, Siyuan and Zheng, Daniel and Talbott, Walter and Gu, Eric},
journal={arXiv preprint arXiv:2303.04248},
year={2023}
}
Setup
Git clone with --recurse-submodules
to initialize EDM submodule properly.
Setup environment variables:
export ML_DATA=~/Data/DDPM-Images
export PYTHONPATH=$PYTHONPATH:.
Then run
sudo apt install python3.8-dev python3.8-venv python3-dev -y
Set up a virtualenv
python3.8 -m venv ~/tract_venv
source ~/tract_venv/bin/activate
or via pyenv
pyenv install 3.8.0
pyenv virtualenv 3.8.0 tract_venv
pyenv local tract_venv
then upgrade pip
pip install --upgrade pip
Install pip pkgs
pip install -r requirements.txt -f https://download.pytorch.org/whl/torch_stable.html
Setup data
Please follow README to setup datasets.
Set up teacher models
Please follow README to setup teacher checkpoints.
Set up EDM
For running with NVIDIA's Elucidated model (EDM), ensure the edm/
submodule has been initialized properly.
Real activation stats for FID
Please follow the Real activation statistics
section in README in order to compute and save the real activation statistics to be used in FID evaluation.
Training
The below commands will reproduce results from our paper when run on a cluster of 8 NVIDIA A100 or V100 GPUs.
Example: Run TC distillation on Cifar10 using distillation time schedule: 1024, 32, 1.
python tc_distill.py --dataset=cifar10 --time_schedule=1024,32,1 --fid_len=50000 --report_fid_len=8M --report_img_len=8M --train_len=96M
Example: Run TC distillation on Cifar10 using EDM teacher
python tc_distill_edm.py --dataset=cifar10 --time_schedule=40,1 --fid_len=50000 --report_fid_len=8M --report_img_len=8M --train_len=96M --batch=512
Getting help
python tc_distill.py --help
Tensorboard outputs are generated in a dir like e/DATASET/MODEL/EXP_NAME/tb/
. For example, you can start tensorboard to view metrics like
tensorboard --logdir e/cifar10/EluDDIM05TCMultiStepx0\(EluUNet\)/aug_prob@0.0_batch@8_dropout@0.0_ema@0.9996_lr@0.001_lr_warmup@None_res@32_sema@0.1_time_schedule@40,1_timesteps@40/tb/ --bind_all
Evaluation
Please follow README to run FID evaluation.