Unverified Commit b4220e97 authored by dg845's avatar dg845 Committed by GitHub
Browse files

Add Community Example Consistency Training Script (#6717)



* initial commit for unconditional/class-conditional consistency training script

* make style

* Add entry for consistency training script in community README.

* Move consistency training script from community to research_projects/consistency_training

* Add requirements.txt and README to research_projects/consistency_training directory.

* Manually revert community README changes for consistency training.

* Fix path to script after moving script to research projects.

* Add option to load U-Net weights from pretrained model.

---------
Co-authored-by: default avatarPatrick von Platen <patrick.v.platen@gmail.com>
parent dc85b578
# Consistency Training
`train_cm_ct_unconditional.py` trains a consistency model (CM) from scratch following the consistency training (CT) algorithm introduced in [Consistency Models](https://arxiv.org/abs/2303.01469) and refined in [Improved Techniques for Training Consistency Models](https://arxiv.org/abs/2310.14189). Both unconditional and class-conditional training are supported.
A usage example is as follows:
```bash
accelerate launch examples/research_projects/consistency_training/train_cm_ct_unconditional.py \
--dataset_name="cifar10" \
--dataset_image_column_name="img" \
--output_dir="/path/to/output/dir" \
--mixed_precision=fp16 \
--resolution=32 \
--max_train_steps=1000 --max_train_samples=10000 \
--dataloader_num_workers=8 \
--noise_precond_type="cm" --input_precond_type="cm" \
--train_batch_size=4 \
--learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
--use_8bit_adam \
--use_ema \
--validation_steps=100 --eval_batch_size=4 \
--checkpointing_steps=100 --checkpoints_total_limit=10 \
--class_conditional --num_classes=10 \
```
\ No newline at end of file
accelerate>=0.16.0
torchvision
transformers>=4.25.1
ftfy
tensorboard
Jinja2
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment