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

3
4
Creating a training image set is [described in a different document](https://huggingface.co/docs/datasets/image_process#image-datasets).

anton-l's avatar
anton-l committed
5
6
### Installing the dependencies

7
Before running the scripts, make sure to install the library's training dependencies:
anton-l's avatar
anton-l committed
8

9
10
11
**Important**

To make sure you can successfully run the latest versions of the example scripts, we highly recommend **installing from source** and keeping the install up to date as we update the example scripts frequently and install some example-specific requirements. To do this, execute the following steps in a new virtual environment:
anton-l's avatar
anton-l committed
12
```bash
13
14
15
git clone https://github.com/huggingface/diffusers
cd diffusers
pip install .
anton-l's avatar
anton-l committed
16
17
```

18
19
20
21
22
23
Then cd in the example folder  and run
```bash
pip install -r requirements.txt
```


24
25
26
27
28
29
And initialize an [🤗Accelerate](https://github.com/huggingface/accelerate/) environment with:

```bash
accelerate config
```

anton-l's avatar
anton-l committed
30
### Unconditional Flowers  
Anton Lozhkov's avatar
Anton Lozhkov committed
31
32
33
34

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

```bash
35
accelerate launch train_unconditional.py \
36
  --dataset_name="huggan/flowers-102-categories" \
Anton Lozhkov's avatar
Anton Lozhkov committed
37
  --resolution=64 \
38
39
  --output_dir="ddpm-ema-flowers-64" \
  --train_batch_size=16 \
Anton Lozhkov's avatar
Anton Lozhkov committed
40
41
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
42
  --use_ema \
43
44
45
46
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no \
  --push_to_hub
Anton Lozhkov's avatar
Anton Lozhkov committed
47
```
Anton Lozhkov's avatar
Anton Lozhkov committed
48
An example trained model: https://huggingface.co/anton-l/ddpm-ema-flowers-64
Anton Lozhkov's avatar
Anton Lozhkov committed
49

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

Anton Lozhkov's avatar
Anton Lozhkov committed
52
<img src="https://user-images.githubusercontent.com/26864830/180248660-a0b143d0-b89a-42c5-8656-2ebf6ece7e52.png" width="700" />
Anton Lozhkov's avatar
Anton Lozhkov committed
53
54


anton-l's avatar
anton-l committed
55
### Unconditional Pokemon 
Anton Lozhkov's avatar
Anton Lozhkov committed
56
57
58
59

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

```bash
60
accelerate launch train_unconditional.py \
61
  --dataset_name="huggan/pokemon" \
Anton Lozhkov's avatar
Anton Lozhkov committed
62
  --resolution=64 \
63
64
  --output_dir="ddpm-ema-pokemon-64" \
  --train_batch_size=16 \
Anton Lozhkov's avatar
Anton Lozhkov committed
65
66
  --num_epochs=100 \
  --gradient_accumulation_steps=1 \
67
  --use_ema \
68
69
70
71
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=no \
  --push_to_hub
Anton Lozhkov's avatar
Anton Lozhkov committed
72
```
Anton Lozhkov's avatar
Anton Lozhkov committed
73
An example trained model: https://huggingface.co/anton-l/ddpm-ema-pokemon-64
Anton Lozhkov's avatar
Anton Lozhkov committed
74

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

Anton Lozhkov's avatar
Anton Lozhkov committed
77
<img src="https://user-images.githubusercontent.com/26864830/180248200-928953b4-db38-48db-b0c6-8b740fe6786f.png" width="700" />
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117


### Using your own data

To use your own dataset, there are 2 ways:
- you can either provide your own folder as `--train_data_dir`
- or you can upload your dataset to the hub (possibly as a private repo, if you prefer so), and simply pass the `--dataset_name` argument.

Below, we explain both in more detail.

#### Provide the dataset as a folder

If you provide your own folders with images, the script expects the following directory structure:

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

In other words, the script will take care of gathering all images inside the folder. You can then run the script like this:

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

Internally, the script will use the [`ImageFolder`](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#imagefolder) feature which will automatically turn the folders into 🤗 Dataset objects.

#### Upload your data to the hub, as a (possibly private) repo

It's very easy (and convenient) to upload your image dataset to the hub using the [`ImageFolder`](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#imagefolder) feature available in 🤗 Datasets. Simply do the following:

```python
from datasets import load_dataset

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

118
# example 2: local files (supported formats are tar, gzip, zip, xz, rar, zstd)
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
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"]})
```

`ImageFolder` will create an `image` column containing the PIL-encoded images.

Next, push it to the hub!

```python
# assuming you have ran the huggingface-cli login command in a terminal
dataset.push_to_hub("name_of_your_dataset")

# if you want to push to a private repo, simply pass private=True:
dataset.push_to_hub("name_of_your_dataset", private=True)
```

and that's it! You can now train your model by simply setting the `--dataset_name` argument to the name of your dataset on the hub.

More on this can also be found in [this blog post](https://huggingface.co/blog/image-search-datasets).
Prathik Rao's avatar
Prathik Rao committed
143
144
145
146
147
148
149
150
151
152
153
154

#### Use ONNXRuntime to accelerate training

In order to leverage onnxruntime to accelerate training, please use train_unconditional_ort.py

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

```bash
accelerate launch train_unconditional_ort.py \
  --dataset_name="huggan/flowers-102-categories" \
  --resolution=64 \
  --output_dir="ddpm-ema-flowers-64" \
155
  --use_ema \
Prathik Rao's avatar
Prathik Rao committed
156
157
158
159
160
161
162
163
164
  --train_batch_size=16 \
  --num_epochs=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_warmup_steps=500 \
  --mixed_precision=fp16
  ```

Please contact Prathik Rao (prathikr), Sunghoon Choi (hanbitmyths), Ashwini Khade (askhade), or Peng Wang (pengwa) on github with any questions.