[Generative-AI] 디퓨전(확산) 모델 개념잡기

0. 이미지생성 모델: GAN & Diffusion

딥러닝 공부를 하며, 생성형 모델을 배우게 된다.

GAN은 자주 들어서 그런지 익숙하다.

생성자와 판별자 두개가 경쟁하며, 이미지를 생성하는 모델.

하지만, 요즘 대세(?)는 디퓨전 모델들인것 같다.

그럼 디퓨전모델들은 도대체 어떤 아이디어로 구성되는 것일까?

 

 

1. 디퓨전 모델의 기본개념

- 원래 이미지에 노이즈를 가해주면 어떻게 될까?

- 원본 이미지에서, 랜덤한 노이즈를 가해주면 위 이미지처럼, 망가지게 됨

- 첫번째 이미지가 X0이면, 두번째 이미지는 X1 = 원본유지비율1*X0 + 잡음비율1*ε1(잡음), X2 = 원본유지비율2*X1 + 잡음비율2*ε2(잡음)....

- diffusion의 핵심은 ε(잡음)을 예측하는 모델 θ를 학습

- 학습한 모델로 εθ를 예측. 원래 ε와 loss_func구성.

- loss_func을 최소화 시키는  θ업데이트

 

 

2. 디퓨전모델 순서대로 정리(with 코드)

def train_step(self, images):
        # 1. 원본이미지, 각시점별 잡음, 잡음이미지(noisy_images)
        # 배치 이미지들의 평균이 0이고, 단위분산
        images = self.normalizer(images, training=True)
        # 입력 이미지 크기에 맞는 잡음 만들기
        noises = tf.random.normal(shape=(BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, 3))

        # 각 이미지마다, 랜덤한 시점의 노이즈 이미지 가져오기 위함.
        diffusion_times = tf.random.uniform(
            shape=(BATCH_SIZE, 1, 1, 1), minval=0.0, maxval=1.0
        )
        # 각 이미지마다 원본, 잡음 rate
        noise_rates, signal_rates = self.diffusion_schedule(diffusion_times)

        # 이미지마다 랜덤한 시점의 잡음 이미지
        noisy_images = signal_rates * images + noise_rates * noises

        with tf.GradientTape() as tape:
            # 2. 잡음 예측, 잡음이미지 예측
            # 잡음 이미지, 원본/잡음rate주고 > 신경망이 노이즈 예측(pred_noises)
            # Unet에 pred_noises사용해서 > 예상원본이미지 생성
            pred_noises, pred_images = self.denoise(
                noisy_images, noise_rates, signal_rates, training=True
            )

            # 3. 로스함수 선언
            # 진짜 노이즈와 신경망이 예측한 노이즈로 loss함수 > 진짜 노이즈를 잘 찾는 신경망 학습
            noise_loss = self.loss(noises, pred_noises)  # 훈련에 사용

        # 4. 학습 = pred_noises함수 업데이트
        gradients = tape.gradient(noise_loss, self.network.trainable_weights)
        self.optimizer.apply_gradients(
            zip(gradients, self.network.trainable_weights)
        )

        self.noise_loss_tracker.update_state(noise_loss)

- 원본이미지와 랜덤한 잡음 만들고,

- 스탭별 원본비율/잡음비율 셋팅

- noisy_images = signal_rates * images + noise_rates * noises > 잡음이미지(잡음이 적용된) 생성

- Unet에 [noisy_images, noise_rates**2]를 입력 > 해당 잡음이미지에 적용된 잡음 εθ(pred_noises) 도출 > εθ를 기반으로 원본 이미지 예측(pred_images)

- noise_loss = self.loss(noises, pred_noises) > 원래 ε(noises)와 Unet으로 도출한 εθ(pred_noises)으로 loss_func구함

- θ업데이트...계속 학습해서 노이즈된 이미지 입력하면, 적용된 ε를 잘 찾아내는 모델 학습

 

 

3. 학습이후, 디퓨전모델 적용

- θ를 학습해서, εθ를 잘 예측하는 모델 완성

- 랜덤한 노이즈 이미지가 입력 > 적용된 ε를 찾아냄

- ε를 적용하여, 그럴듯한 원본이미지 완성

랜덤한 노이즈 이미지 > εθ잘 찾아내고 > 노이즈제거한 이미지