2.5 소프트맥스 회귀로 MNIST 데이터 분류하기

유명곤·2026년 9월 13일

딥 러닝 파이토치 교과서

목록 보기

6/18

손글씨 한 장을 분류하는 일은, 784개의 픽셀 값에서 열 개의 숫자 중 하나를 고르는 일이다. 모델은 픽셀의 밝기를 입력받아 숫자마다 점수를 계산하고, 정답과 비교한 손실을 바탕으로 가중치를 조정한다. 이 과정을 학습에 사용하지 않은 이미지에도 적용하면 분류 성능을 평가할 수 있다.

2.4 소프트맥스 회귀 이해하기에서 로짓·확률·교차 엔트로피의 관계를 살펴보았다. 이번 글에서는 그 관계를 MNIST 데이터에 적용해 데이터 준비 → 모델 계산 → 학습 → 평가의 흐름을 연결한다.

손글씨를 입력과 정답으로 표현하기

MNIST는 0부터 9까지의 손글씨 숫자를 담은 이미지 데이터셋이다. 학습용 이미지 60,000장과 테스트용 이미지 10,000장으로 나뉘며, 각 이미지에는 실제 숫자를 나타내는 레이블(label), 즉 정답이 붙어 있다. 한 이미지는 하나의 숫자에 해당하므로 클래스가 열 개인 분류 문제이다.

이미지는 가로 28, 세로 28픽셀의 흑백 이미지이다. 한 픽셀의 밝기는 0부터 255 사이의 정수로 저장된다. 입력은 이 밝기 값들이고, 정답은 0부터 9 사이의 정수 하나이다. 정답 숫자를 연속적인 크기로 회귀하는 것이 아니라 해당 숫자의 클래스를 고른다.

torchvision은 이미지 데이터셋과 변환 도구 등을 제공하는 PyTorch의 패키지이다. torchvision.datasets.MNIST로 데이터셋을 준비하고, transform에 이미지를 모델 입력으로 바꾸는 규칙을 지정할 수 있다.

밝기 범위와 텐서의 모양 맞추기

이 글에서는 transforms.ToTensor()를 사용한다. MNIST의 8비트 흑백 이미지에 적용하면 픽셀 값을 255로 나누어 0부터 1 사이의 실수로 바꾸고, 채널 축을 앞에 둔 텐서를 만든다. 따라서 0은 0.0, 255는 1.0에 대응한다.

변환된 이미지 하나의 shape는 (1, 28, 28)이다. 맨 앞의 1은 흑백 이미지의 채널 수이다. DataLoader가 이미지 $B$ 장을 묶으면 입력은 (B, 1, 28, 28), 정답은 (B,)가 된다. 여기서 $B$ 는 현재 배치의 실제 표본 수이다. 정답 텐서의 dtype은 정수 클래스 인덱스를 담는 torch.long이다.

학습과 테스트의 입력에는 같은 밝기 변환을 적용해야 한다. 학습은 0~1 범위로 진행하고 테스트에는 0~255 값을 넣으면, 같은 픽셀 밝기가 서로 다른 크기의 입력으로 전달된다. 단일 이미지를 예측할 때도 같은 규칙을 사용한다.

MNIST의 transform은 dataset[index]로 표본을 가져오는 과정에서 적용된다. 따라서 테스트용 DataLoader나 test_dataset[index]를 통해 읽으면 지정한 변환을 거친 이미지를 얻는다. 원시 저장 텐서인 .data에 직접 접근하면 이 변환을 거치지 않으며, .float()만 호출해도 255로 나누는 연산은 수행되지 않는다.

784개의 특성에서 열 개의 로짓 만들기

소프트맥스 회귀에 넣기 위해 이미지의 채널·높이·너비 축을 하나의 특성 축으로 펼친다. 펼치기(flatten)는 픽셀 값과 순서를 유지하면서 텐서의 모양을 바꾸는 연산이다. 이미지마다 $1\times28\times28=784$ 개의 값이 있으므로 (B, 784)가 된다.

PyTorch의 nn.Flatten(start_dim=1)은 0번 축인 배치 축을 유지하고 나머지 축을 펼친다. 이어서 nn.Linear(784, 10)을 적용하면 이미지마다 숫자 0~9에 대응하는 로짓 열 개를 얻는다.

nn.Sequential은 모듈을 나열한 순서대로 적용하므로, 펼치기와 선형 계층을 묶어 이미지 배치에서 로짓까지 계산하는 모델을 만들 수 있다.

단계	shape	의미
변환된 이미지 배치	`(B, 1, 28, 28)`	이미지 $B$ 장, 흑백 채널 하나
펼친 입력	`(B, 784)`	이미지마다 픽셀 특성 784개
선형 계층의 출력	`(B, 10)`	이미지마다 숫자별 로짓 열 개
정답	`(B,)`	이미지마다 정수 정답 하나
예측 클래스	`(B,)`	각 행에서 가장 큰 로짓의 위치

배치를 행으로 쌓은 입력을 $\mathbf{X}\in\mathbb{R}^{B\times784}$ 라 하면, 선형 계층은 다음을 계산한다.

\mathbf{Z}=\mathbf{X}\mathbf{W}^{\mathsf T}+\mathbf{b}, \qquad \mathbf{W}\in\mathbb{R}^{10\times784}, \quad \mathbf{b}\in\mathbb{R}^{10}

가중치 행렬의 각 행은 한 숫자의 점수를 계산하며, 편향은 배치의 모든 이미지에 공유된다. 학습하는 매개변수는 가중치 7,840개와 편향 10개를 합친 7,850개이다. Flatten에는 학습할 매개변수가 없다.

이 모델은 각 픽셀 위치의 밝기에 가중치를 곱해 점수를 더한다. 펼치기 자체가 픽셀 값을 없애는 것은 아니지만, 이 단일 선형 계층에는 주변 픽셀의 모양을 단계적으로 추출하는 연산이 없다. 이 글의 목적은 이 간단한 모델로 분류의 전체 과정을 완성하는 것이다.

손실에는 로짓과 정수 정답을 전달한다

nn.CrossEntropyLoss()에는 (B, 10)의 로짓과 (B,)의 정수 정답을 전달한다. 정답이 7인 이미지에서는 출력의 인덱스 7이 숫자 7에 대응한다. 이 방식에서는 정답을 원-핫 벡터로 바꿀 필요가 없다.

모델 끝에 Softmax를 추가하지 않는다. CrossEntropyLoss는 로짓을 받아 로그 소프트맥스와 정답 클래스의 음의 로그 확률에 해당하는 손실을 안정적으로 계산한다. 소프트맥스 회귀라는 이름에 맞추려고 확률을 먼저 계산해 넣으면 손실 함수가 기대하는 입력과 달라진다.

기본 설정에서는 배치에 속한 표본들의 손실을 평균낸 스칼라가 반환된다. 이 값으로 역전파하고 가중치를 갱신한다. 예측 숫자만 필요할 때는 로짓에 argmax(dim=1)을 적용한다. 클래스 축인 1번 축에서 가장 큰 점수의 위치를 고르는 것이다. 확률 자체를 확인하려는 경우에만 별도로 softmax(dim=1)을 계산하면 된다.

학습할 때와 평가할 때의 차이

학습에서는 미니 배치마다 로짓과 손실을 계산하고, 역전파(backpropagation)로 매개변수의 기울기를 구한 뒤 값을 갱신한다. optimizer.zero_grad()는 이전 기울기를 초기화하고, loss.backward()는 현재 손실의 기울기를 계산하며, optimizer.step()이 실제 매개변수 갱신을 수행한다. 배치를 순회하는 원리는 2.3 미니 배치와 데이터 로더와 같다.

평가에서는 학습을 마친 매개변수를 고정하고 테스트 이미지의 예측을 정답과 비교한다. 테스트 데이터는 이번 학습의 기울기 계산이나 갱신에 사용하지 않는다. 정확도(accuracy)는 평가한 이미지 중 정답을 맞힌 비율이다.

\mathrm{accuracy} =\frac{\text{number of correct predictions}}{\text{number of evaluated samples}}

교차 엔트로피는 정답에 부여한 확률을 평가하고, 정확도는 가장 높은 점수를 받은 클래스가 정답인지만 확인한다. 예측 클래스가 그대로여도 정답 확률은 바뀔 수 있으므로 두 지표는 서로 다른 정보를 담는다.

평가 모드와 기울기 기록은 별개이다

model.eval()은 모델을 평가 모드로 전환한다. torch.no_grad()는 그 안에서 수행하는 계산을 역전파용으로 기록하지 않도록 한다. 따라서 둘은 서로를 대신하지 않는다.

이번 모델의 Flatten과 Linear는 학습·평가 모드에 따라 계산 방식이 바뀌지 않는다. 그래도 학습에서는 model.train(), 평가에서는 model.eval()을 사용해 역할을 드러낸다. 평가 반복문에는 torch.no_grad()를 함께 적용하며, backward()와 step()은 호출하지 않는다.

지표를 집계할 때도 배치와 전체 데이터를 구분해야 한다. 배치 평균 손실에 실제 배치 크기를 곱하면 그 배치의 손실 합이 된다. 이를 누적한 뒤 전체 표본 수로 나누면 마지막 배치 크기가 달라도 표본별 평균을 구할 수 있다. 정확도 역시 배치별 정확도의 단순 평균보다 전체 정답 수를 전체 평가 수로 나누어 계산한다.

MNIST 분류기를 구현하고 결과 확인하기

다음 Python 코드 블록들은 위에서부터 순서대로 이어 실행하는 하나의 예시이다. PyTorch, torchvision, Matplotlib이 설치된 환경을 사용하며, 처음 실행하면 현재 작업 폴더의 data에 MNIST를 내려받는다.

책과 같은 학습률 0.1, 배치 크기 100, 15에포크를 사용한다. 이 값은 실습의 출발 설정이며 최적값을 뜻하지 않는다. 실행 장치는 CPU로 고정한다. 난수 시드는 초기화와 데이터 섞기의 변동을 줄이지만, 버전이나 장치가 달라도 결과가 완전히 같음을 보장하지는 않는다.

데이터와 모델 준비하기

import torch
import torchvision
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

torch.manual_seed(777)
device = torch.device("cpu")
batch_size = 100
epochs = 15
learning_rate = 0.1

transform = transforms.ToTensor()
train_dataset = datasets.MNIST(
    root="data", train=True, transform=transform, download=True
)
test_dataset = datasets.MNIST(
    root="data", train=False, transform=transform, download=True
)
train_loader = DataLoader(
    train_dataset, batch_size=batch_size, shuffle=True
)
test_loader = DataLoader(
    test_dataset, batch_size=batch_size, shuffle=False
)

print(f"PyTorch {torch.__version__}, torchvision {torchvision.__version__}")
print(f"device={device}, train={len(train_dataset)}, test={len(test_dataset)}")

model = nn.Sequential(
    nn.Flatten(start_dim=1),  # (B, 1, 28, 28) -> (B, 784)
    nn.Linear(28 * 28, 10),  # (B, 784) -> (B, 10)
).to(device)
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)

두 데이터셋에 같은 transform을 지정했다. 학습 로더는 순회할 때 표본 순서를 섞고, 테스트 로더는 순서를 유지한다. drop_last를 지정하지 않아 기본값인 False가 적용되므로 남는 표본도 사용한다. 이번에는 60,000장이 배치 크기 100으로 나누어떨어져 한 에포크에 600번 갱신한다.

nn.Sequential 안에 펼치기를 넣었으므로 학습·평가·단일 예측에서 같은 입력 처리 과정을 사용한다. 모델과 옵티마이저는 반복문 밖에서 한 번 만든다.

미니 배치마다 손실을 계산하고 갱신하기

for epoch in range(epochs):
    model.train()
    loss_sum = 0.0
    sample_count = 0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()
        logits = model(images)
        loss = loss_fn(logits, labels)
        loss.backward()
        optimizer.step()

        batch_count = labels.size(0)
        loss_sum += loss.item() * batch_count
        sample_count += batch_count

    print(f"epoch {epoch + 1:02d}: train_loss={loss_sum / sample_count:.4f}")

모델의 매개변수와 입력·정답을 같은 장치에 둔다. 현재 배치에서 images는 (100, 1, 28, 28), logits는 (100, 10), labels는 (100,)이다. 600번의 갱신을 15에포크 동안 반복하므로 총 갱신 횟수는 9,000번이다.

loss.item()은 스칼라 텐서의 값을 Python 숫자로 꺼낸다. 지표 누적에 역전파 그래프를 연결하지 않으면서 손실 합을 기록할 수 있다.

출력한 train_loss는 각 배치를 처리할 당시, 갱신 직전에 계산한 손실들의 평균이다. 에포크가 끝난 최종 모델로 학습 데이터 전체를 다시 평가한 손실은 아니다. 한 에포크 안에서도 매개변수가 계속 바뀐다는 점을 함께 기억해야 한다.

테스트 이미지 전체로 평가하기

model.eval()
test_loss_sum = 0.0
correct = 0
sample_count = 0

with torch.no_grad():
    for images, labels in test_loader:
        images = images.to(device)
        labels = labels.to(device)
        logits = model(images)

        test_loss_sum += loss_fn(logits, labels).item() * labels.size(0)
        predictions = logits.argmax(dim=1)
        correct += (predictions == labels).sum().item()
        sample_count += labels.size(0)

print(f"test_loss={test_loss_sum / sample_count:.4f}")
print(f"test_accuracy={correct / sample_count:.2%} ({correct}/{sample_count})")

predictions == labels는 이미지별 정답 여부를 나타내는 불리언 텐서이다. 이를 합산해 맞힌 이미지 수를 센다. 테스트 데이터 10,000장을 100장씩 읽으므로 평가 반복문은 100번 실행되며, 모든 배치에 같은 최종 모델을 사용한다.

학습 데이터에서 손실이 줄었다는 사실만으로 처음 보는 이미지의 성능을 판단할 수는 없다. 이 테스트 평가는 학습에 쓰지 않은 정해진 분할에서의 성능을 보여준다. 설정을 여러 번 비교해 선택하려면 학습 데이터에서 별도의 검증용 데이터를 나누고, 테스트 데이터는 최종 평가에 사용하는 것이 적절하다.

이미지 한 장을 예측하기

image, label = test_dataset[0]  # image: (1, 28, 28)
model.eval()
with torch.no_grad():
    logits = model(image.unsqueeze(0).to(device))  # (1, 1, 28, 28) -> (1, 10)
    prediction = logits.argmax(dim=1).item()

print(f"sample 0: label={label}, prediction={prediction}")
plt.imshow(image.squeeze(0), cmap="gray")
plt.title(f"Label: {label}, Prediction: {prediction}")
plt.axis("off")
plt.show()

test_dataset[0]은 앞서 지정한 ToTensor()를 적용한 이미지를 반환한다. unsqueeze(0)은 맨 앞에 크기 1의 배치 축을 추가한다. 이미지 한 장도 배치 형태로 넣어야 모델 안의 Flatten(start_dim=1)이 이미지의 세 축을 모두 펼칠 수 있다.

표시할 때는 squeeze(0)으로 채널 축을 제거해 (28, 28)로 만든다. 한 장의 그림은 입력과 예측의 대응을 살펴보기 위한 것이며, 전체 성능은 앞에서 계산한 테스트 정확도로 판단한다.

실행 결과 읽기

아래는 Python 3.13.9, PyTorch 2.13.0, torchvision 0.28.0의 CPU 환경에서 float32로 실행한 결과이다. 위 코드의 실제 출력 중 1·5·10·15에포크와 최종 평가를 발췌했다.

PyTorch 2.13.0, torchvision 0.28.0
device=cpu, train=60000, test=10000
epoch 01: train_loss=0.5352
epoch 05: train_loss=0.3071
epoch 10: train_loss=0.2845
epoch 15: train_loss=0.2744
test_loss=0.2744
test_accuracy=92.26% (9226/10000)
sample 0: label=7, prediction=7

학습 중 기록한 평균 손실은 첫 에포크의 0.5352에서 마지막 에포크의 0.2744로 줄었다. 최종 모델은 테스트 이미지 10,000장 중 9,226장을 맞혔으며, 첫 번째 테스트 이미지의 정답과 예측은 모두 7이었다.

이 결과는 위 설정으로 한 번 실행한 MNIST 테스트 분할의 성능이다. 다른 초기값·실행 환경에서는 수치가 달라질 수 있고, 직접 촬영한 손글씨처럼 입력의 배경이나 크기가 다른 이미지에서도 같은 정확도가 나온다는 뜻은 아니다.

요약과 활용

MNIST 소프트맥스 회귀의 입력은 이미지마다 784개의 픽셀 값이고, 출력은 숫자 0~9에 대응하는 로짓 열 개이다. CrossEntropyLoss는 이 로짓과 정수 정답으로 학습할 손실을 만들고, 예측에서는 클래스 축의 argmax로 숫자를 고른다.

전체 흐름을 복구할 때는 같은 전처리, 배치 축 유지, 로짓과 정답의 shape, 학습에서만 매개변수 갱신, 전체 테스트 표본으로 지표 집계를 확인하면 된다. 이 연결이 갖춰져야 데이터나 모델을 바꾸어도 학습과 평가 코드의 의미를 유지할 수 있다.

참고자료

딥 러닝 파이토치 교과서 — 05-04 소프트맥스 회귀로 MNIST 데이터 분류하기: 실습 주제, MNIST 분할과 기본 학습 설정을 참고했다. 코드는 학습·평가의 입력 처리와 지표 집계를 맞춰 다시 구성했다.
torchvision — MNIST, MNIST 구현: 데이터셋 인자, 표본 접근 시 변환 적용.
torchvision — ToTensor: 이미지의 텐서 변환과 값의 범위.
PyTorch — Linear, Flatten 구현, Sequential: 입력·출력 shape, 매개변수 구성, 모듈의 순차 적용.
PyTorch — CrossEntropyLoss: 로짓, 정답 클래스 인덱스, 기본 평균 손실.
PyTorch — Optimizing Model Parameters: 학습·평가 반복문과 옵티마이저 갱신.
PyTorch — Autograd mechanics: 평가 모드와 기울기 기록의 구분.
PyTorch — Reproducibility: 난수 시드와 재현성의 범위.