What are Diffusion Models?
Diffusion Model 확산 모델은 상대적으로 최근에 '생성 모델'로 알려진 알고리즘 그룹에 추가되었다.
"생성 모델"의 목표 : 여러 훈련 예제가 주어졌을 때 이미지나 오디오와 같은 데이터를 생성하는 방법을 배우는 것이다.
'생성 모델' 의 목표 = 여러 훈련 예제가 주어졌을 때 이미지나 오디오와 같은 데이터를 생성하는 방법을 배우는 것
Diffusion Model의 Main Idea
: Diffusion process의 반복적인(iterative) 특성
생성은 랜덤 노이즈에 시작되지만, 출력 이미지가 나타나기 까지 여러 단계를 걸쳐 점진적으로 개선된다.
각 단계에서, 모델은 우리가 현재 입력에서 완전히 denoise된 버전으로 어떻게 이동하는지를 추정
모델 훈련시키는 과정
모델을 훈련시키는 것은 상대적으로 다른 유형의 생성모델에 비해 간단하다.
- 훈련 데이터에서 일부 이미지를 로드
- 노이즈를 다른 amount 추가
- 모델이 극도로 시끄러운 이미지와 완벽에 가까운 이미지 모두를 어떻게 denoise 하는지를 추정하는 것이 목적!!
- 모델에 노이즈를 추가한 버전을 입력으로 추가
- 모델이 노이즈가 추가된 입력을 denosing하는데 얼마나 성능이 뛰어난지를 평가
- 위의 정보를 사용하여 모델 가중치 업데이트
훈련된 모델로 새로운 이미지를 생성하기 위해선,
완전히 무작위적인 입력으로 시작하여 모델을 통해 반복적으로 공급하고 모델 예측을 기반으로 조금씩 매번 업데이트를 진행해야 한다.
위 과정을 가능한 적은 단계로 좋은 이미지를 생성할 수 있도록 하는 여러 샘플링 방법이 존재한다.
Introduction to Diffusers
! generate images of cute butterflies를 통해 diffusion model 생성 !
이 노트북에서 배울 내용
- 강력한 맞춤형 diffusion model 파이프라인을 직접 실행 가능
- 자신만의 미니 파이프라인 생성 가능
- 확산 모델의 핵심 아이디어 살펴보기
- 훈련을 위해 허브에서 데이터 불러오기
- 스케줄러를 사용하여 이 데이터에 노이즈 추가하는 방법 탐색하기
- UNet 모델 생성 및 훈련
- 이 모든 요소를 결합하여 작동하는 파이프라인 만들기
Step 1 : Setup
%pip install -qq -U diffusers datasets transformers accelerate ftfy pyarrow==9.0.0ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
cudf-cu12 24.4.1 requires pyarrow<15.0.0a0,>=14.0.1, but you have pyarrow 9.0.0 which is incompatible.
오류 발생❗️→ ‘pyarrow’ 버전 불일치로 인해 발생하는 오류라고 한다. 현재 설치된 ‘pyarrow’버전이 ‘cudf-cu12’라이브러리에서 요구하는 버전과 호환되지 않기 때문!
⇒ ‘14.0.1’ 버전으로 업데이트
User Access Token
https://huggingface.co/settings/tokens
해당 링크로 가서 Token 받아오기
위에서 받아온 general_write Token을 이용해서 로그인하기!
from huggingface_hub import notebook_login
notebook_login()모델의 체크포인트를 업로드하기 위해 Git-LFS 설치하기
%%capture
!sudo apt -qq install git-lfs
!git config --global credential.helper store이미지 시각화를 위한 몇 가지 유용한 함수 정의하기
import numpy as np #수치 계산을 위한 라이브러리
import torch #딥러닝을 위한 라이브러리
import torch.nn.functional as F #PyTorch의 신경망 기능
from matplotlib import pyplot as plt #데이터 시각화 라이브러리
from PIL import Image #이미지 처리를 위한 Python Imaging Libaray
#이미지 시각화 함수 정의
# 배치로 제공된 이미지 'x'를 받아 그리드 생성 및 PIL 이미지로 변환
def show_images(x):
"""Given a batch of images x, make a grid and convert to PIL"""
x = x * 0.5 + 0.5 # Map from (-1, 1) back to (0, 1)
grid = torchvision.utils.make_grid(x) #이미지를 그리드로 생성
#그리드를 PIL 이미지로 변환하기 위해 형식 변환
grid_im = grid.detach().cpu().permute(1, 2, 0).clip(0, 1) * 255
grid_im = Image.fromarray(np.array(grid_im).astype(np.uint8))
return grid_im
#PIL 이미지 목록을 받아 하나의 이미지로 병합하여 보기 쉽게 생성
def make_grid(images, size=64):
"""Given a list of PIL images, stack them together into a line for easy viewing"""
# 각 이미지를 가로로 나열할 수 있는 새 이미지 생성
output_im = Image.new("RGB", (size * len(images), size))
# 각 이미지 순회하며 output_im에 붙여넣음
for i, im in enumerate(images):
output_im.paste(im.resize((size, size)), (i * size, 0))
return output_im #병합된 이미지 반환
# Mac users may need device = 'mps' (untested)
# 저는 Mac user이기에 mps를 사용하는 코드로 바꿔주었습니다.
device = torch.device("mps" if torch.backends.mps.is_available() else "cuda" if torch.cuda.is_available() else "cpu")- show_images : 배치로 제공된 이미지를 그리드로 만들고 PIL 이미지로 변환하는 함수
- make_grid : PIL 이미지 목록을 가로로 나열하여 하나의 이미지로 병합하는 함수
- device : 사용할 장치를 설정하는 코드PIL이란?
: 파이썬 이미지 처리 라이브러리 → 이미지 분석 및 처리를 쉽게 할 수 있게 해줌
MVP (Minimum Viable Pipeline)
🤗 Diffusers 의 핵심 API는 세 가지 주요 구성 요소로 나뉩니다.:
- Pipelines: 인기 있는 훈련된 diffusion model에서 샘플을 빠르게 생성할 수 있도록 설계된 고수준 클래스이다. 사용자 친화적인 방식으로 샘플을 생성한다.
=> end-user의 경우 pipeline만 사용하면 됨 - Models: 새로운 diffusion model을 훈련시키기 위한 인기 있는 architecture, e.g. UNet.
- Schedulers: 추론 중 노이즈에서 이미지를 생성하는 다양한 기술과 훈련을 위해 노이즈가 있는 이미지를 생성하는 기술을 포함한다.
⇒ Model과 Scheduler를 이용해서 미세튜닝을 할 수 있어야 함.
End-user는 pipeline을 가져와서 샘플을 생성할 수 있음
from diffusers import DDPMPipeline
# 나비 파이프라인 로드
butterfly_pipeline = DDPMPipeline.from_pretrained(
"johnowhitaker/ddpm-butterflies-32px"
).to(device)
# 이미지 8장 생성
images = butterfly_pipeline(batch_size=8).images
# 결과 보기
make_grid(images)Step 2 : Download a training dataset
import torchvision #컴퓨터 비전 작업을 위한 PyTorch 라이브러리
from datasets import load_dataset
from torchvision import transforms #이미지를 다루는 좋은 기능을 담고 있는 라이브러리
#데이터셋 로드
dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")
# 또는 로컬 폴더에서 이미지 로드
# dataset = load_dataset("imagefolder", data_dir="path/to/folder")
# 32픽셀 정사각형 이미지를 훈련에 사용 ( 더 큰 크기도 시도 가능)
image_size = 32
# GPU 메모리가 부족한 경우 배치 크기를 줄일 수 있음
batch_size = 64
# 데이터 증강 정의
preprocess = transforms.Compose( # 여러 변환을 순차적으로 적용
[
transforms.Resize((image_size, image_size)), # 크기 조정
transforms.RandomHorizontalFlip(), # 랜덤으로 좌우 반전 (데이터 증강)
transforms.ToTensor(), # 텐서로 변환 (0,1)
transforms.Normalize([0.5], [0.5]), # (-1, 1)로 정규화 / 평균:0.5, 표준편차:0.5
]
)
#이미지 변환 함수 정의
def transform(examples):
images = [preprocess(image.convert("RGB")) for image in examples["image"]]
return {"images": images}
#데이터셋에 변환 적용
dataset.set_transform(transform)
# 변환된 이미지를 **배치 단위**로 제공하는 데이터 로더 생성
train_dataloader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size, shuffle=True
)이 코드는 나비 이미지를 로드하고, 이를 훈련에 사용할 수 있도록 적절하게 변환하며, 데이터로더를 생성하여 배치 단위로 이미지를 제공하는 작업을 수행
- transforms : 이미지를 순차적으로 mix-up 가능
- torch.Tensor : 일반적으로 (채널, 높이, 너비) 순서
- PIL.Image : 0-255 정수값 → torch.Tensor : 일반적으로 0-1 실수값
- 왜 파이토치 텐서로 바꾸느냐. → 노이즈는 가우시안 분포이기 때문에 [-1, 1] 범위로 정규화시킴
데이터 로더에서 배치 데이터를 가져와 그 형태를 출력하고, 이미지를 시각화
# 배치 데이터 가져오기
xb = next(iter(train_dataloader))["images"].to(device)[:8]
print("X shape:", xb.shape)
show_images(xb).resize((8 * 64, 64), resample=Image.NEAREST)- next(iter(train_dataloader)): train_dataloader에서 첫 번째 배치를 가져옵니다. 이 배치는 사전 처리된 이미지들을 포함합니다.
- ["images"]: 배치에서 이미지 데이터를 추출합니다.
- .to(device): 이미지를 GPU 또는 CPU로 이동시킵니다. device는 이전에 정의된 장치(GPU 또는 CPU)입니다.
- [:8]: 첫 8개의 이미지만 선택합니다.
- show_images(xb): 선택된 배치 이미지를 그리드 형태로 변환합니다.
- .resize((8 * 64, 64), resample=Image.NEAREST): 이미지를 8개의 이미지가 가로로 배열된 크기(8 * 64, 64)로 조정합니다. resample=Image.NEAREST는 최근접 이웃 보간법을 사용하여 이미지를 재조정합니다.
Step 3 : Define the Scheduler 스케쥴러 정의
훈련의 목표는 입력 이미지에 노이즈를 추가한 후, 노이즈가 추가된 이미지를 모델에 입력을 넣는 것이다. 그리고 추론 단계에서는, 모델 예측을 사용하여 반복적으로 노이즈를 제거할 예정이다. diffusers에서는, 이러한 모든 과정이 모두 scheduler에 의해 처리된다.
💡 [Denoising Diffusion Probablistic Models](https://arxiv.org/abs/2006.11239) ⇒ 논문 살펴보기
노이즈 schedule은 각 다른 timestamp에서 얼마나 많은 노이즈가 추가되는지를 결정한다. 다음은 'DDPM' 훈련 및 샘플링을 위한 기본 설정을 사용하여 Scheduler를 생성하는 방법이다.
Scheduler
: Timestep에 따라 노이즈를 추가하거나 제거하는 과정을 제어하며, 이를 통해 모델 훈련과 추론이 이루어진다.
from diffusers import DDPMScheduler
#노이즈 스케쥴러 생성 - 훈련 시 사용할 타임스탭 = 1000
noise_scheduler = DDPMScheduler(num_train_timesteps=1000)
plt.plot(noise_scheduler.alphas_cumprod.cpu() ** 0.5, label=r"${\sqrt{\bar{\alpha}_t}}$")
plt.plot((1 - noise_scheduler.alphas_cumprod.cpu()) ** 0.5, label=r"$\sqrt{(1 - \bar{\alpha}_t)}$")
plt.legend(fontsize="x-large");노이즈가 추가된 이미지 생성
# timestep 생성 - 0~999 사이를 균등하게 나눈 8개의 값을 생성
# xb : 배치 단위로 불러온 이미지
timesteps = torch.linspace(0, 999, 8).long().to(device)
# 노이즈 생성
noise = torch.randn_like(xb)
# 노이즈가 추가된 이미지 생성 (타임스텝에 맞춰서 노이즈 추가)
noisy_xb = noise_scheduler.add_noise(xb, noise, timesteps)
print("Noisy X shape", noisy_xb.shape)
# 이미지 시각화
show_images(noisy_xb).resize((8 * 64, 64), resample=Image.NEAREST)
Step 4 : Define the Model 모델 정의
대부분의 확산 모델은 U-net의 변형인 아키텍쳐를 사용하며, 여기서도 이것을 사용한다.
요약하자면:
- 모델은 입력 이미지를 여러 개의 ResNet 레이어 블록을 통해 처리하며, 각 블록은 이미지 크기를 절반으로 줄인다.
- 그런 다음 동일한 수의 블록을 통해 이미지를 다시 upsampling 한다.
- downsample 경로의 특징들을 upsample 경로의 일치하는 레이어와 연결하는 skip connection이 존재
이 모델의 주요 특징 중 하나는 입력과 동일한 크기의 이미지를 예측한다는 점이다.
Diffusers 라이브러리는 PyTorch에서 원하는 아키텍쳐를 생성할 수 있는 편리한 UNet2DModel 클래스를 제공한다.
원하는 이미지 크기에 맞는 U-net을 생성해보자. down_block_types는 downsample 블록에 해당하며 (위 다이어그램의 녹색 부분), up_block_types 는 upsample 블력에 해당한다. (위 다이어그램의 빨간색 부분):
UNet에 대한 간략한 설명
U-Net은 주로 이미지 분할 작업에 사용되는 딥러닝 아키텍쳐이다.
U-Net의 구조는 두 가지 주요 부분으로 구성된다. 다운샘플링 경로(Contracting path) & 업샘플링 경로 (UpSampling Path)
- 다운샘플링 경로 (Contracting Path)
- 구조 : 여러 개의 컨볼루션 블록으로 구성되며, 각 블록은 컨볼루션 레이어, 활성화 함수 (ReLU) 및 최대 풀링 (Max Pooling) 레이어로 구성된다.
- 기능 : 입력 이미지의 공간적 크기를 점차적으로 줄여가며 특징을 추출한다. 각 블록은 이미지의 크기를 절반으로 줄이고 특징 맵의 수를 두 배로 증가시킨다.
- 업샘플링 경로 (Expanding Path)
- 구조 : 다운 샘플링 경로의 각 레이어에서 추출된 특징 맵을 업샘플링 경로의 대응하는 레이어에 연결한다.
- 기능 : 원본 이미지의 세부 정보를 보존하며, 업샘플링 과정에서 손실된 공간적 정보를 복원하는데 도움을 준다. 이는 모델이 더 정확한 예측을 할 수 있게 한다.
U-Net의 전체 구조
U-Net의 전체 구조는 “U” 모양을 띠고 있다.
이는 다운샘플링 경로와 업샘플링 경로가 중앙에서 연결되는 형태를 의미한다.
- 입력 이미지 : 원본 이미지가 모델에 입력된다.
- 다운샘플링 경로 : 컨볼루션과 최대 풀링을 통해 이미지의 크기를 줄이고 특징을 추출한다.
- 보틀넥 (BottleNeck) : 다운샘플링 경로의 최종 출력으로, 가장 작은 크기의 특징 맵을 얻는다.
- 업샘플링 경로 : 업컨볼루션과 컨볼루션을 통해 이미지의 크기를 복원한다. 이 과정에서 스킵 연결을 통해 다운샘플링 경로의 특징 맵을 사용한다.
- 출력 이미지 : 최종적으로 입력 이미지와 동일한 크기의 예측 이미지를 출력한다.
U-Net 모델 생성
from diffusers import UNet2DModel
# 모델 생성
model = UNet2DModel(
sample_size=image_size, # 목표 이미지 해상도
in_channels=3, # 입력 채널 수 (이미지 데이터이므로 Channel = 3)
out_channels=3, # 출력 채널 수
layers_per_block=2, # 각 UNet 블록당 사용할 ResNet 레이어 수
block_out_channels=(64, 128, 128, 256), # 각 블록당 출력 채널 수 (채널 수가 많을수록 파라미터 수가 증가)
down_block_types=(
"DownBlock2D", # 일반 ResNet 다운샘플링 블록
"DownBlock2D",
"AttnDownBlock2D", # 공간적 self-attention이 있는 ResNet 다운샘플링 블록
"AttnDownBlock2D",
),
up_block_types=(
"AttnUpBlock2D",
"AttnUpBlock2D", # 공간적 self-attention이 있는 ResNet 업샘플링 블록
"UpBlock2D",
"UpBlock2D", # 일반 ResNet 업샘플링 블록
),
)
model.to(device);- 고해상도 입력을 다룰 때는 더 많은 다운샘플링 블록과 업샘플링 블록을 사용하고, 메모리 사용량을 줄이기 위해 attention 레이어는 가장 낮은 해상도(바닥)레이어에만 유지하는 것이 좋다.
- 데이터 배치와 일부 랜덤 타임스텝을 입력으로 제공하면 입력데이터와 동일한 형태의 출력을 생성하는지 확인 가능하다.
예측 수행 및 검증
with torch.no_grad():
model_prediction = model(noisy_xb, timesteps).sample
model_prediction.shape- with torch.no_grad()::
- 이 블록 내에서는 PyTorch의 자동 미분 기능(gradient calculation)이 비활성화됩니다. 이는 주로 예측이나 평가 단계에서 사용되며, 메모리 사용량을 줄이고 계산 속도를 높일 수 있습니다. 이 블록 내에서는 모델 파라미터가 업데이트되지 않습니다.
- model_prediction = model(noisy_xb, timesteps).sample:
- model에 noisy_xb(노이즈가 추가된 입력 데이터)와 timesteps(타임스텝)을 입력으로 제공하여 예측을 수행합니다.
- .sample은 모델의 예측 결과를 나타냅니다.
- 이 예측 결과는 모델이 입력 데이터의 노이즈를 제거하려고 시도한 결과입니다.
Step 5. Create a Training Loop 📌
아래는 PyTorch에서의 일반적인 최적화 루프입니다. 여기서 우리는 데이터를 배치 단위로 처리하며 각 단계에서 optimizer를 사용해 모델의 파라미터를 업데이트합니다.
- 이 경우, 학습률이 0.0004인 AdamW optimizer를 사용한다.
각 데이터 배치에 대해 다음을 수행한다.
- 랜덤 time steps를 샘플링 → 8개의 배치 사이즈 이미지를 랜덤하게 timestep로 샘플링
- 해당 타임스텝에 따라 데이터를 노이즈화
- 노이즈가 추가된 데이터를 UNet 모델에 전달
- 추가된 노이즈와 예측된 노이즈와 비교, 평균 제곱 오차 (MSE)를 손실함수로 사용
- loss.backward() and optimizer.step()를 통해 모델 파라미터를 업데이트
이 과정에서, 나중에 그래프를 그릴 수 있도록 손실값을 기록한다.
# 노이즈 스케줄러 설정 - DDPM
noise_scheduler = DDPMScheduler(
# 1000개의 타임스텝을 사용하여 훈련. 노이즈 스케줄 방식 설정.
num_train_timesteps=1000, beta_schedule="squaredcos_cap_v2"
)
# 훈련 루프 - AdamW 옵티마이저를 사용하여 학습률을 0.0004로 설정.
optimizer = torch.optim.AdamW(model.parameters(), lr=4e-4)
losses = []
for epoch in range(30):
# 각 에포크 동안 train_dataloader에서 배치를 가져온다.
for step, batch in enumerate(train_dataloader):
# 원본 이미지를 장치로 가져옴
clean_images = batch["images"].to(device)
# 이미지에 추가할 노이즈 샘플링
# torch.randn을 사용해 clean_images와 동일한 형태의 노이즈를 생성
noise = torch.randn(clean_images.shape).to(clean_images.device)
# 배치 크기를 가져온다.
bs = clean_images.shape[0]
# 각 이미지에 대해 랜덤 timesteps 샘플링
timesteps = torch.randint(
0, noise_scheduler.num_train_timesteps, (bs,), device=clean_images.device
).long()
# 각 타임스텝에서 노이즈 크기에 따라 깨끗한 이미지에 노이즈 추가
noisy_images = noise_scheduler.add_noise(clean_images, noise, timesteps)
# 노이즈가 추가된 이미지를 UNet 모델에 통과시켜 예측 얻기
# 노이즈 이미지와 timestep 두 가지를 인자로 받음
noise_pred = model(noisy_images, timesteps, return_dict=False)[0]
# 손실 계산 - MSE 이용
loss = F.mse_loss(noise_pred, noise)
# 손실에 대한 그래디언트 계산
loss.backward(loss)
losses.append(loss.item())
# optimizer로 모델 파라미터 업데이트
optimizer.step()
optimizer.zero_grad()
# 매 5번째 에포크마다 평균 손실을 출력
if (epoch + 1) % 5 == 0:
loss_last_epoch = sum(losses[-len(train_dataloader) :]) / len(train_dataloader)
print(f"Epoch:{epoch+1}, loss: {loss_last_epoch}")- noise_pred : 노이즈를 예측한 것
손실을 그래프로 그려보면, 모델이 처음에는 빠르게 개선되고 이후에는 더 느린 속도로 계속 개선되는 것을 볼 수 있다. (오른쪽의 로그 스케일을 사용하면 더 명확하게 보인다.)
Step 6. Generate Images
다음 블로그 포스트에 이어서 작성하겠다~!! 🚀🚀
'학부연구생 > Diffusion Model' 카테고리의 다른 글
| DDPM (Denoising Diffusion Probabilistic Models) (1) | 2024.07.01 |
|---|---|
| How to create pipeline? (0) | 2024.07.01 |