Skip to content

Commit 86eb2d6

Browse files
daniel-gallodaniel-gallo
committed
Update README.md
1 parent 48e85d5 commit 86eb2d6

File tree

1 file changed

+126
-175
lines changed

1 file changed

+126
-175
lines changed

‎README.md

Lines changed: 126 additions & 175 deletions
Original file line numberDiff line numberDiff line change
@@ -1,202 +1,153 @@
1-
# Early Stopping Diffusion
2-
3-
[![Status](https://github.com/razvanmatisan/early-stopping-diffusion/actions/workflows/python.yml/badge.svg)](https://github.com/razvanmatisan/early-stopping-diffusion/actions/workflows/python.yml)
4-
5-
## Dev instructions
6-
The first time:
7-
1. Configure a virtual environment
8-
```bash
9-
python -m venv venv
10-
source venv/bin/activate
11-
pip install -r src/requirements.txt
12-
```
13-
14-
On Windows:
15-
```bash
16-
python -m venv venv
17-
Set-ExecutionPolicy Unrestricted -Scope Process
18-
venv\Scripts\Activate
19-
pip install -r src\requirements.txt
20-
```
21-
22-
2. Run `pre-commit install`
23-
3. If you use VSCode, it might be helpful to add the following to `settings.json`:
24-
```json
25-
"[python]": {
26-
"editor.formatOnSave": true,
27-
"editor.defaultFormatter": "charliermarsh.ruff",
28-
"editor.codeActionsOnSave": {
29-
"source.fixAll": "explicit",
30-
"source.organizeImports": "explicit"
31-
}
32-
},
33-
```
34-
35-
After that, be sure that all the tests are passing before a commit. Otherwise, GitHub Actions will complain ;) You can check by running
1+
# DuoDiff - Accelerating Diffusion Models with a Dual-Backbone Approach
2+
The first step should be configuring a proper environment. On UNIX,
363
```bash
37-
cd src
38-
python -m pytest tests
39-
```
40-
41-
## Repository structure
42-
- `demos/`: Demos for visualising early stopping diffusion.
43-
- `src/`: Code.
44-
- `CMMD_evaluation/`: Code for calculating the CMMD score of generated samples.
45-
- `benchmarking/`: Files for benchmarking models.
46-
- `checkpoints/`: Checkpoints.
47-
- `datasets/`: Dataset-specific dataloaders.
48-
- `models/`: Model definitions.
49-
- `scripts/`: Scripts for training, generation, evaluation and benchmarking.
50-
- `snellius/`: Files for running experiments on Snellius.
51-
- `tests/`: Unit tests.
52-
- `utils/`:
53-
- `field_utils.py` Getters for time and space embeddings.
54-
- `train_utils.py` Getters for models, optimizers, dataloaders, etc.
55-
- `FID_evaluation.py`: Code for calculating the FID score of generated samples.
56-
- `compute_gflops_and_layer_ratio.py`: Code for computing the average layer ratio and theoretical GFLOP.
57-
- `ddpm_core.py`: Code of the DDPM sampler.
58-
- `get_flops.py`: Code for computing theoretical GFLOPs.
59-
- `requirements.txt`: File with requirements for setting up the virtual environment.
60-
- `train.py`: Code for training models.
61-
- `blogpost.md`: Blogpost about the project.
62-
63-
## Running experiments
64-
All of the experiments should be run inside `src` directory.
4+
python -m venv venv
5+
source venv/bin/activate
6+
pip install -r requirements.txt
657
```
66-
cd src
8+
On Windows,
9+
```powershell
10+
python -m venv venv
11+
Set-ExecutionPolicy Unrestricted -Scope Process
12+
venv\Scripts\Activate
13+
pip install -r src\requirements.txt
6714
```
6815

69-
### Training
70-
Training the models is done using `train.py` script.
71-
Full specification of the script can be found with `python train.py --help` command. Below are sample commands for running training with only the essential arguments.
16+
## Training
17+
In this section, we will see how to train early-exit models and DuoDiff on the CelebA dataset. Training on other datasets is straightforward, and we recommend checking the different options in `main.py`.
7218

73-
#### UViT backbone
19+
<details>
20+
<summary>The first step is to obtain a full-model backbone.</summary>
7421

75-
Command for training the UViT backbone.
76-
```shell
77-
bash scripts/train_uvit.sh
78-
```
79-
or
80-
```shell
22+
```bash
8123
python train.py \
24+
--n_steps 500000 \
25+
--batch_size 128 \
26+
--log_path "${log_path}" \
27+
--dataset "celeba" \
28+
--log_every_n_steps 2500 \
29+
--save_every_n_steps 25000 \
30+
--save_new_every_n_steps 100000 \
31+
--sample_height 64 \
32+
--sample_width 64 \
33+
--img_size 64 \
34+
--patch_size 4 \
35+
--seed 1 \
8236
--model uvit \
83-
--n_steps 100000 \
84-
--batch_size 128
37+
--normalize_timesteps \
38+
--use_amp \
39+
--parametrization "predict_noise"
8540
```
41+
</details>
8642

87-
#### Early-exit models
88-
Command for training the a DeeDiff model:
89-
```shell
90-
bash scripts/train_deediff.sh
91-
```
92-
or
93-
```
94-
python train.py \
95-
--model deediff_uvit \
43+
### Early-exit training (DeeDiff / AdaDiff)
44+
<details>
45+
<summary>Then, we can train an early-exit model based on the full-model backbone.</summary>
46+
47+
We will assume that `load_backbone` points to the weights obtained in the previous step.
48+
```bash
49+
python main.py \
9650
--n_steps 100000 \
9751
--batch_size 128 \
98-
--classifier_type attention_probe \
99-
--normalize_timesteps
100-
```
101-
102-
Below is a specification of how to run training with other settings.
103-
```shell
104-
python train.py \
105-
--model deediff_uvit \
106-
--n_steps ${number_of_training_steps} \
107-
--batch_size ${batch_size} \
108-
--classifier_type ${classifier_type} \
52+
--log_path "${log_path}" \
53+
--log_every_n_steps 2500 \
54+
--save_every_n_steps 2500 \
55+
--save_new_every_n_steps 10000 \
56+
--seed 1 \
57+
--load_backbone "${load_backbone}" \
58+
--model "deediff_uvit" \
59+
--use_amp \
10960
--normalize_timesteps \
110-
[--load_backbone ${checkpoint_path} \]
111-
[--freeze_backbone \]
112-
[--use_unweighted_loss \]
61+
--parametrization "predict_noise" \
62+
--freeze_backbone \
63+
--dataset "celeba" \
64+
--classifier_type "mlp_probe_per_layer" \
65+
--sample_height 64 \
66+
--sample_width 64 \
67+
--img_size 64 \
68+
--patch_size 4 \
69+
--config_path "configs/deediff_celeba.yaml"
11370
```
11471

115-
- `number_of_training_steps`: Number of iterations over the dataloader.
116-
- `batch_size`: Batch size.
117-
- `classifier_type`: Type of the classifier for determining whether to early-exit. Can be one of:
118-
- `attention_probe`: Attention probe.
119-
- `mlp_probe_per_timestep`: Separate MLP probe at each timestep, shared between layers.
120-
- `mlp_probe_per_layer`: Separate MLP probe for each UViT layer, shared between timesteps.
121-
- `mlp_probe_per_layer_per_timestep`: Separate MLP probe for each UViT layer, at each time step (nothing is shared).
122-
- `--freeze_backbone`: If present, then freeze the UViT backbone (train only the classifiers probes).
123-
- `--use_unweighted_loss`: If present, add the unweighted loss to the remaining losses.
124-
- (optional) `checkpoint_path`: Path to the checkpoint with UViT weights. If not specified, then train DeeDiff from scratch.
72+
</details>
12573

126-
### Evaluation
74+
### DuoDiff training
75+
<details>
76+
<summary>Our proposed model, DuoDiff, involves training a shallow model that will be used alongside the full-model during inference. </summary>
12777

128-
For evaluation, you should have a checkpoint. For convenience, we include one that can be downloaded using [git lfs](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage):
12978
```bash
130-
git lfs pull --include "src/checkpoints/frozenBackbone_attention_3losses.pth"
131-
```
132-
133-
#### CMMD
134-
Command for generating samples and calculating the CMMD score:
135-
```shell
136-
bash scripts/cmmd_evaluation.sh
137-
```
138-
or
139-
```shell
140-
python CMMD_evaluation/main.py \
141-
--checkpoint_entry_name frozenBackbone_attention_3losses \
142-
--exit_threshold 0.05 \
143-
--cmmd_batch_size 32 \
144-
--cmmd_max_count 10
79+
python main.py \
80+
--model "uvit" \
81+
--n_steps 500000 \
82+
--batch_size 128 \
83+
--log_path ${log_path} \
84+
--log_every_n_steps 2500 \
85+
--use_amp \
86+
--save_every_n_steps 25000 \
87+
--save_new_every_n_steps 100000 \
88+
--sample_height 64 \
89+
--sample_width 64 \
90+
--seed 1 \
91+
--normalize_timesteps \
92+
--config_path "configs/uvit_celeba_3.yaml" \
93+
--dataset "celeba" \
94+
--parametrization "predict_noise" \
14595
```
14696

147-
- `cmmd_batch_size`: Batch size for embedding generation.
148-
- `cmmd_max_count`: Maximum number of images to read from each directory.
97+
</details>
14998

150-
#### FID
151-
Command for generating samples and calculating the FID score:
152-
```shell
153-
bash scripts/fid_evaluation.sh
99+
## Running inference
100+
In this section, we will see how to generate images using the models trained on the previous section.
101+
### Early-exit sampling
102+
Here, `checkpoint_path` points to the trained early-exit model (not the full model).
103+
```bash
104+
python eesampler.py \
105+
--seed ${seed} \
106+
--checkpoint_path "${checkpoint_path}" \
107+
--batch_size 128 \
108+
--output_folder "${output_folder}" \
109+
--threshold 0.08 \
110+
--config_path "configs/deediff_celeba.yaml"
154111
```
155-
or
156-
```shell
157-
python FID_evaluation.py \
158-
--checkpoint_entry_name frozenBackbone_attention_3losses \
159-
--exit_threshold 0.05
112+
### DuoDiff inference
113+
Notice that we are using two different models, the full one, and the shallow one.
114+
```bash
115+
python sampler.py \
116+
--seed ${seed} \
117+
--checkpoint_path "${shallow_model_path}" \
118+
--checkpoint_path_late "${full_model_path}" \
119+
--batch_size 128 \
120+
--parametrization "predict_noise" \
121+
--output_folder "${output_folder}" \
122+
--config_path "configs/uvit_celeba_3.yaml" \
123+
--config_path_late "configs/uvit_celeba.yaml" \
124+
--t_switch 300
125+
```
126+
## Computing FID scores
127+
We can easily compute the FID scores running the following script.
128+
```bash
129+
python fid.py \
130+
--dataset "celeba"
131+
--samples_path "${samples_path}"
160132
```
161133

162-
### Benchmarking
163-
For computing the theoretical GFLOPs for the MLP probe, attention probe and output head, you can run the following script
164134

165-
```shell
166-
python get_gflops.py
167-
```
168-
169-
Example script for computing the average layer ratio and theoretical GFLOPs:
170-
```shell
171-
python compute_gflops_and_layer_ratio.py \
172-
--indices_by_timestep_directory benchmarking/output/attention_frozen/indices_by_timestep
173-
```
135+
## Dev instructions
136+
The first time:
137+
1. Run `pre-commit install`
138+
2. If you use VSCode, it might be helpful to add the following to `settings.json`:
139+
```json
140+
"[python]": {
141+
"editor.formatOnSave": true,
142+
"editor.defaultFormatter": "charliermarsh.ruff",
143+
"editor.codeActionsOnSave": {
144+
"source.fixAll": "explicit",
145+
"source.organizeImports": "explicit"
146+
}
147+
},
148+
```
174149

175-
For computing the average layer ratio and theoretical GFLOPs for each method, one can run the following script:
176-
```shell
177-
python compute_gflops_and_layer_ratio.py
178-
--indices_by_timestep_directory ${indices_by_timestep_directory} \
150+
After that, be sure that all the tests are passing before a commit. Otherwise, GitHub Actions will complain ;) You can check by running
151+
```bash
152+
python -m pytest tests
179153
```
180-
The parameter ``indices_by_timestep_directory`` is the relative path to the folder which contains files in ``.pt`` format regarding the layers which early exit took place per timestep. These directories can be found in ``benchmarking/output``. Currently, we uploaded only the ``.pt`` files for the model that uses an attention probe and a frozen backbone during training. The reason why we did not include them for all methods is because the files are pretty large. If one would need the files for the other methods, please contact us.
181-
182-
183-
## Resources
184-
### DeeDiff
185-
186-
Tin's code: https://github.com/stases/EarlyDiffusion
187-
188-
Paper: https://arxiv.org/pdf/2309.17074
189-
190-
### Math Diffusion Models
191-
192-
Lecture Notes in Probabilistic Diffusion Models: https://arxiv.org/html/2312.10393v1
193-
194-
Lil'Log blogpost: https://lilianweng.github.io/posts/2021-07-11-diffusion-models
195-
196-
### Backbones Diffusion Models
197-
198-
U-Net: https://arxiv.org/pdf/1505.04597
199-
200-
U-ViT: https://arxiv.org/pdf/2209.12152
201-
202-
Diffusion Transformer (DiT): https://arxiv.org/pdf/2212.09748

0 commit comments

Comments
 (0)