diffusers/v0.43.0/ko/tutorials/basic_training.md

확산 모델 학습

무조건적 이미지 생성은 학습에 사용된 데이터셋의 이미지와 유사한 이미지를 생성하는 diffusion models의 인기 있는 응용 분야입니다. 일반적으로 특정 데이터셋에서 사전 학습된 모델을 fine-tuning할 때 가장 좋은 결과를 얻을 수 있습니다. 이러한 checkpoint 중 많은 것을 Hub에서 찾을 수 있지만, 마음에 드는 것을 찾지 못한다면 직접 훈련할 수도 있습니다!

이 튜토리얼에서는 Smithsonian Butterflies 데이터셋의 하위 집합을 사용하여 UNet2DModel을 처음부터 학습하여 나만의 🦋 나비 🦋를 생성하는 방법을 배웁니다.

💡 이 학습 튜토리얼은 Training with 🧨 Diffusers 노트북을 기반으로 합니다. Diffusion 모델이 어떻게 작동하는지 등 추가적인 세부 사항과 배경을 확인하려면 노트북을 참고하세요!

시작하기 전에 이미지 데이터셋을 로드하고 전처리하는 데 필요한 🤗 Datasets와, GPU 개수에 관계없이 학습을 단순화하는 🤗 Accelerate가 설치되어 있는지 확인하세요. 다음 명령을 실행하면 학습 메트릭을 시각화하는 데 사용할 TensorBoard도 함께 설치됩니다(학습을 추적하는 데 Weights & Biases를 사용할 수도 있습니다).

# uncomment to install the necessary libraries in Colab
#!pip install diffusers[training]

모델을 커뮤니티와 공유하는 것을 권장하며, 이를 위해 Hugging Face 계정에 로그인해야 합니다(계정이 없다면 여기에서 생성하세요!). 노트북에서 로그인하고 프롬프트가 표시되면 토큰을 입력할 수 있습니다. 토큰에 write 역할이 있는지 확인하세요.

>>> from huggingface_hub import notebook_login

>>> notebook_login()

또는 터미널에서 로그인하세요:

hf auth login

모델 체크포인트가 상당히 크기 때문에 이러한 대용량 파일의 버전 관리를 위해 Git-LFS를 설치하세요:

!sudo apt -qq install git-lfs
!git config --global credential.helper store

학습 구성

편의상 학습 하이퍼파라미터를 포함하는 TrainingConfig 클래스를 생성합니다(값은 자유롭게 조정할 수 있습니다):

>>> from dataclasses import dataclass

>>> @dataclass
... class TrainingConfig:
...     image_size = 128  # the generated image resolution
...     train_batch_size = 16
...     eval_batch_size = 16  # how many images to sample during evaluation
...     num_epochs = 50
...     gradient_accumulation_steps = 1
...     learning_rate = 1e-4
...     lr_warmup_steps = 500
...     save_image_epochs = 10
...     save_model_epochs = 30
...     mixed_precision = "fp16"  # `no` for float32, `fp16` for automatic mixed precision
...     output_dir = "ddpm-butterflies-128"  # the model name locally and on the HF Hub

...     push_to_hub = True  # whether to upload the saved model to the HF Hub
...     hub_model_id = "<your-username>/<my-awesome-model>"  # the name of the repository to create on the HF Hub
...     hub_private_repo = None
...     overwrite_output_dir = True  # overwrite the old model when re-running the notebook
...     seed = 0

>>> config = TrainingConfig()

데이터셋 로드

🤗 Datasets 라이브러리를 사용하면 Smithsonian Butterflies 데이터셋을 쉽게 로드할 수 있습니다:

>>> from datasets import load_dataset

>>> config.dataset_name = "huggan/smithsonian_butterflies_subset"
>>> dataset = load_dataset(config.dataset_name, split="train")

💡 HugGan Community Event에서 추가 데이터셋을 찾을 수 있거나, 로컬 ImageFolder를 만들어서 자체 데이터셋을 사용할 수도 있습니다. HugGan Community Event에서 가져온 데이터셋인 경우 config.dataset_name을 해당 데이터셋의 저장소 id로 설정하고, 자체 이미지를 사용하는 경우 imagefolder로 설정하세요.

🤗 Datasets는 Image 기능을 사용하여 이미지 데이터를 자동으로 디코딩하고 시각화할 수 있는 PIL.Image로 로드합니다:

>>> import matplotlib.pyplot as plt

>>> fig, axs = plt.subplots(1, 4, figsize=(16, 4))
>>> for i, image in enumerate(dataset[:4]["image"]):
...     axs[i].imshow(image)
...     axs[i].set_axis_off()
>>> fig.show()

이미지 크기가 모두 다르므로 먼저 전처리해야 합니다:

• Resize는 이미지의 크기를 config.image_size에 정의된 값으로 변경합니다.
• RandomHorizontalFlip는 이미지를 무작위로 좌우 반전시켜 데이터셋을 증강합니다.
• Normalize는 모델이 기대하는 [-1, 1] 범위로 픽셀 값을 다시 스케일링하는 데 중요합니다.

>>> from torchvision import transforms

>>> preprocess = transforms.Compose(
...     [
...         transforms.Resize((config.image_size, config.image_size)),
...         transforms.RandomHorizontalFlip(),
...         transforms.ToTensor(),
...         transforms.Normalize([0.5], [0.5]),
...     ]
... )

🤗 Datasets의 set_transform 메서드를 사용하여 학습 중에 preprocess 함수를 실시간으로 적용합니다:

>>> def transform(examples):
...     images = [preprocess(image.convert("RGB")) for image in examples["image"]]
...     return {"images": images}

>>> dataset.set_transform(transform)

이미지를 다시 시각화하여 크기가 조정되었는지 확인해 보세요. 이제 학습을 위해 데이터셋을 DataLoader로 감싸기 준비가 되었습니다!

>>> import torch

>>> train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=config.train_batch_size, shuffle=True)

UNet2DModel 생성

🧨 Diffusers의 사전 학습 모델은 원하는 매개변수를 사용하여 해당 모델 클래스로 쉽게 생성할 수 있습니다. 예를 들어, UNet2DModel을 생성하려면:

>>> from diffusers import UNet2DModel

>>> model = UNet2DModel(
...     sample_size=config.image_size,  # the target image resolution
...     in_channels=3,  # the number of input channels, 3 for RGB images
...     out_channels=3,  # the number of output channels
...     layers_per_block=2,  # how many ResNet layers to use per UNet block
...     block_out_channels=(128, 128, 256, 256, 512, 512),  # the number of output channels for each UNet block
...     down_block_types=(
...         "DownBlock2D",  # a regular ResNet downsampling block
...         "DownBlock2D",
...         "DownBlock2D",
...         "DownBlock2D",
...         "AttnDownBlock2D",  # a ResNet downsampling block with spatial self-attention
...         "DownBlock2D",
...     ),
...     up_block_types=(
...         "UpBlock2D",  # a regular ResNet upsampling block
...         "AttnUpBlock2D",  # a ResNet upsampling block with spatial self-attention
...         "UpBlock2D",
...         "UpBlock2D",
...         "UpBlock2D",
...         "UpBlock2D",
...     ),
... )

샘플 이미지의 형상이 모델 출력 형상과 일치하는지 빠르게 확인하는 것이 좋은 경우가 많습니다:

>>> sample_image = dataset[0]["images"].unsqueeze(0)
>>> print("Input shape:", sample_image.shape)
Input shape: torch.Size([1, 3, 128, 128])

>>> print("Output shape:", model(sample_image, timestep=0).sample.shape)
Output shape: torch.Size([1, 3, 128, 128])

좋습니다! 다음으로, 이미지에 노이즈를 추가하려면 스케줄러가 필요합니다.

스케줄러 생성

스케줄러는 모델을 학습에 사용하는지 추론에 사용하는지에 따라 동작이 다릅니다. 추론 중에는 스케줄러가 노이즈로부터 이미지를 생성합니다. 학습 중에는 스케줄러가 확산 과정의 특정 시점에서의 모델 출력, 즉 샘플을 받아 노이즈 스케줄과 업데이트 규칙에 따라 이미지에 노이즈를 적용합니다.

DDPMScheduler를 살펴보고 add_noise 메서드를 사용해 이전의 sample_image에 무작위 노이즈를 추가해 보겠습니다:

>>> import torch
>>> from PIL import Image
>>> from diffusers import DDPMScheduler

>>> noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
>>> noise = torch.randn(sample_image.shape)
>>> timesteps = torch.LongTensor([50])
>>> noisy_image = noise_scheduler.add_noise(sample_image, noise, timesteps)

>>> Image.fromarray(((noisy_image.permute(0, 2, 3, 1) + 1.0) * 127.5).type(torch.uint8).numpy()[0])

모델의 학습 목표는 이미지에 추가된 노이즈를 예측하는 것입니다. 이 단계의 손실은 다음 식으로 계산할 수 있습니다:

>>> import torch.nn.functional as F

>>> noise_pred = model(noisy_image, timesteps).sample
>>> loss = F.mse_loss(noise_pred, noise)

모델 학습

지금까지 모델 학습을 시작하는 데 필요한 대부분의 요소를 갖추셨으며, 이제 모든 것을 하나로 통합하는 것만 남았습니다.

먼저, 옵티마이저와 학습률 스케줄러가 필요합니다:

>>> from diffusers.optimization import get_cosine_schedule_with_warmup

>>> optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
>>> lr_scheduler = get_cosine_schedule_with_warmup(
...     optimizer=optimizer,
...     num_warmup_steps=config.lr_warmup_steps,
...     num_training_steps=(len(train_dataloader) * config.num_epochs),
... )

그 다음으로, 모델을 평가하는 방법이 필요합니다. 평가에는 DDPMPipeline을 사용해 샘플 이미지 배치 생성 후 그리드로 저장할 수 있습니다:

>>> from diffusers import DDPMPipeline
>>> from diffusers.utils import make_image_grid
>>> import os

>>> def evaluate(config, epoch, pipeline):
...     # Sample some images from random noise (this is the backward diffusion process).
...     # The default pipeline output type is `List[PIL.Image]`
...     images = pipeline(
...         batch_size=config.eval_batch_size,
...         generator=torch.Generator(device='cpu').manual_seed(config.seed), # Use a separate torch generator to avoid rewinding the random state of the main training loop
...     ).images

...     # Make a grid out of the images
...     image_grid = make_image_grid(images, rows=4, cols=4)

...     # Save the images
...     test_dir = os.path.join(config.output_dir, "samples")
...     os.makedirs(test_dir, exist_ok=True)
...     image_grid.save(f"{test_dir}/{epoch:04d}.png")

이제 🤗 Accelerate를 사용하여 TensorBoard 로깅, 그라디언트 축적, 혼합 정밀도 학습을 쉽게 수행할 수 있는 학습 루프에 이 모든 구성 요소를 묶을 수 있습니다. 모델을 Hub에 업로드하려면 저장소 이름과 정보를 가져오는 함수를 작성한 후 Hub에 push하세요.

💡 아래 학습 루프는 길고 복잡해 보일 수 있지만, 나중에 한 줄의 코드로 학습을 시작할 때 그 가치가 드러날 것입니다. 이미 생성하기를 기다릴 수 없다면, 아래 코드를 복사하여 실행해 보세요. 나중에, 예를 들어 모델 학습이 완료될 때까지 기다리는 동안 언제든 돌아와 학습 루프를 더 자세히 살펴볼 수 있습니다. 🤗

>>> from accelerate import Accelerator
>>> from huggingface_hub import create_repo, upload_folder
>>> from tqdm.auto import tqdm
>>> from pathlib import Path
>>> import os

>>> def train_loop(config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler):
...     # Initialize accelerator and tensorboard logging
...     accelerator = Accelerator(
...         mixed_precision=config.mixed_precision,
...         gradient_accumulation_steps=config.gradient_accumulation_steps,
...         log_with="tensorboard",
...         project_dir=os.path.join(config.output_dir, "logs"),
...     )
...     if accelerator.is_main_process:
...         if config.output_dir is not None:
...             os.makedirs(config.output_dir, exist_ok=True)
...         if config.push_to_hub:
...             repo_id = create_repo(
...                 repo_id=config.hub_model_id or Path(config.output_dir).name, exist_ok=True
...             ).repo_id
...         accelerator.init_trackers("train_example")

...     # Prepare everything
...     # There is no specific order to remember, you just need to unpack the
...     # objects in the same order you gave them to the prepare method.
...     model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
...         model, optimizer, train_dataloader, lr_scheduler
...     )

...     global_step = 0

...     # Now you train the model
...     for epoch in range(config.num_epochs):
...         progress_bar = tqdm(total=len(train_dataloader), disable=not accelerator.is_local_main_process)
...         progress_bar.set_description(f"Epoch {epoch}")

...         for step, batch in enumerate(train_dataloader):
...             clean_images = batch["images"]
...             # Sample noise to add to the images
...             noise = torch.randn(clean_images.shape, device=clean_images.device)
...             bs = clean_images.shape[0]

...             # Sample a random timestep for each image
...             timesteps = torch.randint(
...                 0, noise_scheduler.config.num_train_timesteps, (bs,), device=clean_images.device,
...                 dtype=torch.int64
...             )

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

...             with accelerator.accumulate(model):
...                 # Predict the noise residual
...                 noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
...                 loss = F.mse_loss(noise_pred, noise)
...                 accelerator.backward(loss)

...                 if accelerator.sync_gradients:
...                     accelerator.clip_grad_norm_(model.parameters(), 1.0)
...                 optimizer.step()
...                 lr_scheduler.step()
...                 optimizer.zero_grad()

...             progress_bar.update(1)
...             logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0], "step": global_step}
...             progress_bar.set_postfix(**logs)
...             accelerator.log(logs, step=global_step)
...             global_step += 1

...         # After each epoch you optionally sample some demo images with evaluate() and save the model
...         if accelerator.is_main_process:
...             pipeline = DDPMPipeline(unet=accelerator.unwrap_model(model), scheduler=noise_scheduler)

...             if (epoch + 1) % config.save_image_epochs == 0 or epoch == config.num_epochs - 1:
...                 evaluate(config, epoch, pipeline)

...             if (epoch + 1) % config.save_model_epochs == 0 or epoch == config.num_epochs - 1:
...                 if config.push_to_hub:
...                     upload_folder(
...                         repo_id=repo_id,
...                         folder_path=config.output_dir,
...                         commit_message=f"Epoch {epoch}",
...                         ignore_patterns=["step_*", "epoch_*"],
...                     )
...                 else:
...                     pipeline.save_pretrained(config.output_dir)

후, 꽤 많은 코드였네요! 하지만 이제 🤗 Accelerate의 notebook_launcher 함수로 학습을 시작할 준비가 끝났습니다. 학습 루프, 모든 학습 인자, 그리고 학습에 사용할 프로세스 수(이 값은 사용 가능한 GPU 수로 변경할 수 있습니다)를 함수에 전달하세요:

>>> from accelerate import notebook_launcher

>>> args = (config, model, noise_scheduler, optimizer, train_dataloader, lr_scheduler)

>>> notebook_launcher(train_loop, args, num_processes=1)

학습이 완료되면, diffusion 모델이 생성한 최종 🦋 이미지 🦋를 확인해 보세요!

>>> import glob

>>> sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
>>> Image.open(sample_images[-1])

다음 단계

무조건적 이미지 생성은 학습할 수 있는 작업의 한 예시입니다. 🧨 Diffusers Training Examples 페이지를 방문하여 다른 작업과 학습 기법을 탐색해 보세요. 아래는 배울 수 있는 내용들의 예시입니다:

• Textual Inversion는 모델이 특정 시각 개념을 학습하도록 하고 이를 생성된 이미지에 통합하는 알고리즘입니다.
• 주제에 대한 여러 입력 이미지를 제공하여 해당 주제의 개인화된 이미지를 생성하는 기법인 DreamBooth.
• 가이드: 자체 데이터셋으로 Stable Diffusion 모델을 파인튜닝하는 방법
• 매우 큰 모델을 더 빠르게 미세 조정하기 위한 메모리 효율적인 기법인 LoRA 사용 가이드.