train_unconditional.py 16.7 KB
Newer Older
anton-l's avatar
anton-l committed
1
import argparse
2
import inspect
3
import math
anton-l's avatar
anton-l committed
4
import os
5
6
from pathlib import Path
from typing import Optional
anton-l's avatar
anton-l committed
7
8
9
10

import torch
import torch.nn.functional as F

11
from accelerate import Accelerator
12
from accelerate.logging import get_logger
anton-l's avatar
anton-l committed
13
from datasets import load_dataset
14
from diffusers import DDPMPipeline, DDPMScheduler, UNet2DModel, __version__
15
from diffusers.optimization import get_scheduler
anton-l's avatar
anton-l committed
16
from diffusers.training_utils import EMAModel
17
from diffusers.utils import deprecate
18
from huggingface_hub import HfFolder, Repository, whoami
19
from packaging import version
anton-l's avatar
anton-l committed
20
from torchvision.transforms import (
Patrick von Platen's avatar
Patrick von Platen committed
21
    CenterCrop,
anton-l's avatar
anton-l committed
22
23
    Compose,
    InterpolationMode,
anton-l's avatar
anton-l committed
24
    Normalize,
anton-l's avatar
anton-l committed
25
26
27
28
    RandomHorizontalFlip,
    Resize,
    ToTensor,
)
anton-l's avatar
anton-l committed
29
from tqdm.auto import tqdm
anton-l's avatar
anton-l committed
30
31


32
logger = get_logger(__name__)
33
diffusers_version = version.parse(version.parse(__version__).base_version)
anton-l's avatar
anton-l committed
34
35


36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
def _extract_into_tensor(arr, timesteps, broadcast_shape):
    """
    Extract values from a 1-D numpy array for a batch of indices.

    :param arr: the 1-D numpy array.
    :param timesteps: a tensor of indices into the array to extract.
    :param broadcast_shape: a larger shape of K dimensions with the batch
                            dimension equal to the length of timesteps.
    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
    """
    if not isinstance(arr, torch.Tensor):
        arr = torch.from_numpy(arr)
    res = arr[timesteps].float().to(timesteps.device)
    while len(res.shape) < len(broadcast_shape):
        res = res[..., None]
    return res.expand(broadcast_shape)


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
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
def parse_args():
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
    parser.add_argument(
        "--dataset_name",
        type=str,
        default=None,
        help=(
            "The name of the Dataset (from the HuggingFace hub) to train on (could be your own, possibly private,"
            " dataset). It can also be a path pointing to a local copy of a dataset in your filesystem,"
            " or to a folder containing files that HF Datasets can understand."
        ),
    )
    parser.add_argument(
        "--dataset_config_name",
        type=str,
        default=None,
        help="The config of the Dataset, leave as None if there's only one config.",
    )
    parser.add_argument(
        "--train_data_dir",
        type=str,
        default=None,
        help=(
            "A folder containing the training data. Folder contents must follow the structure described in"
            " https://huggingface.co/docs/datasets/image_dataset#imagefolder. In particular, a `metadata.jsonl` file"
            " must exist to provide the captions for the images. Ignored if `dataset_name` is specified."
        ),
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="ddpm-model-64",
        help="The output directory where the model predictions and checkpoints will be written.",
    )
    parser.add_argument("--overwrite_output_dir", action="store_true")
    parser.add_argument(
        "--cache_dir",
        type=str,
        default=None,
        help="The directory where the downloaded models and datasets will be stored.",
    )
    parser.add_argument(
        "--resolution",
        type=int,
        default=64,
        help=(
            "The resolution for input images, all the images in the train/validation dataset will be resized to this"
            " resolution"
        ),
    )
    parser.add_argument(
        "--train_batch_size", type=int, default=16, help="Batch size (per device) for the training dataloader."
    )
    parser.add_argument(
108
109
110
111
112
113
114
115
116
117
        "--eval_batch_size", type=int, default=16, help="The number of images to generate for evaluation."
    )
    parser.add_argument(
        "--dataloader_num_workers",
        type=int,
        default=0,
        help=(
            "The number of subprocesses to use for data loading. 0 means that the data will be loaded in the main"
            " process."
        ),
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
    )
    parser.add_argument("--num_epochs", type=int, default=100)
    parser.add_argument("--save_images_epochs", type=int, default=10, help="How often to save images during training.")
    parser.add_argument(
        "--save_model_epochs", type=int, default=10, help="How often to save the model during training."
    )
    parser.add_argument(
        "--gradient_accumulation_steps",
        type=int,
        default=1,
        help="Number of updates steps to accumulate before performing a backward/update pass.",
    )
    parser.add_argument(
        "--learning_rate",
        type=float,
        default=1e-4,
        help="Initial learning rate (after the potential warmup period) to use.",
    )
    parser.add_argument(
        "--lr_scheduler",
        type=str,
        default="cosine",
        help=(
            'The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
            ' "constant", "constant_with_warmup"]'
        ),
    )
    parser.add_argument(
        "--lr_warmup_steps", type=int, default=500, help="Number of steps for the warmup in the lr scheduler."
    )
    parser.add_argument("--adam_beta1", type=float, default=0.95, help="The beta1 parameter for the Adam optimizer.")
    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
    parser.add_argument(
        "--adam_weight_decay", type=float, default=1e-6, help="Weight decay magnitude for the Adam optimizer."
    )
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer.")
    parser.add_argument(
        "--use_ema",
        action="store_true",
        default=True,
        help="Whether to use Exponential Moving Average for the final model weights.",
    )
    parser.add_argument("--ema_inv_gamma", type=float, default=1.0, help="The inverse gamma value for the EMA decay.")
    parser.add_argument("--ema_power", type=float, default=3 / 4, help="The power value for the EMA decay.")
    parser.add_argument("--ema_max_decay", type=float, default=0.9999, help="The maximum decay magnitude for EMA.")
    parser.add_argument("--push_to_hub", action="store_true", help="Whether or not to push the model to the Hub.")
    parser.add_argument("--hub_token", type=str, default=None, help="The token to use to push to the Model Hub.")
    parser.add_argument(
        "--hub_model_id",
        type=str,
        default=None,
        help="The name of the repository to keep in sync with the local `output_dir`.",
    )
    parser.add_argument(
        "--hub_private_repo", action="store_true", help="Whether or not to create a private repository."
    )
    parser.add_argument(
        "--logging_dir",
        type=str,
        default="logs",
        help=(
            "[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
            " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."
        ),
    )
    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default="no",
        choices=["no", "fp16", "bf16"],
        help=(
            "Whether to use mixed precision. Choose"
            "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
            "and an Nvidia Ampere GPU."
        ),
    )

196
    parser.add_argument(
197
198
199
200
        "--predict_epsilon",
        action="store_true",
        default=True,
        help="Whether the model should predict the 'epsilon'/noise error or directly the reconstructed image 'x0'.",
201
202
203
204
205
    )

    parser.add_argument("--ddpm_num_steps", type=int, default=1000)
    parser.add_argument("--ddpm_beta_schedule", type=str, default="linear")

206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
    args = parser.parse_args()
    env_local_rank = int(os.environ.get("LOCAL_RANK", -1))
    if env_local_rank != -1 and env_local_rank != args.local_rank:
        args.local_rank = env_local_rank

    if args.dataset_name is None and args.train_data_dir is None:
        raise ValueError("You must specify either a dataset name from the hub or a train data directory.")

    return args


def get_full_repo_name(model_id: str, organization: Optional[str] = None, token: Optional[str] = None):
    if token is None:
        token = HfFolder.get_token()
    if organization is None:
        username = whoami(token)["name"]
        return f"{username}/{model_id}"
    else:
        return f"{organization}/{model_id}"


anton-l's avatar
anton-l committed
227
def main(args):
228
    logging_dir = os.path.join(args.output_dir, args.logging_dir)
229
    accelerator = Accelerator(
230
        gradient_accumulation_steps=args.gradient_accumulation_steps,
231
232
233
234
        mixed_precision=args.mixed_precision,
        log_with="tensorboard",
        logging_dir=logging_dir,
    )
anton-l's avatar
anton-l committed
235

anton-l's avatar
anton-l committed
236
237
    model = UNet2DModel(
        sample_size=args.resolution,
238
239
        in_channels=3,
        out_channels=3,
anton-l's avatar
anton-l committed
240
241
242
243
244
245
246
247
248
        layers_per_block=2,
        block_out_channels=(128, 128, 256, 256, 512, 512),
        down_block_types=(
            "DownBlock2D",
            "DownBlock2D",
            "DownBlock2D",
            "DownBlock2D",
            "AttnDownBlock2D",
            "DownBlock2D",
249
        ),
anton-l's avatar
anton-l committed
250
251
252
253
254
255
256
        up_block_types=(
            "UpBlock2D",
            "AttnUpBlock2D",
            "UpBlock2D",
            "UpBlock2D",
            "UpBlock2D",
            "UpBlock2D",
257
        ),
anton-l's avatar
anton-l committed
258
    )
259
260
261
262
263
264
265
266
267
268
269
    accepts_predict_epsilon = "predict_epsilon" in set(inspect.signature(DDPMScheduler.__init__).parameters.keys())

    if accepts_predict_epsilon:
        noise_scheduler = DDPMScheduler(
            num_train_timesteps=args.ddpm_num_steps,
            beta_schedule=args.ddpm_beta_schedule,
            predict_epsilon=args.predict_epsilon,
        )
    else:
        noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)

270
271
272
273
274
275
276
    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )
anton-l's avatar
anton-l committed
277
278
279

    augmentations = Compose(
        [
anton-l's avatar
anton-l committed
280
            Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),
anton-l's avatar
anton-l committed
281
            CenterCrop(args.resolution),
anton-l's avatar
anton-l committed
282
283
            RandomHorizontalFlip(),
            ToTensor(),
anton-l's avatar
anton-l committed
284
            Normalize([0.5], [0.5]),
anton-l's avatar
anton-l committed
285
286
        ]
    )
287
288
289
290
291
292
293
294
295
296

    if args.dataset_name is not None:
        dataset = load_dataset(
            args.dataset_name,
            args.dataset_config_name,
            cache_dir=args.cache_dir,
            split="train",
        )
    else:
        dataset = load_dataset("imagefolder", data_dir=args.train_data_dir, cache_dir=args.cache_dir, split="train")
anton-l's avatar
anton-l committed
297
298
299
300
301

    def transforms(examples):
        images = [augmentations(image.convert("RGB")) for image in examples["image"]]
        return {"input": images}

302
303
    logger.info(f"Dataset size: {len(dataset)}")

anton-l's avatar
anton-l committed
304
    dataset.set_transform(transforms)
305
306
307
    train_dataloader = torch.utils.data.DataLoader(
        dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=args.dataloader_num_workers
    )
anton-l's avatar
anton-l committed
308

anton-l's avatar
anton-l committed
309
    lr_scheduler = get_scheduler(
310
        args.lr_scheduler,
anton-l's avatar
anton-l committed
311
        optimizer=optimizer,
312
        num_warmup_steps=args.lr_warmup_steps,
anton-l's avatar
anton-l committed
313
        num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,
anton-l's avatar
anton-l committed
314
315
316
317
318
319
    )

    model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        model, optimizer, train_dataloader, lr_scheduler
    )

320
321
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)

322
    ema_model = EMAModel(model, inv_gamma=args.ema_inv_gamma, power=args.ema_power, max_value=args.ema_max_decay)
anton-l's avatar
anton-l committed
323

324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
    # Handle the repository creation
    if accelerator.is_main_process:
        if args.push_to_hub:
            if args.hub_model_id is None:
                repo_name = get_full_repo_name(Path(args.output_dir).name, token=args.hub_token)
            else:
                repo_name = args.hub_model_id
            repo = Repository(args.output_dir, clone_from=repo_name)

            with open(os.path.join(args.output_dir, ".gitignore"), "w+") as gitignore:
                if "step_*" not in gitignore:
                    gitignore.write("step_*\n")
                if "epoch_*" not in gitignore:
                    gitignore.write("epoch_*\n")
        elif args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)
anton-l's avatar
anton-l committed
340

341
342
343
344
    if accelerator.is_main_process:
        run = os.path.split(__file__)[-1].split(".")[0]
        accelerator.init_trackers(run)

anton-l's avatar
anton-l committed
345
    global_step = 0
anton-l's avatar
anton-l committed
346
    for epoch in range(args.num_epochs):
anton-l's avatar
anton-l committed
347
        model.train()
348
        progress_bar = tqdm(total=num_update_steps_per_epoch, disable=not accelerator.is_local_main_process)
349
350
351
        progress_bar.set_description(f"Epoch {epoch}")
        for step, batch in enumerate(train_dataloader):
            clean_images = batch["input"]
352
353
            # Sample noise that we'll add to the images
            noise = torch.randn(clean_images.shape).to(clean_images.device)
354
            bsz = clean_images.shape[0]
355
356
            # Sample a random timestep for each image
            timesteps = torch.randint(
357
                0, noise_scheduler.config.num_train_timesteps, (bsz,), device=clean_images.device
358
            ).long()
359

360
            # Add noise to the clean images according to the noise magnitude at each timestep
361
            # (this is the forward diffusion process)
362
363
364
365
            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

            with accelerator.accumulate(model):
                # Predict the noise residual
366
367
                model_output = model(noisy_images, timesteps).sample

368
                if args.predict_epsilon:
369
                    loss = F.mse_loss(model_output, noise)  # this could have different weights!
370
                else:
371
372
373
374
375
376
377
378
379
                    alpha_t = _extract_into_tensor(
                        noise_scheduler.alphas_cumprod, timesteps, (clean_images.shape[0], 1, 1, 1)
                    )
                    snr_weights = alpha_t / (1 - alpha_t)
                    loss = snr_weights * F.mse_loss(
                        model_output, clean_images, reduction="none"
                    )  # use SNR weighting from distillation paper
                    loss = loss.mean()

380
                accelerator.backward(loss)
381

382
383
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
384
385
                optimizer.step()
                lr_scheduler.step()
386
387
                if args.use_ema:
                    ema_model.step(model)
388
                optimizer.zero_grad()
389

390
391
392
393
394
            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                progress_bar.update(1)
                global_step += 1

395
396
397
398
399
            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
            if args.use_ema:
                logs["ema_decay"] = ema_model.decay
            progress_bar.set_postfix(**logs)
            accelerator.log(logs, step=global_step)
400
        progress_bar.close()
anton-l's avatar
anton-l committed
401

anton-l's avatar
anton-l committed
402
        accelerator.wait_for_everyone()
anton-l's avatar
anton-l committed
403

anton-l's avatar
anton-l committed
404
        # Generate sample images for visual inspection
anton-l's avatar
anton-l committed
405
        if accelerator.is_main_process:
anton-l's avatar
anton-l committed
406
            if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
407
408
409
                pipeline = DDPMPipeline(
                    unet=accelerator.unwrap_model(ema_model.averaged_model if args.use_ema else model),
                    scheduler=noise_scheduler,
anton-l's avatar
anton-l committed
410
                )
anton-l's avatar
anton-l committed
411

412
413
414
415
416
                deprecate("todo: remove this check", "0.10.0", "when the most used version is >= 0.8.0")
                if diffusers_version < version.parse("0.8.0"):
                    generator = torch.manual_seed(0)
                else:
                    generator = torch.Generator(device=pipeline.device).manual_seed(0)
anton-l's avatar
anton-l committed
417
                # run pipeline in inference (sample random noise and denoise)
418
419
420
421
422
                images = pipeline(
                    generator=generator,
                    batch_size=args.eval_batch_size,
                    output_type="numpy",
                ).images
anton-l's avatar
anton-l committed
423

anton-l's avatar
anton-l committed
424
425
426
427
428
                # denormalize the images and save to tensorboard
                images_processed = (images * 255).round().astype("uint8")
                accelerator.trackers[0].writer.add_images(
                    "test_samples", images_processed.transpose(0, 3, 1, 2), epoch
                )
anton-l's avatar
anton-l committed
429

430
431
            if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
                # save the model
432
                pipeline.save_pretrained(args.output_dir)
433
                if args.push_to_hub:
434
                    repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=False)
anton-l's avatar
anton-l committed
435
        accelerator.wait_for_everyone()
anton-l's avatar
anton-l committed
436

437
438
    accelerator.end_training()

anton-l's avatar
anton-l committed
439
440

if __name__ == "__main__":
441
    args = parse_args()
anton-l's avatar
anton-l committed
442
    main(args)