train_ddpm.py 6.08 KB
Newer Older
anton-l's avatar
anton-l committed
1
import argparse
anton-l's avatar
anton-l committed
2
import os
anton-l's avatar
anton-l committed
3
4
5
6

import torch
import torch.nn.functional as F

anton-l's avatar
anton-l committed
7
import PIL.Image
anton-l's avatar
anton-l committed
8
9
10
from accelerate import Accelerator
from datasets import load_dataset
from diffusers import DDPM, DDPMScheduler, UNetModel
anton-l's avatar
anton-l committed
11
12
13
14
from torchvision.transforms import (
    Compose,
    InterpolationMode,
    Lambda,
anton-l's avatar
anton-l committed
15
    CenterCrop,
anton-l's avatar
anton-l committed
16
17
18
19
    RandomHorizontalFlip,
    Resize,
    ToTensor,
)
anton-l's avatar
anton-l committed
20
21
22
23
from tqdm.auto import tqdm
from transformers import get_linear_schedule_with_warmup


anton-l's avatar
anton-l committed
24
25
26
27
28
29
30
31
32
33
def main(args):
    accelerator = Accelerator(mixed_precision=args.mixed_precision)

    model = UNetModel(
        attn_resolutions=(16,),
        ch=128,
        ch_mult=(1, 2, 4, 8),
        dropout=0.0,
        num_res_blocks=2,
        resamp_with_conv=True,
anton-l's avatar
anton-l committed
34
        resolution=args.resolution,
anton-l's avatar
anton-l committed
35
36
    )
    noise_scheduler = DDPMScheduler(timesteps=1000)
anton-l's avatar
anton-l committed
37
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
anton-l's avatar
anton-l committed
38
39
40

    augmentations = Compose(
        [
anton-l's avatar
anton-l committed
41
            Resize(args.resolution, interpolation=InterpolationMode.BILINEAR),
anton-l's avatar
anton-l committed
42
            CenterCrop(args.resolution),
anton-l's avatar
anton-l committed
43
44
45
46
47
            RandomHorizontalFlip(),
            ToTensor(),
            Lambda(lambda x: x * 2 - 1),
        ]
    )
anton-l's avatar
anton-l committed
48
    dataset = load_dataset(args.dataset, split="train")
anton-l's avatar
anton-l committed
49
50
51
52
53
54

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

    dataset.set_transform(transforms)
anton-l's avatar
anton-l committed
55
    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
anton-l's avatar
anton-l committed
56
57
58

    lr_scheduler = get_linear_schedule_with_warmup(
        optimizer=optimizer,
anton-l's avatar
anton-l committed
59
60
        num_warmup_steps=args.warmup_steps,
        num_training_steps=(len(train_dataloader) * args.num_epochs) // args.gradient_accumulation_steps,
anton-l's avatar
anton-l committed
61
62
63
64
65
66
    )

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

anton-l's avatar
anton-l committed
67
    for epoch in range(args.num_epochs):
anton-l's avatar
anton-l committed
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
        model.train()
        with tqdm(total=len(train_dataloader), unit="ba") as pbar:
            pbar.set_description(f"Epoch {epoch}")
            for step, batch in enumerate(train_dataloader):
                clean_images = batch["input"]
                noisy_images = torch.empty_like(clean_images)
                noise_samples = torch.empty_like(clean_images)
                bsz = clean_images.shape[0]

                timesteps = torch.randint(0, noise_scheduler.timesteps, (bsz,), device=clean_images.device).long()
                for idx in range(bsz):
                    noise = torch.randn(clean_images.shape[1:]).to(clean_images.device)
                    noise_samples[idx] = noise
                    noisy_images[idx] = noise_scheduler.forward_step(clean_images[idx], noise, timesteps[idx])

anton-l's avatar
anton-l committed
83
                if step % args.gradient_accumulation_steps != 0:
anton-l's avatar
anton-l committed
84
85
                    with accelerator.no_sync(model):
                        output = model(noisy_images, timesteps)
anton-l's avatar
anton-l committed
86
                        # predict the noise residual
anton-l's avatar
anton-l committed
87
88
89
90
                        loss = F.mse_loss(output, noise_samples)
                        accelerator.backward(loss)
                else:
                    output = model(noisy_images, timesteps)
anton-l's avatar
anton-l committed
91
                    # predict the noise residual
anton-l's avatar
anton-l committed
92
93
94
95
96
97
98
99
100
101
102
                    loss = F.mse_loss(output, noise_samples)
                    accelerator.backward(loss)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                    optimizer.step()
                    lr_scheduler.step()
                    optimizer.zero_grad()
                pbar.update(1)
                pbar.set_postfix(loss=loss.detach().item(), lr=optimizer.param_groups[0]["lr"])

                optimizer.step()

anton-l's avatar
anton-l committed
103
        # Generate a sample image for visual inspection
anton-l's avatar
anton-l committed
104
105
106
107
        torch.distributed.barrier()
        if args.local_rank in [-1, 0]:
            model.eval()
            with torch.no_grad():
anton-l's avatar
anton-l committed
108
109
110
111
112
113
114
                if isinstance(model, torch.nn.parallel.DistributedDataParallel):
                    pipeline = DDPM(unet=model.module, noise_scheduler=noise_scheduler)
                else:
                    pipeline = DDPM(unet=model, noise_scheduler=noise_scheduler)
                pipeline.save_pretrained(args.output_path)

                generator = torch.manual_seed(0)
anton-l's avatar
anton-l committed
115
116
117
118
119
120
121
122
123
124
                # run pipeline in inference (sample random noise and denoise)
                image = pipeline(generator=generator)

                # process image to PIL
                image_processed = image.cpu().permute(0, 2, 3, 1)
                image_processed = (image_processed + 1.0) * 127.5
                image_processed = image_processed.type(torch.uint8).numpy()
                image_pil = PIL.Image.fromarray(image_processed[0])

                # save image
anton-l's avatar
anton-l committed
125
126
127
                test_dir = os.path.join(args.output_path, "test_samples")
                os.makedirs(test_dir, exist_ok=True)
                image_pil.save(f"{test_dir}/{epoch}.png")
anton-l's avatar
anton-l committed
128
129
130
131
        torch.distributed.barrier()


if __name__ == "__main__":
anton-l's avatar
anton-l committed
132
    parser = argparse.ArgumentParser(description="Simple example of a training script.")
anton-l's avatar
anton-l committed
133
    parser.add_argument("--local_rank", type=int)
anton-l's avatar
anton-l committed
134
135
136
137
138
    parser.add_argument("--dataset", type=str, default="huggan/flowers-102-categories")
    parser.add_argument("--resolution", type=int, default=64)
    parser.add_argument("--output_path", type=str, default="ddpm-model")
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--num_epochs", type=int, default=100)
anton-l's avatar
anton-l committed
139
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
anton-l's avatar
anton-l committed
140
141
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--warmup_steps", type=int, default=500)
anton-l's avatar
anton-l committed
142
143
144
145
146
    parser.add_argument(
        "--mixed_precision",
        type=str,
        default="no",
        choices=["no", "fp16", "bf16"],
147
148
149
150
151
        help=(
            "Whether to use mixed precision. Choose"
            "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
            "and an Nvidia Ampere GPU."
        ),
anton-l's avatar
anton-l committed
152
153
154
155
156
157
158
159
    )

    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

    main(args)