📌 JAX로 GAN(생성적 적대 신경망) 구현 - 이미지 생성 프로젝트

📌 JAX로 GAN(생성적 적대 신경망) 구현 - 이미지 생성 프로젝트

---

🚀 GAN(생성적 적대 신경망)이란?

GAN(Generative Adversarial Network)은 **생성자(Generator)**와 **판별자(Discriminator)**라는 두 개의 신경망을 사용하여 데이터를 학습하고 생성하는 모델입니다.

• 생성자(G): 랜덤한 잡음 벡터로부터 실제 같은 데이터를 생성
• 판별자(D): 데이터가 실제(real)인지 생성된(fake) 것인지 판별
• 목표: 생성자는 판별자를 속일 정도로 실제 같은 데이터를 생성

---

💡 1. GAN의 학습 과정

📐 GAN 알고리즘 구조

1. 생성자 학습: 판별자를 속이도록 가짜 데이터를 생성
2. 판별자 학습: 실제 데이터와 생성 데이터를 구분
3. 적대적 훈련: 두 모델이 서로 경쟁하며 발전

🔄 손실 함수 정의

• 생성자 손실:

LG=−E[log⁡(D(G(z)))]L_G = -\mathbb{E}[\log(D(G(z)))]

• 판별자 손실:

LD=−E[log⁡(D(x))]−E[log⁡(1−D(G(z)))]L_D = -\mathbb{E}[\log(D(x))] - \mathbb{E}[\log(1 - D(G(z)))]

---

🔧 2. 라이브러리 설치

pip install jax jaxlib flax optax tensorflow-datasets

---

💾 3. 데이터 준비 - MNIST 손글씨 데이터

import tensorflow_datasets as tfds

# MNIST 데이터 불러오기
ds = tfds.load("mnist", split="train", as_supervised=True)

# 데이터 전처리 함수
def preprocess(image, label):
    image = (image / 255.0) * 2 - 1  # [-1, 1]로 정규화
    return jnp.array(image, dtype=jnp.float32).reshape(-1), label

# 데이터 변환
train_data = [preprocess(image, label) for image, label in tfds.as_numpy(ds)]
print(f"훈련 데이터 수: {len(train_data)}")

---

🧠 4. GAN 모델 구현

반응형

📝 생성자 모델 (Generator)

from flax import linen as nn
import jax.numpy as jnp

class Generator(nn.Module):
    latent_dim: int

    @nn.compact
    def __call__(self, z):
        x = nn.Dense(128)(z)
        x = nn.relu(x)
        x = nn.Dense(256)(x)
        x = nn.relu(x)
        x = nn.Dense(512)(x)
        x = nn.relu(x)
        x = nn.Dense(28 * 28)(x)
        x = nn.tanh(x)  # [-1, 1] 범위로
        return x.reshape((-1, 28, 28, 1))

---

📝 판별자 모델 (Discriminator)

class Discriminator(nn.Module):

    @nn.compact
    def __call__(self, x):
        x = x.reshape(-1, 28 * 28)  # 평탄화
        x = nn.Dense(512)(x)
        x = nn.leaky_relu(x, 0.2)
        x = nn.Dense(256)(x)
        x = nn.leaky_relu(x, 0.2)
        x = nn.Dense(1)(x)
        return x

---

📉 5. 손실 함수와 최적화 설정

📝 손실 함수

import optax

# 생성자 손실
def generator_loss(logits_fake):
    return -jnp.mean(jax.nn.sigmoid(logits_fake))

# 판별자 손실
def discriminator_loss(logits_real, logits_fake):
    loss_real = -jnp.mean(jax.nn.log_sigmoid(logits_real))
    loss_fake = -jnp.mean(jax.nn.log_sigmoid(-logits_fake))
    return loss_real + loss_fake

# 옵티마이저
optimizer = optax.adam(1e-4)

---

🏗️ 6. 모델 초기화

from flax.training import train_state
import jax
from jax import random

key = random.PRNGKey(42)

# 모델 초기화
gen_model = Generator(latent_dim=100)
dis_model = Discriminator()

gen_params = gen_model.init(key, random.normal(key, (1, 100)))
dis_params = dis_model.init(key, random.normal(key, (1, 28, 28, 1)))

# 학습 상태 초기화
gen_state = train_state.TrainState.create(apply_fn=gen_model.apply, params=gen_params, tx=optimizer)
dis_state = train_state.TrainState.create(apply_fn=dis_model.apply, params=dis_params, tx=optimizer)

---

🔁 7. 학습 루프 정의

📝 업데이트 함수

@jax.jit
def train_step(gen_state, dis_state, real_images, z):
    def gen_loss_fn(params):
        fake_images = gen_model.apply(params, z)
        logits_fake = dis_model.apply(dis_state.params, fake_images)
        return generator_loss(logits_fake)

    def dis_loss_fn(params):
        logits_real = dis_model.apply(params, real_images)
        fake_images = gen_model.apply(gen_state.params, z)
        logits_fake = dis_model.apply(params, fake_images)
        return discriminator_loss(logits_real, logits_fake)

    # 기울기 계산
    gen_loss, gen_grads = jax.value_and_grad(gen_loss_fn)(gen_state.params)
    dis_loss, dis_grads = jax.value_and_grad(dis_loss_fn)(dis_state.params)

    # 파라미터 업데이트
    gen_state = gen_state.apply_gradients(grads=gen_grads)
    dis_state = dis_state.apply_gradients(grads=dis_grads)
    
    return gen_state, dis_state, gen_loss, dis_loss

---

🏃 8. 모델 학습

epochs = 50
batch_size = 64

for epoch in range(epochs):
    for i in range(0, len(train_data), batch_size):
        batch = train_data[i:i + batch_size]
        real_images, _ = zip(*batch)
        real_images = jnp.stack(real_images)

        # 랜덤 잡음 생성
        z = random.normal(key, (batch_size, 100))

        # 학습 스텝
        gen_state, dis_state, g_loss, d_loss = train_step(gen_state, dis_state, real_images, z)

    print(f"Epoch {epoch+1}, Gen Loss: {g_loss:.4f}, Dis Loss: {d_loss:.4f}")

---

✅ 9. 이미지 생성 결과 확인

import matplotlib.pyplot as plt

def generate_images(gen_state, num_images=16):
    z = random.normal(key, (num_images, 100))
    generated = gen_model.apply(gen_state.params, z)
    generated = (generated + 1) / 2  # [-1, 1] -> [0, 1]

    plt.figure(figsize=(4, 4))
    for i in range(num_images):
        plt.subplot(4, 4, i+1)
        plt.imshow(generated[i, :, :, 0], cmap='gray')
        plt.axis('off')
    plt.show()

generate_images(gen_state)

---

📌 다음 글 예고: JAX로 VAE(변이형 오토인코더) 구현

다음 글에서는 JAX를 사용하여 VAE 모델을 구축하고,
잠재 공간에서의 의미 있는 데이터 생성 방식을 탐구하겠습니다.

---

 

JAX, GAN, 이미지 생성, 딥러닝, Python, 고속 연산, 생성적 적대 신경망, 고성능 컴퓨팅, 학습 루프, 모델 학습, MNIST, 이미지 생성 모델, 랜덤 잡음, Generator, Discriminator, 인공지능 모델링