Skip to content

Commit b4220e9

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: Patrick von Platen <[email protected]>
1 parent dc85b57 commit b4220e9

File tree

3 files changed

+1462
-0
lines changed

3 files changed

+1462
-0
lines changed
Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# Consistency Training
2+
3+
`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.
4+
5+
A usage example is as follows:
6+
7+
```bash
8+
accelerate launch examples/research_projects/consistency_training/train_cm_ct_unconditional.py \
9+
--dataset_name="cifar10" \
10+
--dataset_image_column_name="img" \
11+
--output_dir="/path/to/output/dir" \
12+
--mixed_precision=fp16 \
13+
--resolution=32 \
14+
--max_train_steps=1000 --max_train_samples=10000 \
15+
--dataloader_num_workers=8 \
16+
--noise_precond_type="cm" --input_precond_type="cm" \
17+
--train_batch_size=4 \
18+
--learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
19+
--use_8bit_adam \
20+
--use_ema \
21+
--validation_steps=100 --eval_batch_size=4 \
22+
--checkpointing_steps=100 --checkpoints_total_limit=10 \
23+
--class_conditional --num_classes=10 \
24+
```
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
accelerate>=0.16.0
2+
torchvision
3+
transformers>=4.25.1
4+
ftfy
5+
tensorboard
6+
Jinja2

0 commit comments

Comments
 (0)