train_unconditional.py 16.9 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
        "--prediction_type",
        type=str,
        default="epsilon",
        choices=["epsilon", "sample"],
201
        help="Whether the model should predict the 'epsilon'/noise error or directly the reconstructed image 'x0'.",
202
203
204
205
206
    )

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

207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
    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
228
def main(args):
229
    logging_dir = os.path.join(args.output_dir, args.logging_dir)
230
    accelerator = Accelerator(
231
        gradient_accumulation_steps=args.gradient_accumulation_steps,
232
233
234
235
        mixed_precision=args.mixed_precision,
        log_with="tensorboard",
        logging_dir=logging_dir,
    )
anton-l's avatar
anton-l committed
236

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

262
    if accepts_prediction_type:
263
264
265
        noise_scheduler = DDPMScheduler(
            num_train_timesteps=args.ddpm_num_steps,
            beta_schedule=args.ddpm_beta_schedule,
266
            prediction_type=args.prediction_type,
267
268
269
270
        )
    else:
        noise_scheduler = DDPMScheduler(num_train_timesteps=args.ddpm_num_steps, beta_schedule=args.ddpm_beta_schedule)

271
272
273
274
275
276
277
    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
278
279
280

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

    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
298
299
300
301
302

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

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

anton-l's avatar
anton-l committed
305
    dataset.set_transform(transforms)
306
307
308
    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
309

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

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

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

323
324
325
326
327
328
    ema_model = EMAModel(
        accelerator.unwrap_model(model),
        inv_gamma=args.ema_inv_gamma,
        power=args.ema_power,
        max_value=args.ema_max_decay,
    )
anton-l's avatar
anton-l committed
329

330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
    # 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
346

347
348
349
350
    if accelerator.is_main_process:
        run = os.path.split(__file__)[-1].split(".")[0]
        accelerator.init_trackers(run)

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

366
            # Add noise to the clean images according to the noise magnitude at each timestep
367
            # (this is the forward diffusion process)
368
369
370
371
            noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)

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

374
                if args.prediction_type == "epsilon":
375
                    loss = F.mse_loss(model_output, noise)  # this could have different weights!
376
                elif args.prediction_type == "sample":
377
378
379
380
381
382
383
384
                    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()
385
386
                else:
                    raise ValueError(f"Unsupported prediction type: {args.prediction_type}")
387

388
                accelerator.backward(loss)
389

390
391
                if accelerator.sync_gradients:
                    accelerator.clip_grad_norm_(model.parameters(), 1.0)
392
393
                optimizer.step()
                lr_scheduler.step()
394
395
                if args.use_ema:
                    ema_model.step(model)
396
                optimizer.zero_grad()
397

398
399
400
401
402
            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                progress_bar.update(1)
                global_step += 1

403
404
405
406
407
            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)
408
        progress_bar.close()
anton-l's avatar
anton-l committed
409

anton-l's avatar
anton-l committed
410
        accelerator.wait_for_everyone()
anton-l's avatar
anton-l committed
411

anton-l's avatar
anton-l committed
412
        # Generate sample images for visual inspection
anton-l's avatar
anton-l committed
413
        if accelerator.is_main_process:
anton-l's avatar
anton-l committed
414
            if epoch % args.save_images_epochs == 0 or epoch == args.num_epochs - 1:
415
416
417
                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
418
                )
anton-l's avatar
anton-l committed
419

420
421
422
423
424
                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
425
                # run pipeline in inference (sample random noise and denoise)
426
427
428
429
430
                images = pipeline(
                    generator=generator,
                    batch_size=args.eval_batch_size,
                    output_type="numpy",
                ).images
anton-l's avatar
anton-l committed
431

anton-l's avatar
anton-l committed
432
433
434
435
436
                # 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
437

438
439
            if epoch % args.save_model_epochs == 0 or epoch == args.num_epochs - 1:
                # save the model
440
                pipeline.save_pretrained(args.output_dir)
441
                if args.push_to_hub:
442
                    repo.push_to_hub(commit_message=f"Epoch {epoch}", blocking=False)
anton-l's avatar
anton-l committed
443
        accelerator.wait_for_everyone()
anton-l's avatar
anton-l committed
444

445
446
    accelerator.end_training()

anton-l's avatar
anton-l committed
447
448

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