DreamBooth를 활용한 이미지 파인튜닝: 코드와 구현 방법

최근 몇 년 동안 딥러닝 기술의 발전은 이미지 생성 및 변형 작업을 한층 더 혁신적으로 만들었습니다. 특히, DreamBooth와 같은 기술을 활용하면 특정한 이미지 스타일이나 특성을 학습하여 새로운 이미지를 생성할 수 있습니다. 이 글에서는 ShivamShrirao의 GitHub 저장소에서 제공하는 DreamBooth 코드를 활용하여 이미지 파인튜닝을 구현하는 방법에 대해 다룹니다.

 

파인튜닝 방법이 궁금하시면 아래 글을 참고해주세요

Stable Diffusion과 파인 튜닝 방법: 완벽 가이드

1. DreamBooth란 무엇인가?

DreamBooth는 이미지 생성 및 변형을 위해 딥러닝 모델을 훈련시키는 기법입니다. 이 기법은 특히 특정 스타일이나 특성을 가진 이미지를 학습하고, 이를 바탕으로 새로운 이미지를 생성하는 데 유용합니다. 이를 통해 예술가, 디자이너, 연구자들은 기존의 이미지를 변형하거나 새로운 스타일의 이미지를 쉽게 생성할 수 있습니다.

스타일 Fine-tuning 예시

2. DreamBooth의 주요 기능

• 스타일 전이: 특정 스타일의 이미지를 학습하여 새로운 이미지에 그 스타일을 적용.
• 이미지 변형: 입력 이미지의 특성을 바탕으로 새로운 이미지를 생성.
• 고해상도 이미지 생성: 기존 모델보다 더 높은 해상도의 이미지를 생성.

3. 필요한 라이브러리 설치

DreamBooth를 사용하려면 몇 가지 필수 라이브러리를 설치해야 합니다. 아래 명령어를 통해 필요한 패키지를 설치할 수 있습니다.

pip install diffusers transformers accelerate
pip install torch torchvision torchaudio
pip install datasets
pip install wandb

 

4. 데이터셋 준비

DreamBooth를 활용하기 위해서는 훈련에 사용할 데이터셋이 필요합니다. 데이터셋은 다양한 스타일의 이미지를 포함해야 하며, 이를 통해 모델이 다양한 특성을 학습할 수 있습니다. 데이터셋은 로컬에 저장하거나, 공공 데이터셋을 사용할 수 있습니다.

5. 코드 설명 및 구현

이제 본격적으로 DreamBooth의 코드를 살펴보겠습니다. 아래 코드는 train_dreambooth.py 파일의 주요 부분을 설명합니다.

5.1. 라이브러리 임포트

import argparse
import logging
import os
import torch
from accelerate import Accelerator
from diffusers import StableDiffusionPipeline
from transformers import AutoTokenizer
from datasets import load_dataset
from torch.utils.data import DataLoader

여기서는 필요한 라이브러리를 임포트합니다. argparse는 명령행 인자를 처리하고, logging은 로그를 기록하는 데 사용됩니다. torch는 PyTorch 라이브러리로 딥러닝 모델을 구현하고 훈련하는 데 사용됩니다.

5.2. 하이퍼파라미터 설정

parser = argparse.ArgumentParser(description="Training script for DreamBooth")
parser.add_argument("--pretrained_model_name_or_path", type=str, required=True, help="Path to pretrained model")
parser.add_argument("--train_data_dir", type=str, required=True, help="Directory containing training data")
parser.add_argument("--output_dir", type=str, required=True, help="Directory to save the trained model")
parser.add_argument("--num_train_epochs", type=int, default=4, help="Number of training epochs")
parser.add_argument("--learning_rate", type=float, default=5e-5, help="Learning rate")
args = parser.parse_args()

여기서는 하이퍼파라미터를 설정합니다. pretrained_model_name_or_path는 사전 학습된 모델의 경로를 지정하고, train_data_dir는 훈련 데이터가 저장된 디렉토리를 지정합니다. output_dir는 훈련된 모델을 저장할 디렉토리를 지정합니다.

5.3. 데이터 로드 및 전처리

dataset = load_dataset('imagefolder', data_dir=args.train_data_dir)
train_dataloader = DataLoader(dataset['train'], batch_size=16, shuffle=True)

여기서는 데이터셋을 로드하고, 데이터 로더를 설정합니다. load_dataset 함수는 주어진 디렉토리에서 이미지를 로드하며, DataLoader는 배치 단위로 데이터를 불러오는 역할을 합니다.

5.4. 모델 초기화

model = StableDiffusionPipeline.from_pretrained(args.pretrained_model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name_or_path)

여기서는 사전 학습된 모델과 토크나이저를 초기화합니다. StableDiffusionPipeline은 이미지 생성 모델을 불러오고, AutoTokenizer는 텍스트 데이터를 토큰화하는 데 사용됩니다.

5.5. 훈련 루프

optimizer = torch.optim.AdamW(model.parameters(), lr=args.learning_rate)
accelerator = Accelerator()

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

for epoch in range(args.num_train_epochs):
    model.train()
    for batch in train_dataloader:
        outputs = model(**batch)
        loss = outputs.loss
        accelerator.backward(loss)
        optimizer.step()
        optimizer.zero_grad()

    print(f"Epoch {epoch} completed.")
    model.save_pretrained(args.output_dir)

여기서는 모델 훈련을 위한 루프를 설정합니다. optimizer는 AdamW 옵티마이저를 사용하여 모델의 파라미터를 최적화합니다. accelerator는 모델과 옵티마이저, 데이터로더를 가속화합니다. 각 에포크마다 모델을 훈련시키고, 손실을 계산하여 역전파를 수행합니다. 훈련이 끝나면 모델을 저장합니다.

5.6. 모델 저장 및 평가

훈련이 완료되면, 모델을 저장하고 평가합니다. 저장된 모델은 새로운 이미지 생성을 위해 사용할 수 있습니다.

6. DreamBooth를 활용한 실제 예제

위 코드를 이용하면 DreamBooth를 사용하여 다양한 스타일의 이미지를 생성해볼 수 있습니다. 예를 들어, 특정 예술가의 스타일을 학습한 모델을 통해 새로운 작품을 생성할 수 있습니다. 이를 위해 다양한 예술가의 작품을 데이터셋으로 사용하고, DreamBooth를 활용하여 모델을 훈련합니다.

7. 최적화 및 튜닝

모델의 성능을 최적화하기 위해 하이퍼파라미터를 조정하고, 더 많은 데이터를 활용할 수 있습니다. 또한, 모델의 구조를 변경하거나, 추가적인 정규화 기법을 도입하여 성능을 향상시킬 수 있습니다.

8. 결론

이 글에서는 DreamBooth를 활용하여 이미지 파인튜닝을 구현하는 방법에 대해 자세히 살펴보았습니다. 제공된 코드를 통해 직접 모델을 훈련시키고, 다양한 스타일의 이미지를 생성해볼 수 있습니다. DreamBooth는 이미지 생성 및 변형 작업에서 강력한 도구이며, 이를 통해 창의적인 작업을 더욱 효과적으로 수행할 수 있습니다. 이 블로그 글이 DreamBooth를 시작하는 데 도움이 되기를 바랍니다. DreamBooth를 활용하여 더 많은 가능성을 탐구하고, 창의적인 프로젝트를 성공적으로 완수하시길 바랍니다.

 

 IP Adapter를 활용한 방법은 아래 포스팅을 참고해주세요.

파인튜닝을 위한 IP-Adapter 활용: Stable Diffusion 개선하기