📌 JAX를 활용한 신경망 모델 구축 - MLP를 이용한 손글씨 인식 (MNIST)

📌 JAX를 활용한 신경망 모델 구축 - MLP를 이용한 손글씨 인식 (MNIST)

---

🚀 JAX로 신경망 모델 만들기

이전 글에서는 JAX의 자동 미분과 JIT 컴파일의 기본 사용법을 배웠습니다.
이번에는 이를 활용하여 간단한 다층 퍼셉트론(MLP) 모델을 구축하여 손글씨 데이터(MNIST)를 분류해보겠습니다.

---

📝 1. 데이터셋 준비

💾 MNIST 데이터 불러오기

JAX에서는 직접 데이터를 불러와야 하므로 tensorflow_datasets를 사용하여 데이터를 로드합니다.

pip install tensorflow-datasets

import tensorflow_datasets as tfds
import jax.numpy as jnp

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

# 데이터 전처리 함수
def preprocess(image, label):
    image = jnp.array(image, dtype=jnp.float32) / 255.0
    label = jnp.array(label, dtype=jnp.int32)
    return image.reshape(-1), label

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

---

💡 2. 신경망 모델 정의

🧠 MLP 모델 구조

• 입력층: 784 (28x28 이미지 펼침)
• 은닉층: 128 (ReLU 활성화 함수)
• 출력층: 10 (Softmax 활성화 함수)

import jax
from jax import random, jit, grad
import jax.numpy as jnp

# 파라미터 초기화 함수
def init_params(layer_sizes, key):
    params = []
    for n_in, n_out in zip(layer_sizes[:-1], layer_sizes[1:]):
        key, subkey = random.split(key)
        weights = random.normal(subkey, (n_in, n_out)) * 0.01
        biases = jnp.zeros(n_out)
        params.append((weights, biases))
    return params

# MLP 모델 함수
def predict(params, x):
    for w, b in params[:-1]:
        x = jnp.dot(x, w) + b
        x = jnp.maximum(x, 0)  # ReLU 활성화 함수
    final_w, final_b = params[-1]
    logits = jnp.dot(x, final_w) + final_b
    return logits - jax.scipy.special.logsumexp(logits, axis=1, keepdims=True)

---

⚙️ 3. 손실 함수와 정확도 계산

📉 손실 함수 (Cross Entropy)

def cross_entropy_loss(params, x, y):
    logits = predict(params, x)
    one_hot = jax.nn.one_hot(y, num_classes=10)
    return -jnp.mean(jnp.sum(one_hot * logits, axis=1))

🧮 정확도 계산 함수

반응형

def accuracy(params, x, y):
    logits = predict(params, x)
    predictions = jnp.argmax(logits, axis=1)
    return jnp.mean(predictions == y)

---

🔧 4. 학습 루프 정의

🔁 훈련 단계

learning_rate = 0.01
epochs = 5
batch_size = 128
key = random.PRNGKey(42)

# 모델 초기화
params = init_params([784, 128, 10], key)

# 기울기 계산 함수
grad_loss = jit(grad(cross_entropy_loss))

# 파라미터 업데이트 함수
@jit
def update(params, x, y, lr):
    grads = grad_loss(params, x, y)
    return [(w - lr * dw, b - lr * db) for (w, b), (dw, db) in zip(params, grads)]

---

🏃 5. 모델 학습

📊 학습 루프

for epoch in range(epochs):
    # 미니배치 학습
    for i in range(0, len(train_data), batch_size):
        batch = train_data[i:i + batch_size]
        x_batch, y_batch = zip(*batch)
        x_batch = jnp.stack(x_batch)
        y_batch = jnp.array(y_batch)

        # 파라미터 업데이트
        params = update(params, x_batch, y_batch, learning_rate)

    # 에포크별 손실 및 정확도 출력
    train_loss = cross_entropy_loss(params, x_batch, y_batch)
    train_acc = accuracy(params, x_batch, y_batch)
    print(f"Epoch {epoch + 1}, Loss: {train_loss:.4f}, Accuracy: {train_acc:.4f}")

---

✅ 6. 학습 결과 평가

학습이 완료된 모델을 사용하여 테스트 데이터를 평가합니다.

🧩 테스트 데이터 정확도 계산

# 테스트 데이터 로드
ds_test = tfds.load('mnist', split='test', as_supervised=True)
test_data = [(preprocess(image, label)) for image, label in tfds.as_numpy(ds_test)]

# 평가
x_test, y_test = zip(*test_data)
x_test = jnp.stack(x_test)
y_test = jnp.array(y_test)

test_acc = accuracy(params, x_test, y_test)
print(f"테스트 정확도: {test_acc:.4f}")

---

🌟 JAX의 장점 활용

1. 자동 미분
   • grad()를 사용하여 기울기 계산을 간단하게 수행합니다.
2. JIT 컴파일
   • jit()를 사용하여 학습 속도를 대폭 향상시켰습니다.
3. 함수 벡터화
   • JAX의 벡터화 기능을 통해 배치 데이터를 효과적으로 처리했습니다.

---

💡 모델 개선 아이디어

• 데이터 증강: 학습 데이터를 늘려 모델 성능 향상
• 드롭아웃 추가: 과적합 방지
• 고급 옵티마이저 사용: Adam이나 RMSprop으로 성능 개선

---

📌 다음 글 예고: JAX로 CNN 구현하기

다음 글에서는 JAX를 사용하여 **합성곱 신경망(CNN)**을 구축하고, 이미지 분류 문제를 더 복잡하게 해결해보겠습니다.
GPU 활용을 극대화하여 성능을 비교 분석해보겠습니다.

---

 

JAX, 신경망, MLP, MNIST, 딥러닝, 머신러닝, Python, 자동 미분, JIT 컴파일, 학습 루프, 데이터 전처리, 모델 학습, 정확도 계산, GPU 활용, 고성능 모델, 함수 벡터화, 모델 평가