unconditional_training.mdx 7.4 KB
Newer Older
Patrick von Platen's avatar
Patrick von Platen committed
1
<!--Copyright 2023 The HuggingFace Team. All rights reserved.
Nathan Lambert's avatar
Nathan Lambert committed
2
3
4
5
6
7
8
9
10
11
12

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
specific language governing permissions and limitations under the License.
-->

13
# Unconditional image generation
Patrick von Platen's avatar
Patrick von Platen committed
14

15
Unconditional image generation is not conditioned on any text or images, unlike text- or image-to-image models. It only generates images that resemble its training data distribution.
Patrick von Platen's avatar
Patrick von Platen committed
16

17
18
19
20
21
22
<iframe
	src="https://stevhliu-ddpm-butterflies-128.hf.space"
	frameborder="0"
	width="850"
	height="550"
></iframe>
Patrick von Platen's avatar
Patrick von Platen committed
23

24
25
26
27

This guide will show you how to train an unconditional image generation model on existing datasets as well as your own custom dataset. All the training scripts for unconditional image generation can be found [here](https://github.com/huggingface/diffusers/tree/main/examples/unconditional_image_generation) if you're interested in learning more about the training details.

Before running the script, make sure you install the library's training dependencies:
Patrick von Platen's avatar
Patrick von Platen committed
28

29
30
```bash
pip install diffusers[training] accelerate datasets
Patrick von Platen's avatar
Patrick von Platen committed
31
```
32

33
Next, initialize an 🤗 [Accelerate](https://github.com/huggingface/accelerate/) environment with:
34
35
36
37
38

```bash
accelerate config
```

39
40
41
42
43
To setup a default 🤗 Accelerate environment without choosing any configurations:

```bash
accelerate config default
```
44

45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
Or if your environment doesn't support an interactive shell like a notebook, you can use:

```bash
from accelerate.utils import write_basic_config

write_basic_config()
```

## Upload model to Hub

You can upload your model on the Hub by adding the following argument to the training script:

```bash
--push_to_hub
```

## Save and load checkpoints

It is a good idea to regularly save checkpoints in case anything happens during training. To save a checkpoint, pass the following argument to the training script:

```bash
--checkpointing_steps=500
```

The full training state is saved in a subfolder in the `output_dir` every 500 steps, which allows you to load a checkpoint and resume training if you pass the `--resume_from_checkpoint` argument to the training script:

```bash
--resume_from_checkpoint="checkpoint-1500"
```

## Finetuning

You're ready to launch the [training script](https://github.com/huggingface/diffusers/blob/main/examples/unconditional_image_generation/train_unconditional.py) now! Specify the dataset name to finetune on with the `--dataset_name` argument and then save it to the path in `--output_dir`.

<Tip>

💡 A full training run takes 2 hours on 4xV100 GPUs.

</Tip>

For example, to finetune on the [Oxford Flowers](https://huggingface.co/datasets/huggan/flowers-102-categories) dataset:
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

```bash
accelerate launch train_unconditional.py \
  --dataset_name="huggan/flowers-102-categories" \
  --resolution=64 \
  --output_dir="ddpm-ema-flowers-64" \
  --train_batch_size=16 \
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no \
  --push_to_hub
```

101
102
103
<div class="flex justify-center">
    <img src="https://user-images.githubusercontent.com/26864830/180248660-a0b143d0-b89a-42c5-8656-2ebf6ece7e52.png"/>
</div>
104

105
Or if you want to train your model on the [Pokemon](https://huggingface.co/datasets/huggan/pokemon) dataset:
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

```bash
accelerate launch train_unconditional.py \
  --dataset_name="huggan/pokemon" \
  --resolution=64 \
  --output_dir="ddpm-ema-pokemon-64" \
  --train_batch_size=16 \
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no \
  --push_to_hub
```

121
122
123
<div class="flex justify-center">
    <img src="https://user-images.githubusercontent.com/26864830/180248200-928953b4-db38-48db-b0c6-8b740fe6786f.png"/>
</div>
124

125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
### Training with multiple GPUs

`accelerate` allows for seamless multi-GPU training. Follow the instructions [here](https://huggingface.co/docs/accelerate/basic_tutorials/launch)
for running distributed training with `accelerate`. Here is an example command:

```bash
accelerate launch --mixed_precision="fp16" --multi_gpu train_unconditional.py \
  --dataset_name="huggan/pokemon" \
  --resolution=64 --center_crop --random_flip \
  --output_dir="ddpm-ema-pokemon-64" \
  --train_batch_size=16 \
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
  --use_ema \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision="fp16" \
  --logger="wandb"
```

145
## Finetuning with your own data
146

147
There are two ways to finetune a model on your own dataset:
148

149
150
- provide your own folder of images to the `--train_data_dir` argument
- upload your dataset to the Hub and pass the dataset repository id to the `--dataset_name` argument.
151

152
<Tip>
153

154
155
156
💡 Learn more about how to create an image dataset for training in the [Create an image dataset](https://huggingface.co/docs/datasets/image_dataset) guide. 

</Tip>
157
158
159
160
161

Below, we explain both in more detail.

### Provide the dataset as a folder

162
If you provide your own dataset as a folder, the script expects the following directory structure:
163
164
165
166
167
168
169

```bash
data_dir/xxx.png
data_dir/xxy.png
data_dir/[...]/xxz.png
```

170
Pass the path to the folder containing the images to the `--train_data_dir` argument and launch the training:
171
172
173
174
175
176
177

```bash
accelerate launch train_unconditional.py \
    --train_data_dir <path-to-train-directory> \
    <other-arguments>
```

178
179
180
Internally, the script uses the [`ImageFolder`](https://huggingface.co/docs/datasets/image_load#imagefolder) to automatically build a dataset from the folder.

### Upload your data to the Hub
181

182
<Tip>
183

184
185
186
187
188
💡 For more details and context about creating and uploading a dataset to the Hub, take a look at the [Image search with 🤗 Datasets](https://huggingface.co/blog/image-search-datasets) post.

</Tip>

To upload your dataset to the Hub, you can start by creating one with the [`ImageFolder`](https://huggingface.co/docs/datasets/image_load#imagefolder) feature, which creates an `image` column containing the PIL-encoded images, from 🤗 Datasets:
189
190
191
192
193
194
195

```python
from datasets import load_dataset

# example 1: local folder
dataset = load_dataset("imagefolder", data_dir="path_to_your_folder")

196
# example 2: local files (supported formats are tar, gzip, zip, xz, rar, zstd)
197
198
199
200
201
202
203
204
205
206
207
208
dataset = load_dataset("imagefolder", data_files="path_to_zip_file")

# example 3: remote files (supported formats are tar, gzip, zip, xz, rar, zstd)
dataset = load_dataset(
    "imagefolder",
    data_files="https://download.microsoft.com/download/3/E/1/3E1C3F21-ECDB-4869-8368-6DEBA77B919F/kagglecatsanddogs_3367a.zip",
)

# example 4: providing several splits
dataset = load_dataset(
    "imagefolder", data_files={"train": ["path/to/file1", "path/to/file2"], "test": ["path/to/file3", "path/to/file4"]}
)
Patrick von Platen's avatar
Patrick von Platen committed
209
210
```

211
Then you can use the [`~datasets.Dataset.push_to_hub`] method to upload it to the Hub:
Patrick von Platen's avatar
Patrick von Platen committed
212

213
214
215
```python
# assuming you have ran the huggingface-cli login command in a terminal
dataset.push_to_hub("name_of_your_dataset")
Patrick von Platen's avatar
Patrick von Platen committed
216

217
218
219
# if you want to push to a private repo, simply pass private=True:
dataset.push_to_hub("name_of_your_dataset", private=True)
```
Patrick von Platen's avatar
Patrick von Platen committed
220

221
Now train your model by simply setting the `--dataset_name` argument to the name of your dataset on the Hub.