Diffusion Model을 이해해보자 2편 - chatGPT로 튜토리얼 생성 망한 이유

Diffusion Model을 이해해보자 1편 - 태초마을 (tistory.com) 에서 이어집니다.

 

 

뭐, 해볼라면 해볼수도 있었겠지만, introduction 이후의 논문은 잠시 미루기로 했다. 사실 내 머리만으로는 이런 복잡한 수식들을 이해하기도 힘들고, 언제까지 기억할 수 있을런지도 확신이 들지 않는다. 그래서 파이썬 실습을 통해 구조를 익힌 뒤에 다시 논문으로 돌아오기로 했다.

 

파이썬 실습을 하는 방법은, 다양하게 있겠지만 난 요즘 그렇듯 chatGPT를 활용했다. 질문하기에도 훨씬 편하고, 원하는 수준의 튜토리얼을 설정할 수 있다. 참 공부하기 쉬워진 세상 아닌가? 그래서 더 공부하게 되지 않는 것 같기도 하지만.

 

전체 코드는 아래와 같다.

 

class DiffusionModel(tf.keras.Model):
    def __init__(self, timesteps=1000):
        super(DiffusionModel, self).__init__()
        self.timesteps = timesteps
        self.model = keras.Sequential([
            keras.layers.Dense(128, activation='relu'),
            keras.layers.Dense(128, activation='relu'),
            keras.layers.Dense(1)
        ])

    def sample_noise(self, shape):
        return tf.random.normal(shape)

    def q_sample(self, x_start, t):
        beta_t = 0.02 * tf.sqrt(1.0 - (1.0 - t/self.timesteps))
        return x_start * (1 - beta_t) + self.sample_noise(x_start.shape) * tf.sqrt(beta_t)

    def predict_noise(self, x_t, t):
        return self.model(tf.concat([x_t, tf.fill(x_t.shape[:-1] + (1,), t)], -1))

    def p_sample(self, x_t, t):
        pred_noise = self.predict_noise(x_t, t)
        beta_t = 0.02 * tf.sqrt(1.0 - (1.0 - t/self.timesteps))
        return (x_t - tf.sqrt(beta_t) * pred_noise) / (1 - beta_t)
        
def train_step(model, x_start, optimizer):
    with tf.GradientTape() as tape:
        t = tf.random.uniform([], maxval=model.timesteps, dtype=tf.int32)
        x_t = model.q_sample(x_start, t)
        predicted_noise = model.predict_noise(x_t, t)
        loss = tf.reduce_mean((predicted_noise - model.sample_noise(x_start.shape))**2)
    
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return loss

# Example usage:
model = DiffusionModel()
optimizer = tf.keras.optimizers.Adam()
for epoch in range(10):
    for x_start in dataset:  # replace 'dataset' with actual data
        loss = train_step(model, x_start, optimizer)
    print(f"Epoch {epoch}: Loss {loss.numpy()}")

def generate_sample(model, shape):
    x_t = model.sample_noise(shape)
    for t in reversed(range(model.timesteps)):
        x_t = model.p_sample(x_t, t)
    return x_t

# Generate an image
generated_image = generate_sample(model, (1, 28, 28))
plt.imshow(generated_image.numpy().squeeze(), cmap='gray')
plt.show()

 

막상 파이썬을 켜고 보니, 일단 텐서플로 구조에 더 익숙해져야겠구나 싶다. gradient tape는 예전에 본 것도 같지만, 아무리 해도 익숙해지지가 않는다. gradient tape의 역할은 일종의 옵저버라고 생각된다. 학습 과정에서 back propagation 같은 편미분과 학습 파라미터 전달 과정을 쉽게 도와주는 녀석인데, 코드 중간에 덩그러니 "gradient tape를 활용해서 코드를 이렇게 짰습니다" 라는 말투가 전혀 와닿지가 않는다. 빼서는 안되는 필수적인 녀석인데, 왜 쓸 때마다 저렇게 소개해줘야 하지? 싶은 것이다.

 

뭐 그것만 빼고는, 스무스 하다. 1편에서 익혔듯, p 구조와 q 구조가 들어가 있고, 둘을 연결해주는 predict_noise가 중간에 있다. 그 부분을 학습 시켜주는 것이다. CIFAR10 데이터를 통해 전체를 구현하면 아래와 같다. 

---

위까지가 chatGPT예제를 직접 돌려보기 전에 쓴 글이고, 돌려본 이후에는 chatGPT만으로 튜토리얼 학습은 절대 불가능하다는 결론이 나왔다. 끝없이 쏟아지는 에러에 두손 두발 다들었다. 네트워크 구조가 복잡해지다보니, 레이어 간 하이퍼파라미터 통일에도 애를 먹는 듯 하다. 구조를 아무래도 직접 뜯어보고, 설정을 변경해줘야겠다.