확산 모델 훈련하기
확산 모델 훈련하기
무조건적 이미지 생성(unconditional image generation)은 확산 모델(diffusion model)에서 특히 인기 있는 활용 분야예요. 훈련에 쓴 데이터셋에 들어 있던 이미지와 비슷한 이미지를 만들어 내는 작업이죠. 보통은 사전 훈련된 모델을 특정 데이터셋으로 파인튜닝(finetuning)하는 게 가장 좋은 결과를 줘요. 이런 체크포인트는 Hub에서 많이 찾을 수 있는데, 마음에 드는 게 없다면 언제든 직접 훈련하면 됩니다.
이 튜토리얼에서는 Smithsonian Butterflies 데이터셋의 일부로 [UNet2DModel]을 처음부터 훈련해서 나만의 🦋 나비 🦋를 만들어 볼 거예요.
[!TIP] 💡 이 훈련 튜토리얼은 Training with 🧨 Diffusers 노트북을 바탕으로 만들어졌어요. 확산 모델이 어떻게 동작하는지 같은 더 자세한 내용과 설명을 보려면 그 노트북을 확인해 보세요!
시작하기 전에, 이미지 데이터셋을 불러오고 전처리하려면 🤗 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()
데이터셋 불러오기
Smithsonian Butterflies 데이터셋은 🤗 Datasets 라이브러리로 손쉽게 불러올 수 있어요:
>>> from datasets import load_dataset
>>> config.dataset_name = "huggan/smithsonian_butterflies_subset"
>>> dataset = load_dataset(config.dataset_name, split="train")
[!TIP] 💡 HugGan Community Event에서 추가 데이터셋을 찾아볼 수도 있고, 로컬
ImageFolder를 만들어 나만의 데이터셋을 쓸 수도 있어요. HugGan Community Event에서 가져온 데이터셋이면config.dataset_name을 그 데이터셋의 레포지토리 id로, 나만의 이미지를 쓴다면imagefolder로 설정하면 됩니다.
🤗 Datasets는 [~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의 [~datasets.Dataset.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",
... ),
... )
샘플 이미지의 shape가 모델 출력 shape와 일치하는지 가볍게 확인해 보는 것도 좋은 습관이에요:
>>> 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])
좋아요! 이제 이미지에 노이즈를 더할 스케줄러(scheduler)가 필요하네요.
스케줄러 만들기
스케줄러는 모델을 훈련에 쓰느냐, 추론에 쓰느냐에 따라 다르게 동작해요. 추론 중에는 스케줄러가 노이즈로부터 이미지를 만들어 내고, 훈련 중에는 확산 과정의 특정 시점에서 나온 모델 출력(샘플)을 받아 *노이즈 스케줄(noise schedule)*과 *업데이트 규칙(update rule)*에 따라 이미지에 노이즈를 적용해요.
[DDPMScheduler]를 살펴보면서, 아까 그 sample_image에 add_noise 메서드로 무작위 노이즈를 더해 볼게요:
>>> 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]으로 샘플 이미지 한 배치를 생성한 뒤 그리드(grid)로 저장할 수 있습니다:
>>> 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하면 됩니다.
[!TIP] 💡 아래 훈련 루프는 길고 겁이 나 보일 수 있지만, 나중에 단 한 줄의 코드로 훈련을 시작하게 될 때 그만한 가치가 있답니다. 기다리기 어렵고 바로 이미지를 만들어 보고 싶다면 아래 코드를 복사해 실행해도 좋아요. 나중에 모델 훈련이 끝나기를 기다리는 동안 이 훈련 루프를 더 자세히 들여다보면 됩니다. 🤗
>>> 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의 [~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)
훈련이 끝나면 나만의 확산 모델이 만들어 낸 최종 🦋 이미지 🦋를 확인해 보세요!
>>> import glob
>>> sample_images = sorted(glob.glob(f"{config.output_dir}/samples/*.png"))
>>> Image.open(sample_images[-1])
다음 단계
무조건적 이미지 생성은 훈련할 수 있는 작업 중 한 예시일 뿐이에요. 다른 작업과 훈련 기법을 탐험해 보고 싶다면 🧨 Diffusers Training Examples 페이지를 방문해 보세요. 배울 수 있는 예시 몇 가지를 소개할게요:
- Textual Inversion는 모델에 특정 시각적 개념을 가르쳐 생성 이미지에 통합하는 알고리즘이에요.
- DreamBooth는 피사체의 입력 이미지 몇 장으로 그 피사체의 개인화된 이미지를 생성하는 기법입니다.
- Guide는 나만의 데이터셋으로 Stable Diffusion 모델을 파인튜닝하는 방법을 다뤄요.
- Guide는 아주 큰 모델을 더 빠르게 파인튜닝하는 메모리 효율적 기법인 LoRA 사용법을 설명합니다.
더 알아보기 (Learn more)
- 🧨 Diffusers Training Examples: 무조건적 생성 외에도 훈련할 수 있는 다양한 작업과 기법.
- Textual Inversion: 특정 시각 개념을 모델에 가르치는 알고리즘.
- DreamBooth: 피사체 이미지 몇 장으로 개인화된 이미지를 생성하는 기법.
- LoRA 가이드: 메모리 효율적으로 대형 모델을 파인튜닝하는 기법.