ReLU로 MNIST 손글씨 분류하기 (TensorFlow)

Sigmoid 함수의 Vanishing Gradient 문제를 ReLU로 해결하고,
실제 MNIST 손글씨 데이터셋으로 숫자를 분류하는 모델을 구현해보겠습니다.

---

전체 코드

import tensorflow as tf
import numpy as np
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.datasets import mnist

def load_mnist():
    (train_data, train_labels), (test_data, test_labels) = mnist.load_data()
    train_data = np.reshape(train_data, [-1, 28, 28, 1])
    test_data = np.expand_dims(test_data, axis=-1)
    train_data, test_data = normalize(train_data, test_data)
    train_labels = to_categorical(train_labels, 10)
    test_labels = to_categorical(test_labels, 10)
    return train_data, train_labels, test_data, test_labels

def normalize(train_data, test_data):
    train_data = train_data.astype(np.float32) / 255.0
    test_data = test_data.astype(np.float32) / 255.0
    return train_data, test_data

def flatten():
    return tf.keras.layers.Flatten()

def dense(channel, weight_init):
    return tf.keras.layers.Dense(units=channel, use_bias=True, kernel_initializer=weight_init)

def relu():
    return tf.keras.layers.Activation(tf.keras.activations.relu)

class create_model(tf.keras.Model):
    def __init__(self, label_dim):
        super(create_model, self).__init__()
        weight_init = tf.keras.initializers.RandomNormal()
        self.model = tf.keras.Sequential()
        self.model.add(flatten())
        for i in range(2):
            self.model.add(dense(256, weight_init))
            self.model.add(relu())
        self.model.add(dense(label_dim, weight_init))

    def call(self, x, training=None, mask=None):
        x = self.model(x)
        return x

def loss_fn(model, images, labels):
    logits = model(images, training=True)
    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=labels))
    return loss

def accuracy_fn(model, images, labels):
    logits = model(images, training=True)
    predicted = tf.equal(tf.argmax(logits, -1), tf.argmax(labels, -1))
    accuracy = tf.reduce_mean(tf.cast(predicted, dtype=tf.float32))
    return accuracy

def grad(model, images, labels):
    with tf.GradientTape() as tape:
        loss = loss_fn(model, images, labels)
    return tape.gradient(loss, model.trainable_variables)

train_x, train_y, test_x, test_y = load_mnist()

learning_rate = 0.001
batch_size = 128
training_epochs = 1
training_iterations = len(train_x) // batch_size
label_dim = 10

train_dataset = tf.data.Dataset.from_tensor_slices((train_x, train_y)) \
    .shuffle(buffer_size=100000).batch(batch_size).prefetch(buffer_size=batch_size)
test_dataset = tf.data.Dataset.from_tensor_slices((test_x, test_y)).batch(batch_size)

network = create_model(label_dim)
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)
checkpoint = tf.train.Checkpoint(dnn=network)

for epoch in range(training_epochs):
    for step, (train_input, train_label) in enumerate(train_dataset):
        grads = grad(network, train_input, train_label)
        optimizer.apply_gradients(zip(grads, network.trainable_variables))
        train_loss = loss_fn(network, train_input, train_label)
        train_accuracy = accuracy_fn(network, train_input, train_label)
        for test_input, test_label in test_dataset.take(1):
            test_accuracy = accuracy_fn(network, test_input, test_label)
        print(f"Epoch: {epoch}, Step: {step}, Loss: {train_loss:.4f}, Train Acc: {train_accuracy:.4f}, Test Acc: {test_accuracy:.4f}")

---

1. MNIST 데이터셋

MNIST는 0~9까지의 손글씨 숫자 이미지 데이터셋입니다.
60,000개의 학습 데이터와 10,000개의 테스트 데이터로 구성되어 있습니다.

(train_data, train_labels), (test_data, test_labels) = mnist.load_data()

각 이미지는 28×28 픽셀 크기입니다.

---

2. 데이터 전처리

Shape 변환

train_data = np.reshape(train_data, [-1, 28, 28, 1])
test_data = np.expand_dims(test_data, axis=-1)

TensorFlow가 이미지를 받는 shape는 [batch_size, height, width, channel] 입니다.
MNIST는 흑백(그레이스케일) 이미지이기 때문에 채널이 1입니다. (컬러는 RGB로 3)
원래 shape가 [N, 28, 28] 이기 때문에 채널 차원을 추가해 [N, 28, 28, 1] 로 만들어줍니다.

두 가지 방법이 있습니다:

• np.reshape : 전체 구조를 강제로 재배열 (데이터 순서가 꼬일 수 있음)
• np.expand_dims : 차원 하나만 추가 (더 안전, 일반적으로 더 많이 사용)

정규화 (Normalization)

train_data = train_data.astype(np.float32) / 255.0
test_data = test_data.astype(np.float32) / 255.0

이미지 픽셀값은 0~255 사이의 값인데, 이것을 0~1 사이로 정규화해야 합니다.
데이터의 범위가 너무 크면 학습이 잘 되지 않기 때문에 255로 나눠서 스케일을 맞춰줍니다.

One-Hot Encoding

train_labels = to_categorical(train_labels, 10)  # [N,] → [N, 10]
test_labels = to_categorical(test_labels, 10)

정답 레이블(0~9)을 One-Hot 형태로 변환합니다.
Loss 함수를 계산할 때 정수값보다 One-Hot 형태가 훨씬 편리하기 때문입니다.

7 → [0, 0, 0, 0, 0, 0, 0, 1, 0, 0]
3 → [0, 0, 0, 1, 0, 0, 0, 0, 0, 0]

---

3. 네트워크 구성 함수

def flatten():
    return tf.keras.layers.Flatten()

def dense(channel, weight_init):
    return tf.keras.layers.Dense(units=channel, use_bias=True, kernel_initializer=weight_init)

def relu():
    return tf.keras.layers.Activation(tf.keras.activations.relu)

네트워크를 짤 때 어떤 함수를 사용할지 먼저 정의합니다.

• Flatten : shape를 펼쳐주는 역할 [N, 28, 28, 1] → [N, 784]
• Dense : Fully Connected Layer (모든 뉴런이 연결된 레이어)
• ReLU : Activation Function f(x) = max(0, x)

Sigmoid 대신 ReLU를 사용하는 이유는 Sigmoid의 Gradient가 항상 0~0.25로 작아서 네트워크가 깊어질수록 Vanishing Gradient가 발생하기 때문입니다.
ReLU는 양수 구간에서 Gradient = 1 이기 때문에 깊은 네트워크에서도 Gradient가 잘 전달됩니다.

---

4. 모델 클래스

class create_model(tf.keras.Model):
    def __init__(self, label_dim):
        super(create_model, self).__init__()
        weight_init = tf.keras.initializers.RandomNormal()
        self.model = tf.keras.Sequential()
        self.model.add(flatten())
        for i in range(2):
            self.model.add(dense(256, weight_init))
            self.model.add(relu())
        self.model.add(dense(label_dim, weight_init))

    def call(self, x, training=None, mask=None):
        x = self.model(x)
        return x

tf.keras.Model 을 상속받아 모델 클래스를 만듭니다.

• weight_init = tf.keras.initializers.RandomNormal() : 평균 0, 분산 1인 가우시안 분포로 W를 랜덤하게 초기화합니다
• tf.keras.Sequential() : 레이어를 순서대로 쌓아나가는 리스트 자료구조 같은 것입니다

모델의 레이어 구조:

[N, 28, 28, 1]
    ↓ Flatten
[N, 784]
    ↓ Dense(256) + ReLU
[N, 256]
    ↓ Dense(256) + ReLU
[N, 256]
    ↓ Dense(10)
[N, 10]   ← 숫자 0~9 각각의 점수(logit)

call 함수는 모델을 호출했을 때 실행되는 함수입니다.

---

5. Gradient와 Backpropagation

def grad(model, images, labels):
    with tf.GradientTape() as tape:
        loss = loss_fn(model, images, labels)
    return tape.gradient(loss, model.trainable_variables)

여기서 자주 헷갈리는 두 개념을 정리하면:

Backpropagation 은 Gradient를 계산하는 방법입니다.
Gradient Descent 는 계산된 Gradient를 사용해서 W를 업데이트하는 방법입니다.

1. Forward
   Input → Network → Output → Loss 계산

2. Backpropagation  ← tape.gradient()
   Loss를 미분 → 각 W의 Gradient 계산 (Chain Rule)

3. Gradient Descent  ← optimizer.apply_gradients()
   W := W - α * Gradient

tape.gradient 가 Backpropagation을 수행하고,
optimizer.apply_gradients 가 Gradient Descent로 W를 업데이트합니다.

이전 코드에서는 W.assign_sub(learning_rate * W_grad) 로 직접 업데이트했는데,
optimizer.apply_gradients 가 이 과정을 자동으로 처리해주는 것입니다.

---

6. 데이터셋 파이프라인

train_dataset = tf.data.Dataset.from_tensor_slices((train_x, train_y)) \
    .shuffle(buffer_size=100000) \
    .batch(batch_size) \
    .prefetch(buffer_size=batch_size)

60,000개의 데이터를 한 번에 메모리에 올리면 부담이 크기 때문에 batch_size(128개)씩 쪼개서 네트워크에 넣습니다.

• shuffle(buffer_size) : 데이터를 섞습니다. buffer_size는 데이터셋 크기보다 크면 됩니다
• batch(128) : 128개씩 묶어서 네트워크에 넣습니다
• prefetch(buffer_size) : 현재 배치를 학습하는 동안 다음 배치를 미리 메모리에 올려둬서 속도를 높입니다

training_iterations = len(train_x) // batch_size
# 60,000 // 128 = 468
# 전체 데이터를 128개씩 나누면 468번의 스텝이 필요

---

7. Checkpoint

checkpoint = tf.train.Checkpoint(dnn=network)

Checkpoint는 두 가지 역할을 합니다:

• 학습 중간에 끊겼을 때 저장된 Weight를 불러와 이어서 학습 가능
• 학습이 끝난 후 테스트 이미지에 대한 정확도 확인 시 저장된 모델 불러오기 가능

---

8. Adam Optimizer

optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

이전 코드에서는 SGD(확률적 경사 하강법)를 사용했는데, 이번엔 Adam 을 사용합니다.

Adam은 SGD보다 대부분의 경우 더 빠르고 안정적으로 수렴합니다.
학습률을 자동으로 조절하는 기능이 있어 실무에서 가장 많이 사용되는 옵티마이저입니다.
일반적으로 learning_rate = 0.001 또는 0.0003 을 사용합니다.

---

학습 결과

Epoch: 0, Step: 0,   Loss: 2.1713, Train Acc: 0.4453, Test Acc: 0.2891
Epoch: 0, Step: 10,  Loss: 1.2406, Train Acc: 0.7969, Test Acc: 0.7500
Epoch: 0, Step: 50,  Loss: 0.3108, Train Acc: 0.9297, Test Acc: 0.8672
Epoch: 0, Step: 100, Loss: 0.3457, Train Acc: 0.8906, Test Acc: 0.9375
Epoch: 0, Step: 200, Loss: 0.1797, Train Acc: 0.9531, Test Acc: 0.9766
Epoch: 0, Step: 300, Loss: 0.1506, Train Acc: 0.9375, Test Acc: 0.9844
Epoch: 0, Step: 400, Loss: 0.1634, Train Acc: 0.9609, Test Acc: 0.9844
Epoch: 0, Step: 468, Loss: 0.0909, Train Acc: 0.9688, Test Acc: 0.9844

단 1 Epoch(전체 데이터를 한 번 학습) 만에 약 98%의 정확도를 달성했습니다.

---

전체 흐름 정리

MNIST 이미지 로드 (28×28, 흑백)
    ↓
Shape 변환 : [N, 28, 28] → [N, 28, 28, 1]
    ↓
정규화 : 0~255 → 0~1
    ↓
One-Hot Encoding : 7 → [0,0,0,0,0,0,0,1,0,0]
    ↓
Flatten : [N, 28, 28, 1] → [N, 784]
    ↓
Dense(256) + ReLU → Dense(256) + ReLU → Dense(10)
    ↓
Softmax Cross-Entropy Loss 계산
    ↓
GradientTape → Backpropagation → Gradient 계산
    ↓
Adam optimizer → Gradient Descent → W 업데이트
    ↓
1 Epoch 학습 → Test Accuracy ~98% ✅

---

정리

개념 설명 코드

Flatten	이미지를 1차원으로 펼침	tf.keras.layers.Flatten()
Dense	Fully Connected Layer	tf.keras.layers.Dense(256)
ReLU	Vanishing Gradient 해결	tf.keras.activations.relu
normalize	0~255 → 0~1 정규화	data / 255.0
prefetch	다음 배치 미리 메모리에 올림	.prefetch(buffer_size)
Checkpoint	학습 중단 시 모델 저장/복원	tf.train.Checkpoint(dnn=network)
Adam	학습률 자동 조절 옵티마이저	tf.keras.optimizers.Adam(0.001)
Backpropagation	Gradient를 계산하는 방법	tape.gradient(loss, variables)
Gradient Descent	Gradient로 W를 업데이트하는 방법	optimizer.apply_gradients(...)

ReLU와 Adam 옵티마이저를 사용해 단 1 Epoch만에 98%의 정확도를 달성했습니다. 🚀