Main/딥러닝

[논문공부]Diffusion model 코드와 함께 공부하기 - 1

역학맨 2025. 4. 23. 13:47

안녕하세요 역학맨입니다. 오늘은 최근에 Diffusion model에 관해 관심이 생겨 여러가지 논문을 읽고 있습니다. 

논문에서 나타나는 수식만 보다보니 머리가 너무 어지럽고 잘 와닿지않아서 코드를 함께 이해하면서 진행해보려고 합니다. 아무쪼록 이 게시글이 이해하는데 도움이 되셨으면 좋겠습니다. 

 

*딥러닝 분야를 공부한지 얼마되지않아 코드와 관련한 내용에 대해 모르는 부분이 많아 최대한 자세하게 정리했습니다.

 

Diffusion model 

 

import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import make_swiss_roll
from helper_plot import hdr_plot_style
import torch.autograd as autograd
import torch
import torch.nn as nn
import torch.optim as optim

hdr_plot_style()
# Sample a batch from the swiss roll
def sample_batch(size, noise=1.0):
    x, _= make_swiss_roll(size, noise=noise)
    return x[:, [0, 2]] / 10.0
# Plot it
data = sample_batch(10**4).T
plt.figure(figsize=(16, 12))
plt.scatter(*data, alpha=0.5, color='red', edgecolor='white', s=40)
plt.show()

해당 코드를 진행하게되면 noise가 진행되지않은 Swiss roll 원본 이미지를 다음과 같이 확인할 수 있습니다. 

 

 

 

Forward Process

이번 포스팅에서는 Markov chain, Bayes' rule을 활용하여 수식적으로 증명하는 과정을 생략하고 결론만을 이용해서 코드 상에서 어떻게 나타나는지 확인해보도록 할게요 ( 추후 수식적 증명을 포스팅하겠습니다)

 

$$q(x_{t}|x_{t-1})= T_{\pi)}(x_{t}|x_{t-1} ; \beta _{t}), q(x_{t}|x_{t-1})= N(x_{t};\sqrt{1-\beta _{t}}x_{t-1},\beta_{t}I$$

 

정해진 스텝 수와 노이즈 스케쥴을 통해 점진적으로 원본 이미지에 노이즈를 추가하여 이미지를 파괴하는 과정을 나타냅니다. 

 

이제 이 원본이미지에 정해진 $\beta$값을(=noise schedule) 지정하여 노이즈를 스텝에 따라 어떻게 나타나게 되는지 확인해볼 것입니다. 

 

 

 

def forward_process(x_start, n_steps, noise=None):
    """ Diffuse the data (t == 0 means diffused for 1 step) """
    x_seq = [x_start]
    for n in range(n_steps):
        x_seq.append((torch.sqrt(1 - betas[n]) * x_seq[-1]) + (torch.sqrt(betas[n]) * torch.rand_like(x_start)))
    return x_seq
n_steps = 100
betas = torch.tensor([0.035] * n_steps)
dataset = torch.Tensor(data.T).float()
x_seq = forward_process(dataset, n_steps, betas)
fig, axs = plt.subplots(1, 10, figsize=(28, 3))
for i in range(10):
    axs[i].scatter(x_seq[int((i / 10.0) * n_steps)][:, 0], x_seq[int((i / 10.0) * n_steps)][:, 1], s=10);
    axs[i].set_axis_off(); axs[i].set_title('$q(\mathbf{x}_{'+str(int((i / 10.0) * n_steps))+'})$')
plt.show()

 

 

forward_process의 정의를 보게되면 x_seq에 초기 이미지(x_start = dataset)를 저장한 이후에 반복문을 통해 정해진 스텝에 따라  append 함수를 사용하여 x_seq[-1](가장 최근 이미지)를 불러와 노이즈를 더하고 정해진 스케쥴에 따른 랜덤한 노이즈를 더한 이후에 x_seq에 노이즈를 씌운 이미지를 저장하게 됩니다. 

$$x_{t}=\sqrt{\alpha_{t}}x_{t-1}+\sqrt{1-\alpha_{t}}\epsilon$$

$$\beta = 1-\alpha$$

 

 

*Code 정리

append: x = [1,2,3] 일 때 x.append[4]를 진행하게 되면 x = [1,2,3,4]로 값이 변하게 됩니다. 

x_seq[-1] : list내에서 가장 마지막 성분(가장 최근 이미지)을 불러옵니다. 

torch.rand_like(x_start) : x_start와 같은 크기의 랜덤 숫자를 불러오는 명령

torch.tensor는 GPU의 사용을 할 수 있게 리스트를 구성하는 것입니다.

torch.Tensor(data.T): .T는 Transpose를 의미한다. tensor는 3차원 형태 이상의 데이터

 

# calculations for diffusion q(x_t | x_{t-1}) and others
self.sqrt_alphas_cumprod = tf.constant(np.sqrt(alphas_cumprod), dtype=tf_dtype)
self.sqrt_one_minus_alphas_cumprod = tf.constant(np.sqrt(1. - alphas_cumprod), dtype=tf_dtype)
self.log_one_minus_alphas_cumprod = tf.constant(np.log(1. - alphas_cumprod), dtype=tf_dtype)
self.sqrt_recip_alphas_cumprod = tf.constant(np.sqrt(1. / alphas_cumprod), dtype=tf_dtype)
self.sqrt_recipm1_alphas_cumprod = tf.constant(np.sqrt(1. / alphas_cumprod - 1), dtype=tf_dtype)
  
  def q_sample(self, x_start, t, noise=None):
    """
    Diffuse the data (t == 0 means diffused for 1 step)
    """
    if noise is None:
      noise = tf.random_normal(shape=x_start.shape)
    assert noise.shape == x_start.shape
    return (
        self._extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
        self._extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
    )

<Denoising Diffusion Probabilistic Model (DDPM) 논문 코드>

 

 

Noise Schedule $(\beta_{t})$ 설정

def make_beta_schedule(schedule='linear', n_timesteps=1000, start=1e-5, end=1e-2):
    if schedule == 'linear':
        betas = torch.linspace(start, end, n_timesteps)
    elif schedule == "quad":
        betas = torch.linspace(start ** 0.5, end ** 0.5, n_timesteps) ** 2
    elif schedule == "sigmoid":
        betas = torch.linspace(-6, 6, n_timesteps)
        betas = torch.sigmoid(betas) * (end - start) + start
    return betas

 

앞서 살펴본 예제에서는 Beta값이 0.035로 고정되어있었다. 하지만 실제 논문에서는 고정된 값을 사용하지않고 quad, linear, sigmoid 등의 함수를 이용해서 beta값이 점진적으로 증가할 수 있도록 설계한다. (그랬을 때 성능이 더 향상되고 계산도 더 빨라진다고 말한다)

 

betas = make_beta_schedule(schedule='sigmoid', n_timesteps=n_steps, start=1e-5, end=1e-2)
alphas = 1 - betas
alphas_prod = torch.cumprod(alphas, 0)
alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod[:-1]], 0)
alphas_bar_sqrt = torch.sqrt(alphas_prod)
one_minus_alphas_bar_log = torch.log(1 - alphas_prod)
one_minus_alphas_bar_sqrt = torch.sqrt(1 - alphas_prod)

def extract(input, t, x):
    shape = x.shape
    out = torch.gather(input, 0, t.to(input.device))
    reshape = [t.shape[0]] + [1] * (len(shape) - 1)
    return out.reshape(*reshape)
def q_sample(x_0, t, noise=None):
    if noise is None:
        noise = torch.randn_like(x_0)
    alphas_t = extract(alphas_bar_sqrt, t, x_0)
    alphas_1_m_t = extract(one_minus_alphas_bar_sqrt, t, x_0)
    return (alphas_t * x_0 + alphas_1_m_t * noise)
fig, axs = plt.subplots(1, 10, figsize=(28, 3))
for i in range(10):
    q_i = q_sample(dataset, torch.tensor([i * 10]))
    axs[i].scatter(q_i[:, 0], q_i[:, 1], s=10);
    axs[i].set_axis_off(); axs[i].set_title('$q(\mathbf{x}_{'+str(i*10)+'})$')

 

*Code 정리

torch.cumprod (tensor, dim) : 주어진 tensor의 값을 차례대로 곱하는 과정이다. 이를 이용해서 $\bar{\alpha_{t}}=\alpha_{1}\cdot\alpha_{2}\cdots\alpha_{t}$를 의미하게 된다.

ex) x=torch.tensor([1,2,3,4]) y= torch.cumprod(x,dim=0) 

print(y)를 하게 되면 tensor([1,2,6,24])로 출력되게된다.

*dim=0 라는건 행 방향으로 이어붙이겠다는 의미이다.

 

torch.cat(tensor,dim) : 서로 연결되어있지않은 tensor를 연결하여 하나의 tensor(?)로 이어붙여주는 명령어다. (cat=concatenate) 

위 코드 alphas_prod_p = torch.cat([torch.tensor([1]).float(), alphas_prod([:-1]],0)를 한번 살펴보면 먼저 tensor값으로 첫 열에 1을 지정해주고 그 뒤로 alphas_prod값을 맨 마지막 값을 제외하고 전부 이어붙이는 것을 나타내고 있다. 

예를 들어 alphas_prod=([2,3,4,5])이면 alphas_prod_p는 ([1,2,3,4])가 되는 것이다. 즉 $\bar{\alpha_{t-1}}$을 나타낸 것이다.

ex) z= torch.cat([x,y],dim=0) 

print(z)를 하게되면 tensor([1,2,3,4,1,2,6,24])로 출력이 된다. 

import torch

alpha_prod = torch.tensor([2,3,4,5])
alpha_prod_p = torch.cat([torch.tensor([1]).float(),alpha_prod[:-1]],0)
print('alpha_prod_p=', alpha_prod_p)

 

예시 실행결과

 

 

Beta schedule에 따른 차이점 (LINEAR, QUAD, SIGMOID)

 

위의 코드를 통해 noise schedule을 원하는 방식으로 지정해서 사용할 수 있습니다. 해당 코드를 조금 수정해서 subplot을 이용해 비교하기 쉽게 3가지 schedule이 나타날 수 있도록 실행해보았습니다. 사실 step의 수가 세분화되어 나뉘어져있지않기때문에 큰 차이를 느끼기는 어려울 수 있지만 조금 더 세밀하게 나타내게 된다면 차이를 볼 수 있을 것 같습니다.

 

 

이번 포스팅은 여기까지.. Forward process에 대해서 예제 코드를  활용해서 한번 알아보았습니다. 

다음 번에는 Reverse process를 통해 어떻게 학습이 진행되고 어떻게 이미지를 복원해내는지에 대해 한번 알아보도록 하겠습니다. 딥러닝과 관련된 코딩은 처음 해보는 부분이라 앞으로도 조금 자세하게 문법적인 부분도 정리할 겸해서 계속 서술해나갈 계획입니다. 저처럼 처음하시는 분들이 도움을 받을 수 있으면 좋겠습니다. 

그럼 오늘도 좋은 하루 되시길 바랍니다! ㅎㅎ 

 

 

*Reference

https://github.com/acids-ircam

 

 

 

 

 

 

'Main > 딥러닝' 카테고리의 다른 글

강화학습(Reinforcement Learning) - Deep Q Learning(DQN)  (2) 2025.06.16