노이즈 스케줄러부터 꽃 생성까지 확장한 Diffusion 구현기

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

TL;DR

작성자는 CIFAR 이미지 분류기를 만든 뒤 Diffusion model의 Noise Scheduler, Forward Diffusion, UNET을 직접 구현하며 이미지 생성 프로젝트를 확장했습니다. PyTorch indexing 오류와 장시간 학습 문제를 겪었지만 TensorBoard로 loss 추세를 확인하고 SGD를 AdamW로 바꾸었으며, M2 Air의 MPS를 활용해 연산 속도도 개선했습니다. CIFAR의 32×32 출력에서 Flowers102의 256×256 이미지로 전환하는 과정에서는 VAE를 UNET과 결합했다가 색이 뒤섞인 결과를 얻어 VAE를 제거하고 UNET을 다시 구성했습니다. 약 561번째 epoch와 반복 학습 100회 이후 꽃 형태가 나타나기 시작했지만 아직 출력 품질과 loss 정체 문제는 남아 있습니다.

실용적 조언

학습 loss를 직전 값과 단순 비교하는 대신 TensorBoard로 전체 추세를 기록하면 장기 학습에서 발산이나 정체 시점을 더 쉽게 파악할 수 있습니다.
PyTorch에서 연산하는 tensor와 model이 같은 device에 있는지 확인해야 하며, MPS나 CUDA를 사용할 때 device 이동을 일관되게 관리해야 합니다.
새 데이터셋으로 확장할 때는 이미지 해상도와 UNET 구조를 함께 조정하고, VAE를 결합하는 경우 압축 표현의 shape와 복원 결과가 UNET의 입력·출력과 맞는지 단계별로 확인해야 합니다.

섹션별 상세

작성자는 CIFAR 데이터셋으로 작은 이미지 분류기를 만든 뒤 이미지 생성 프로젝트에 도전하면서 Diffusion model의 구성 요소를 단계별로 학습했습니다. Noise Scheduler가 이미지에 노이즈를 점진적으로 추가하고, Forward Diffusion과 학습 루프를 거쳐 UNET이 노이즈를 줄이는 구조를 먼저 정리했습니다. 이론을 읽은 뒤 다른 구현의 알고리즘과 프로그램 흐름만 참고하고 코드를 직접 작성했으며, PyTorch indexing과 변수 처리에서 반복적으로 오류를 겪었습니다.

Noise Scheduler 구현 뒤에는 Forward Diffusion을 추가했고, 가장 큰 시간은 데이터셋에 맞는 UNET을 조정하는 데 들었습니다. 병목층 길이, convolution 수, 추가 layer 같은 구조를 바꾸며 실험했지만 학습이 오래 걸려 500~600번째 epoch에 모델이 제대로 작동하지 않는다는 사실을 발견한 경우도 있었습니다. 처음에는 loss가 직전 값보다 커지면 빨간색, 작아지면 초록색으로 콘솔에 표시했지만 전체 추세를 파악하기 어려웠고, TensorBoard 그래프로 바꾸면서 loss 변화를 더 명확히 확인했습니다.

작성자는 학습 중 SGD를 AdamW로 교체하고 M2 Air에서 PyTorch의 MPS를 사용해 속도를 높였습니다. 두 tensor가 같은 device에 있어야 연산할 수 있다는 조건을 몰라 여러 부분이 깨졌지만, 코드를 수정해 GPU 연산을 실행할 수 있게 됐습니다. 이후 CIFAR의 32×32 이미지에서 Flowers102의 256×256 이미지로 확장하면서 학습 속도와 메모리 부담이 커졌고, 반복적인 재학습과 구조 수정이 필요했습니다.

Flowers102에 맞추되 UNET 변경을 줄이기 위해 VAE를 먼저 학습하고 연결했지만, 결합 뒤 출력이 색의 덩어리처럼 나타났고 어느 시점에는 loss도 더 이상 감소하지 않았습니다. 작성자는 며칠간 디버깅한 뒤 VAE를 제거하고 256×256 이미지를 직접 처리하도록 UNET을 다시 작성했습니다. 약 561번째 epoch에 이르러 꽃과 비슷한 형태가 나타나기 시작했지만 여러 문제가 남아 있으며, 현재 구현은 GitHub 저장소 TinyDPPM에서 확인할 수 있습니다.

용어 해설

노이즈 스케줄러(Noise Scheduler): — Diffusion model에서 이미지에 노이즈를 어느 단계에서 얼마나 추가할지 정하는 구성 요소입니다. Forward Diffusion 동안 원본 이미지에 점진적으로 노이즈를 섞어 순수한 노이즈에 가깝게 만들며, 이후 UNET이 각 단계의 노이즈를 줄이는 학습 과정에 필요한 시간별 입력을 만듭니다.
순방향 확산(Forward Diffusion): — 원본 이미지에 여러 단계로 노이즈를 추가해 점차 알아보기 어려운 상태로 바꾸는 과정입니다. 글에서는 Noise Scheduler가 이 과정을 제어하고, 학습 중 UNET이 각 단계에서 제거해야 할 노이즈 패턴을 익히도록 입력을 생성하는 흐름으로 다뤄집니다.
UNET: — Diffusion model에서 노이즈가 섞인 이미지의 노이즈를 예측하거나 줄이는 데 사용하는 신경망 구조입니다. 기본 구조만으로는 데이터셋에 맞는 출력 품질을 보장하기 어려워 병목층 길이, convolution 수, 추가 layer 등을 데이터 해상도와 형태에 맞게 조정해야 합니다.
Variational Autoencoder: — 이미지를 더 작고 고차원적인 표현으로 압축했다가 복원하는 생성 모델 계열입니다. 글에서는 VAE를 먼저 학습한 뒤 UNET과 결합했지만 색이 뒤섞인 출력이 나왔고, 결국 VAE를 제거하고 256×256 입력을 직접 처리하도록 UNET을 다시 작성했습니다.
TensorBoard: — 학습 중 기록한 loss 같은 지표를 시간에 따른 그래프로 확인하는 도구입니다. 작성자는 콘솔 색상으로 직전 값과 비교하던 방식 대신 TensorBoard에서 loss의 전체 추세를 확인하면서 모델이 제대로 학습되는지 더 쉽게 판단할 수 있었다고 말합니다.
MPS: — Apple Silicon GPU에서 PyTorch 연산을 실행하기 위한 API입니다. 글에서는 M2 Air에서 MPS를 사용했고, 서로 연산하는 두 tensor가 같은 device에 있어야 한다는 조건을 뒤늦게 파악한 뒤 코드를 수정해 학습 속도를 높였습니다.

언급된 도구

PyTorch중립

Tensor 연산, UNET 학습, MPS와 CUDA 기반 device 실행에 사용했습니다.

TensorBoard추천

학습 loss의 전체 추세를 그래프로 확인하는 데 사용했습니다.

Zed중립

Diffusion model 구현을 시작하고 코드를 작성하는 편집 환경으로 사용했습니다.

언급된 리소스

GitHubTinyDPPM

문서Deep Learning with PyTorch: A Step-by-Step Beginner's Guide