README.md 1.38 KB
Newer Older
Anton Lozhkov's avatar
Anton Lozhkov committed
1
2
## Training examples

anton-l's avatar
anton-l committed
3
4
5
6
7
8
9
10
### Installing the dependencies

Before running the scipts, make sure to install the library's training dependencies:

```bash
pip install diffusers[training] accelerate datasets
```

anton-l's avatar
anton-l committed
11
### Unconditional Flowers  
Anton Lozhkov's avatar
Anton Lozhkov committed
12
13
14
15

The command to train a DDPM UNet model on the Oxford Flowers dataset:

```bash
16
accelerate launch train_unconditional.py \
Anton Lozhkov's avatar
Anton Lozhkov committed
17
18
  --dataset="huggan/flowers-102-categories" \
  --resolution=64 \
19
20
  --output_dir="ddpm-ema-flowers-64" \
  --train_batch_size=16 \
Anton Lozhkov's avatar
Anton Lozhkov committed
21
22
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
23
24
25
26
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no \
  --push_to_hub
Anton Lozhkov's avatar
Anton Lozhkov committed
27
28
```

anton-l's avatar
anton-l committed
29
A full training run takes 2 hours on 4xV100 GPUs.
Anton Lozhkov's avatar
Anton Lozhkov committed
30
31
32
33

<img src="https://user-images.githubusercontent.com/26864830/173855866-5628989f-856b-4725-a944-d6c09490b2df.png" width="500" />


anton-l's avatar
anton-l committed
34
### Unconditional Pokemon 
Anton Lozhkov's avatar
Anton Lozhkov committed
35
36
37
38

The command to train a DDPM UNet model on the Pokemon dataset:

```bash
39
accelerate launch train_unconditional.py \
Anton Lozhkov's avatar
Anton Lozhkov committed
40
41
  --dataset="huggan/pokemon" \
  --resolution=64 \
42
43
  --output_dir="ddpm-ema-pokemon-64" \
  --train_batch_size=16 \
Anton Lozhkov's avatar
Anton Lozhkov committed
44
45
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
46
47
48
49
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no \
  --push_to_hub
Anton Lozhkov's avatar
Anton Lozhkov committed
50
51
```

anton-l's avatar
anton-l committed
52
A full training run takes 2 hours on 4xV100 GPUs.
Anton Lozhkov's avatar
Anton Lozhkov committed
53
54

<img src="https://user-images.githubusercontent.com/26864830/173856733-4f117f8c-97bd-4f51-8002-56b488c96df9.png" width="500" />